"""Compose existing data, training, and raw evaluation stages without supplied weights."""

from __future__ import annotations

import json
import os
import shutil
import subprocess
from copy import deepcopy
from pathlib import Path
from typing import Any

from omegaconf import DictConfig, OmegaConf
from softasp_core.config.validation import validate_config
from softasp_core.training.schedule import CORRECTED_PROTOCOL

from softasp.cli.train import run_training
from softasp.instantiations.runtime import evaluate as evaluate_release
from softasp.sudoku_backbone import extract_backbone
from softasp.sudoku_reporting import file_sha256
from softasp.sudoku_scratch_data import prepare_scratch_data


def validate_scratch_config(cfg: DictConfig) -> None:
    """Fail before downloads/training when the requested workflow is incompatible."""
    if (
        cfg.dataset.name != "sudoku"
        or list(cfg.dataset.splits) != ["train", "val", "test"]
        or cfg.model.name != "softasp_sudoku"
        or cfg.model.params.backbone != "default"
        or cfg.model.params.oracle_perception
        or cfg.backbone.model.name != "perception_only_sudoku"
        or cfg.backbone.model.params.d_model != cfg.model.params.d_model
        or cfg.release_profile != "sudoku_release1"
        or int(os.environ.get("WORLD_SIZE", "1")) != 1
    ):
        raise ValueError(
            "From-scratch workflow requires single-process Sudoku AS2 and matching default CNN"
        )
    if cfg.pipeline.data_mode not in ("generate", "existing"):
        raise ValueError("pipeline.data_mode must be generate or existing")
    subset = cfg.preprocessing.subset_size
    if subset is not None and (
        isinstance(subset, bool) or not isinstance(subset, int) or subset < 1
    ):
        raise ValueError("preprocessing.subset_size must be a positive integer or null")
    name = cfg.pipeline.name
    if not isinstance(name, str) or not name or name in (".", "..") or Path(name).name != name:
        raise ValueError("pipeline.name must be a single directory name")
    if cfg.pipeline.data_mode == "generate" and (
        cfg.data_artifact.path or cfg.data_artifact.sha256 or cfg.data_artifact.protocol_id
    ):
        raise ValueError("Generate mode cannot use an existing data_artifact; select existing mode")
    for model, trainer in ((cfg.model, cfg.trainer), (cfg.backbone.model, cfg.backbone.trainer)):
        if model.pretrained_from or trainer.resume_from:
            raise ValueError("From-scratch training cannot accept pretrained_from or resume_from")
        if (
            trainer.protocol != CORRECTED_PROTOCOL
            or trainer.seed != 42
            or trainer.distributed != "false"
            or trainer.compile
            or trainer.gradient_accumulation_steps != 1
            or trainer.epochs < 1
            or trainer.eval_on_start
            or not 1 <= trainer.eval_every <= trainer.epochs
            or trainer.batch_size_per_gpu < 1
            or trainer.num_workers < 0
        ):
            raise ValueError(
                "From-scratch training requires seed 42, corrected single-device training "
                "and post-training validation"
            )
        validate_config(
            OmegaConf.create({"trainer": OmegaConf.to_container(trainer, resolve=True)})
        )


def _stage_config(cfg: DictConfig, root: Path, *, backbone: bool) -> DictConfig:
    stage = deepcopy(cfg)
    if backbone:
        stage.model = deepcopy(cfg.backbone.model)
        stage.trainer = deepcopy(cfg.backbone.trainer)
        # This stage intentionally trains a fresh perception model, not AS2.
        stage.release_profile = None
    del stage.backbone
    del stage.pipeline
    stage.paths.output_dir = str(root / ("backbone" if backbone else "as2"))
    return stage


def train_from_scratch(config: DictConfig) -> Path:
    """Run a fresh two-stage training workflow and record exact artifact lineage."""
    cfg = deepcopy(config)
    OmegaConf.resolve(cfg)
    validate_scratch_config(cfg)
    root = Path(cfg.pipeline.run_dir).expanduser().resolve()
    models = Path(cfg.pipeline.model_dir).expanduser().resolve()
    if root == models or root in models.parents or models in root.parents:
        raise ValueError("Workflow run_dir and model_dir must be separate directories")
    for directory in (root, models):
        if directory.exists():
            raise FileExistsError(f"Use a fresh workflow directory: {directory}")
    root.mkdir(parents=True, exist_ok=False)
    models.mkdir(parents=True, exist_ok=False)
    manifest_path = root / "manifest.json"
    source_root = Path(__file__).resolve().parents[2]
    try:
        revision = subprocess.check_output(
            ["git", "rev-parse", "HEAD"], cwd=source_root, text=True, stderr=subprocess.DEVNULL
        ).strip()
        dirty = bool(
            subprocess.check_output(
                ["git", "status", "--porcelain"], cwd=source_root, text=True
            ).strip()
        )
    except (OSError, subprocess.CalledProcessError):
        revision, dirty = None, None
    record: dict[str, Any] = {
        "status": "running",
        "supplied_weights": False,
        "git_commit": revision,
        "git_dirty": dirty,
        "uv_lock_sha256": (
            file_sha256(source_root / "uv.lock")
            if (source_root / "uv.lock").is_file() else None
        ),
        "configuration": OmegaConf.to_container(cfg, resolve=True),
        "stages": {},
    }

    def save() -> None:
        manifest_path.write_text(json.dumps(record, indent=2, allow_nan=False) + "\n")

    save()
    try:
        prepare_scratch_data(cfg)
        record["dataset"] = OmegaConf.to_container(cfg.resolved_dataset, resolve=True)
        OmegaConf.save(cfg, root / "resolved.yaml", resolve=True)
        save()

        backbone_cfg = _stage_config(cfg, root, backbone=True)
        backbone_run = run_training(backbone_cfg)
        backbone_path = models / "backbone.pt"
        record["stages"]["backbone"] = extract_backbone(backbone_run / "best.pt", backbone_path)
        save()

        as2_cfg = _stage_config(cfg, root, backbone=False)
        as2_cfg.model.pretrained_from = str(backbone_path)
        # The release guard checks the complete tensor set, dtype, shape, and finiteness.
        as2_run = run_training(as2_cfg)
        selected = as2_run / "best.pt"
        checkpoint = models / "as2.pt"
        with selected.open("rb") as source, checkpoint.open("xb") as target:
            shutil.copyfileobj(source, target)
        record["stages"]["as2"] = {
            "source_checkpoint": str(selected),
            "path": str(checkpoint),
            "sha256": file_sha256(checkpoint),
            "backbone_sha256": file_sha256(backbone_path),
        }
        save()

        evaluation_cfg = deepcopy(as2_cfg)
        evaluation_cfg.dataset.splits = ["test"]
        evaluation_cfg.evaluator.checkpoint = str(checkpoint)
        evaluation_cfg.paths.output_dir = str(root / "evaluation")
        reports = evaluate_release(evaluation_cfg)
        record["stages"]["evaluation"] = {
            "report_dir": str(root / "evaluation" / "eval"),
            "metrics": reports["test"]["stages"],
        }
        record["status"] = "complete"
        save()
    except BaseException as error:
        record["status"] = "failed"
        record["error"] = f"{type(error).__name__}: {error}"
        save()
        raise
    return manifest_path