#!/usr/bin/env python3
from __future__ import annotations

import os
import runpy
import sys
from pathlib import Path


PLUGIN_DIR = Path(__file__).resolve().parents[1]
BACKEND_DIR = Path(os.environ.get("ANIMAFORGE_BACKEND_DIR", PLUGIN_DIR / "backend" / "sd-scripts")).expanduser().resolve()
TRAIN_PYTHON = Path(
    os.environ.get(
        "ANIMAFORGE_TRAIN_PYTHON",
        BACKEND_DIR / ".venv" / ("Scripts/python.exe" if os.name == "nt" else "bin/python"),
    )
).expanduser()
TRAIN_VENV = TRAIN_PYTHON.parent.parent.resolve()


def _value_after(args: list[str], name: str) -> str | None:
    if name not in args:
        return None
    index = args.index(name)
    if index + 1 >= len(args):
        return None
    return args[index + 1]


def _pop_arg(args: list[str], name: str) -> str | None:
    if name not in args:
        return None
    index = args.index(name)
    if index + 1 >= len(args):
        return None
    value = args[index + 1]
    del args[index:index + 2]
    return value


def _pop_optional_int(args: list[str], name: str, default: int) -> int:
    value = _pop_arg(args, name)
    if value is None:
        return default
    try:
        return max(1, int(value))
    except Exception:
        return default


def _ensure_backend_ready() -> None:
    if not BACKEND_DIR.exists():
        raise RuntimeError(
            f"sd-scripts backend was not found: {BACKEND_DIR}\n"
            "Run scripts/setup_sd_scripts_backend.ps1 on Windows before training."
        )
    if not (BACKEND_DIR / "anima_train_network.py").exists():
        raise RuntimeError(f"anima_train_network.py was not found under backend: {BACKEND_DIR}")
    if not TRAIN_PYTHON.exists():
        raise RuntimeError(
            f"training Python was not found: {TRAIN_PYTHON}\n"
            "Run scripts/setup_sd_scripts_backend.ps1 to create the backend virtual environment."
        )


def _ensure_dataset_config(args: list[str]) -> list[str]:
    if "--dataset_config" in args:
        return args

    num_repeats = _pop_optional_int(args, "--anima_num_repeats", 1)
    train_data_dir = _pop_arg(args, "--train_data_dir")
    if not train_data_dir:
        return args

    image_dir = Path(train_data_dir).expanduser().resolve()
    output_dir = Path(_value_after(args, "--output_dir") or image_dir / "anima_lora_output").expanduser().resolve()
    output_name = _value_after(args, "--output_name") or image_dir.name
    resolution = _value_after(args, "--resolution") or "1024,1024"
    batch_size = _value_after(args, "--train_batch_size") or "1"
    caption_extension = _value_after(args, "--caption_extension") or ".txt"
    caption_separator = _value_after(args, "--caption_separator") or ","
    keep_tokens = _value_after(args, "--keep_tokens") or "1"

    try:
        width, height = [item.strip() for item in resolution.split(",", 1)]
        resolution_toml = f"[{int(width)}, {int(height)}]"
    except Exception:
        resolution_toml = "1024"

    output_dir.mkdir(parents=True, exist_ok=True)
    config_path = output_dir / f"{output_name}.dataset.toml"
    config_path.write_text(
        "\n".join(
            [
                "[general]",
                "shuffle_caption = true" if "--shuffle_caption" in args else "shuffle_caption = false",
                f"caption_extension = {caption_extension!r}",
                f"caption_separator = {caption_separator!r}",
                f"keep_tokens = {int(keep_tokens)}",
                "",
                "[[datasets]]",
                f"resolution = {resolution_toml}",
                f"batch_size = {int(batch_size)}",
                "enable_bucket = true",
                "bucket_no_upscale = false",
                "",
                "  [[datasets.subsets]]",
                f"  image_dir = {str(image_dir)!r}",
                f"  num_repeats = {num_repeats}",
                f"  class_tokens = {image_dir.name!r}",
                "",
            ]
        ),
        encoding="utf-8",
    )
    args.extend(["--dataset_config", str(config_path)])
    return args


def main() -> None:
    _ensure_backend_ready()
    if os.path.realpath(sys.prefix) != os.path.realpath(TRAIN_VENV):
        os.execv(str(TRAIN_PYTHON), [str(TRAIN_PYTHON), __file__, *sys.argv[1:]])

    sys.argv[1:] = _ensure_dataset_config(list(sys.argv[1:]))
    os.chdir(BACKEND_DIR)
    sys.path.insert(0, str(BACKEND_DIR))
    runpy.run_path(str(BACKEND_DIR / "anima_train_network.py"), run_name="__main__")


if __name__ == "__main__":
    main()