"""Fixed-invocation command builder — the enforcement point for INV-T7.

arbo hands the worker a set of BOUNDED parameters; this module turns them into the ONE
`accelerate launch sdxl_train_network.py …` command the worker is allowed to run. Every
parameter is validated (type, range, allowlist, path-containment) before a single argv
element is produced. There is NO path from an arbo request to a free-form argument — the
flag set is a constant, only the *values* of vetted parameters vary.

`build_command` is a pure function (no I/O, no subprocess) so it is trivially unit-testable:
feed it a request dict, assert the argv + env. `InvalidTrainRequest` maps to HTTP 422.
"""

from __future__ import annotations

import re
from pathlib import Path

from . import config

# Safe token for names that land in a filename / kohya folder path. No shell metachars,
# no path separators, no leading dot — even though we never use a shell (argv list), this
# also protects the on-disk paths built from `output_name` / `trigger` / `subject_class`.
_SAFE_TOKEN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9 _.\-]{0,63}$")
_SAFE_OUTPUT_NAME = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_\-]{0,63}$")  # filename stem, no spaces


class InvalidTrainRequest(ValueError):
    """A train request failed validation → HTTP 422 (never a silent downgrade)."""


def _require(cond: bool, msg: str) -> None:
    if not cond:
        raise InvalidTrainRequest(msg)


def _validate_under(path_str: str, roots: tuple[Path, ...], field: str) -> Path:
    """Resolve `path_str` and require it to sit under one of `roots` (no traversal escape)."""
    _require(isinstance(path_str, str) and path_str.startswith("/"), f"{field} must be an absolute path")
    p = Path(path_str)
    # Reject traversal explicitly even before resolve() (belt): no '..' components.
    _require(".." not in p.parts, f"{field} must not contain '..'")
    resolved = p.resolve()
    for root in roots:
        try:
            resolved.relative_to(root.resolve())
            return resolved
        except ValueError:
            continue
    allowed = ", ".join(str(r) for r in roots)
    raise InvalidTrainRequest(f"{field} {resolved} is not under an allowed root ({allowed})")


def validate_request(req: dict) -> dict:
    """Validate + normalize an arbo /train request. Returns a clean param dict or raises."""
    # ---- tier -> (steps, dim, res) --------------------------------------------------
    tier_raw = req.get("tier")
    _require(tier_raw in config.TIERS, f"tier must be one of {sorted(config.TIERS)}; got {tier_raw!r}")
    tier = str(tier_raw)
    spec = config.TIERS[tier]

    # ---- device_index (int, 0=3090 / 1=A6000 under PCI_BUS_ID) ----------------------
    device_index = req.get("device_index")
    _require(isinstance(device_index, int) and device_index in (config.DEVICE_3090, config.DEVICE_A6000),
             f"device_index must be {config.DEVICE_3090} (3090) or {config.DEVICE_A6000} (A6000)")

    # ---- tier/device fit: quality (1024) is A6000-only (§4.6) -----------------------
    if tier == "quality":
        _require(device_index in config.QUALITY_ONLY_DEVICES,
                 "tier 'quality' (1024) requires the A6000 (device_index=1); the 3090 cannot fit it")

    # ---- names (safe tokens; used in argv AND on-disk paths) ------------------------
    trigger = req.get("trigger")
    subject_class = req.get("subject_class")
    output_name = req.get("output_name")
    _require(isinstance(trigger, str) and bool(_SAFE_TOKEN.match(trigger)), "trigger is not a safe token")
    _require(isinstance(subject_class, str) and bool(_SAFE_TOKEN.match(subject_class)),
             "subject_class is not a safe token")
    _require(isinstance(output_name, str) and bool(_SAFE_OUTPUT_NAME.match(output_name)),
             "output_name is not a safe filename stem")

    # ---- repeats + seed (bounded ints) ----------------------------------------------
    repeats = req.get("repeats")
    _require(isinstance(repeats, int) and 1 <= repeats <= 100, "repeats must be an int in [1, 100]")
    seed = req.get("seed", 42)
    _require(isinstance(seed, int) and 0 <= seed <= 2**31 - 1, "seed must be a non-negative int32")

    # ---- train_id (optional; arbo's id — the publish-path namespace, Phase 2) -------
    # Explicit over the derive-from-output_dir.parent fallback (decouples the worker from
    # arbo's handoff layout). Path-safe token since it lands in the loras/trained/{train_id}/ path.
    train_id = req.get("train_id")
    if train_id is not None:
        _require(isinstance(train_id, str) and bool(_SAFE_OUTPUT_NAME.match(train_id)),
                 "train_id is not a safe token")

    # ---- paths (containment-checked) ------------------------------------------------
    dataset_dir = _validate_under(req.get("dataset_dir", ""), (config.HANDOFF_ROOT,), "dataset_dir")
    output_dir = _validate_under(req.get("output_dir", ""), (config.HANDOFF_ROOT,), "output_dir")
    base_model_path = _validate_under(req.get("base_model_path", ""), config.ALLOWED_MODEL_ROOTS, "base_model_path")

    return {
        "tier": tier,
        "steps": spec["steps"],
        "dim": spec["dim"],
        "resolution": spec["resolution"],
        "device_index": device_index,
        "trigger": trigger,
        "subject_class": subject_class,
        "output_name": output_name,
        "repeats": repeats,
        "seed": seed,
        "dataset_dir": str(dataset_dir),
        "output_dir": str(output_dir),
        "base_model_path": str(base_model_path),
        "train_id": train_id,
    }


def published_relative_path(output_dir: str, output_name: str, train_id: str | None = None) -> str:
    """The ComfyUI-relative loras path for a succeeded LoRA (Phase 2 publish step).

    Uses the explicit `train_id` (arbo now sends it — decoupled from the path layout); falls
    back to deriving it from the handoff layout `HANDOFF_ROOT/{train_id}/output` (the name of the
    dir the LoRA was written into) when absent. Returns e.g. `trained/392cf898ac03/name.safetensors`,
    the exact string a ComfyUI `LoraLoader.lora_name` widget takes. Pure — no filesystem touch."""
    tid = train_id or Path(output_dir).parent.name
    return f"{config.PUBLISH_SUBDIR}/{tid}/{output_name}.safetensors"


def build_command(req: dict) -> tuple[list[str], dict[str, str], dict]:
    """Validate `req` and return `(argv, env_overlay, params)` for the fixed sd-scripts invocation.

    argv is a LIST (never a shell string) — no metachar interpretation is possible. env_overlay
    is merged onto os.environ by the caller (the CUDA ordering + allocator knobs from §4.6).
    `params` is the vetted, normalized parameter dict (so the caller need not re-validate).
    """
    p = validate_request(req)
    h = config.SDXL_HPARAMS
    batch = h["train_batch_size_full"] if p["device_index"] == config.DEVICE_A6000 else h["train_batch_size_lean"]
    alpha = int(round(p["dim"] * float(h["network_alpha_ratio"])))
    res = f"{p['resolution']},{p['resolution']}"

    argv = [
        str(config.ACCELERATE_BIN), "launch",
        "--num_processes", "1",
        "--num_cpu_threads_per_process", "1",
        "--mixed_precision", h["mixed_precision"],
        "--dynamo_backend", "no",
        str(config.SDXL_TRAIN_SCRIPT),
        "--pretrained_model_name_or_path", p["base_model_path"],
        "--train_data_dir", p["dataset_dir"],
        "--output_dir", p["output_dir"],
        "--output_name", p["output_name"],
        "--save_model_as", "safetensors",
        "--save_precision", h["save_precision"],
        "--mixed_precision", h["mixed_precision"],
        # LoRA network (§4.6 tier dim; unet-only per the lean recipe)
        "--network_module", "networks.lora",
        "--network_dim", str(p["dim"]),
        "--network_alpha", str(alpha),
        "--network_train_unet_only",
        # steps + resolution (§4.6 tier table)
        "--max_train_steps", str(p["steps"]),
        "--resolution", res,
        "--enable_bucket",
        "--train_batch_size", batch,
        # memory-lean base (§4.6 — the proven Sindra flag set)
        "--gradient_checkpointing",
        "--cache_latents", "--cache_latents_to_disk",
        "--cache_text_encoder_outputs", "--cache_text_encoder_outputs_to_disk",
        "--optimizer_type", "adamw8bit",
        "--no_half_vae",
        "--sdpa",
        "--caption_extension", ".txt",  # CRITICAL: sd-scripts defaults to .caption → would skip our captions
        # optimizer schedule (agent-discretion defaults; confirm vs Sindra)
        "--learning_rate", h["learning_rate"],
        "--lr_scheduler", h["lr_scheduler"],
        "--lr_warmup_steps", h["lr_warmup_steps"],
        "--min_snr_gamma", h["min_snr_gamma"],
        "--noise_offset", h["noise_offset"],
        "--max_data_loader_n_workers", h["max_data_loader_n_workers"],
        "--persistent_data_loader_workers",
        "--seed", str(p["seed"]),
    ]

    env_overlay = {
        "CUDA_DEVICE_ORDER": "PCI_BUS_ID",  # index 0=3090, 1=A6000 (recipe-pinned ordering)
        "CUDA_VISIBLE_DEVICES": str(p["device_index"]),
        "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
    }
    return argv, env_overlay, p