"""Declarative, resumable multi-stage training curricula.

A curriculum is a JSON file declaring an ordered list of training *stages* (each a
pure-LM or mixed LM+UD+SRL+chart run) applied to a single model architecture. The
runner (:mod:`scripts.run_curriculum`) executes each stage via ``torchrun`` on the
existing ``scripts.train_lm`` / ``scripts.train_mixed`` entry points, chaining each
stage's warm-start (``init_from``) to the previous stage's final checkpoint.

This module is the *pure* core — parsing, init-chain resolution, filesystem-based
completion detection, action planning (skip/run/resume/conflict), and torchrun
command construction. It imports no torch and runs no training, so it is fully
unit-testable. The orchestration side effects live in the runner script.

Design guarantees:
* **No silent overwrite** — a stage whose output already has ``train_result.json``
  is treated as complete and skipped; a partially-written stage is a *conflict*
  unless ``resume`` is requested.
* **Resumable** — an interrupted stage resumes from its own latest checkpoint;
  completed earlier stages are skipped and only feed their checkpoint forward.
* **Reproducible** — every stage records model/train config, corpus, resolved init
  and expected output checkpoints, objective weights, and the relation-vocab
  signature in a manifest.
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass, field
from pathlib import Path

VALID_KINDS = ("lm", "mixed")
PREVIOUS = "@previous"
_CKPT_RE = re.compile(r"^ckpt-(\d+)$")


@dataclass(frozen=True)
class Stage:
    """One training stage in a curriculum."""

    name: str
    kind: str  # "lm" | "mixed"
    train_config: str
    corpus: str
    run_id: str
    init_from: str | None = None  # None | "@previous" | explicit checkpoint dir
    max_steps: int | None = None
    micro_batch_size: int | None = None
    # mixed-only knobs (ignored for kind == "lm"):
    lm_ratio: float | None = None
    graph_languages: str | None = None
    graph_max_sentences_per_language: int | None = None

    def __post_init__(self) -> None:
        if self.kind not in VALID_KINDS:
            raise ValueError(f"stage {self.name!r}: kind must be one of {VALID_KINDS}, got {self.kind!r}")
        if not self.run_id:
            raise ValueError(f"stage {self.name!r}: run_id is required")
        if self.max_steps is not None and self.max_steps <= 0:
            raise ValueError(f"stage {self.name!r}: max_steps must be positive")


@dataclass(frozen=True)
class Curriculum:
    """An ordered sequence of stages over one model architecture."""

    name: str
    model_config: str
    relation_vocab_signature: str
    stages: tuple[Stage, ...]

    def __post_init__(self) -> None:
        if not self.stages:
            raise ValueError("curriculum must declare at least one stage")
        names = [s.name for s in self.stages]
        run_ids = [s.run_id for s in self.stages]
        if len(set(names)) != len(names):
            raise ValueError("stage names must be unique")
        if len(set(run_ids)) != len(run_ids):
            raise ValueError("stage run_ids must be unique")
        if self.stages[0].init_from == PREVIOUS:
            raise ValueError("first stage cannot init_from '@previous'")


def load_curriculum(path: str | Path) -> Curriculum:
    """Parse and validate a curriculum JSON file."""

    data = json.loads(Path(path).read_text(encoding="utf-8"))
    for key in ("name", "model_config", "relation_vocab_signature", "stages"):
        if key not in data:
            raise ValueError(f"curriculum is missing required key {key!r}")
    stages = tuple(Stage(**_stage_fields(s)) for s in data["stages"])
    return Curriculum(
        name=data["name"],
        model_config=data["model_config"],
        relation_vocab_signature=data["relation_vocab_signature"],
        stages=stages,
    )


_STAGE_KEYS = {f.name for f in Stage.__dataclass_fields__.values()}  # type: ignore[attr-defined]


def _stage_fields(raw: dict) -> dict:
    unknown = set(raw) - _STAGE_KEYS
    if unknown:
        raise ValueError(f"stage {raw.get('name')!r}: unknown keys {sorted(unknown)}")
    return raw


@dataclass(frozen=True)
class ResolvedStage:
    """A stage with its init/output checkpoints and filesystem status resolved."""

    stage: Stage
    output_dir: Path
    init_checkpoint: Path | None
    expected_checkpoint: Path | None
    completed: bool
    started: bool
    latest_checkpoint: Path | None


def latest_checkpoint(output_dir: Path) -> Path | None:
    """Return the ``ckpt-NNNNNNN`` with the highest step in ``output_dir``, or None."""

    if not output_dir.is_dir():
        return None
    best: tuple[int, Path] | None = None
    for child in output_dir.iterdir():
        match = _CKPT_RE.match(child.name)
        if match and child.is_dir():
            step = int(match.group(1))
            if best is None or step > best[0]:
                best = (step, child)
    return best[1] if best else None


def _expected_checkpoint(output_dir: Path, stage: Stage, default_max_steps: int | None) -> Path | None:
    steps = stage.max_steps if stage.max_steps is not None else default_max_steps
    return output_dir / f"ckpt-{steps:07d}" if steps is not None else None


def resolve_stages(
    curriculum: Curriculum, outputs_root: str | Path, *, default_max_steps: int | None = None
) -> list[ResolvedStage]:
    """Resolve init-chain, expected outputs, and completion status for each stage.

    ``@previous`` init resolves to the prior stage's *expected* final checkpoint,
    so the chain can be planned before any stage has run.
    """

    outputs_root = Path(outputs_root)
    resolved: list[ResolvedStage] = []
    prev_expected: Path | None = None
    for stage in curriculum.stages:
        output_dir = outputs_root / stage.run_id
        if stage.init_from == PREVIOUS:
            if prev_expected is None:
                raise ValueError(f"stage {stage.name!r}: '@previous' has no resolvable prior checkpoint")
            init_ckpt: Path | None = prev_expected
        elif stage.init_from:
            init_ckpt = Path(stage.init_from)
        else:
            init_ckpt = None
        expected = _expected_checkpoint(output_dir, stage, default_max_steps)
        completed = (output_dir / "train_result.json").is_file()
        latest = latest_checkpoint(output_dir)
        resolved.append(
            ResolvedStage(
                stage=stage,
                output_dir=output_dir,
                init_checkpoint=init_ckpt,
                expected_checkpoint=expected,
                completed=completed,
                started=output_dir.is_dir() and any(output_dir.iterdir()) if output_dir.is_dir() else False,
                latest_checkpoint=latest,
            )
        )
        prev_expected = expected
    return resolved


def plan_actions(resolved: list[ResolvedStage], *, resume: bool = False) -> list[tuple[ResolvedStage, str]]:
    """Assign each stage an action.

    * ``skip`` — completed (``train_result.json`` present); never overwritten.
    * ``run`` — not started; execute fresh.
    * ``resume`` — started but not completed and ``resume=True``; restart from its
      latest checkpoint.
    * ``conflict`` — started but not completed and ``resume=False``; the caller must
      refuse to proceed rather than overwrite a partial run.
    """

    actions: list[tuple[ResolvedStage, str]] = []
    for r in resolved:
        if r.completed:
            actions.append((r, "skip"))
        elif not r.started:
            actions.append((r, "run"))
        elif resume:
            actions.append((r, "resume"))
        else:
            actions.append((r, "conflict"))
    return actions


def build_stage_command(
    resolved: ResolvedStage,
    *,
    model_config: str,
    nproc: int,
    sp_model: str,
    action: str,
    python: str = "python",
) -> list[str]:
    """Build the ``torchrun`` argv for a stage (pure; no execution).

    ``action`` is ``"run"`` (use warm-start ``init_from``) or ``"resume"`` (restart
    from the stage's latest checkpoint instead).
    """

    stage = resolved.stage
    module = "scripts.train_lm" if stage.kind == "lm" else "scripts.train_mixed"
    cmd = [
        "torchrun", "--standalone", f"--nproc_per_node={nproc}", "-m", module,
        "--train-config", stage.train_config,
        "--model-config", model_config,
        "--run-id", stage.run_id,
    ]
    # LM stages read a pre-packed token corpus and need no tokenizer; mixed stages
    # tokenise UD/UP treebanks on the fly and require the SentencePiece model.
    cmd += ["--corpus" if stage.kind == "lm" else "--lm-corpus", stage.corpus]
    if stage.kind == "mixed":
        cmd += ["--tokenizer", "sentencepiece", "--sp-model", sp_model]
        if stage.lm_ratio is not None:
            cmd += ["--lm-ratio", str(stage.lm_ratio)]
        if stage.graph_languages:
            cmd += ["--graph-languages", stage.graph_languages]
        if stage.graph_max_sentences_per_language is not None:
            cmd += ["--graph-max-sentences-per-language", str(stage.graph_max_sentences_per_language)]
    if stage.max_steps is not None:
        cmd += ["--max-steps", str(stage.max_steps)]
    if stage.micro_batch_size is not None:
        cmd += ["--micro-batch-size", str(stage.micro_batch_size)]

    if action == "resume":
        if resolved.latest_checkpoint is None:
            raise ValueError(f"stage {stage.name!r}: resume requested but no checkpoint found")
        cmd += ["--resume-from", str(resolved.latest_checkpoint)]
    elif resolved.init_checkpoint is not None:
        if stage.kind == "lm":
            raise ValueError(f"stage {stage.name!r}: train_lm does not support warm-start init_from")
        cmd += ["--init-from", str(resolved.init_checkpoint)]
    return cmd


def stage_manifest(resolved: ResolvedStage, *, model_config: str, relation_vocab_signature: str) -> dict:
    """Reproducibility record for one stage (written by the runner)."""

    s = resolved.stage
    return {
        "name": s.name,
        "kind": s.kind,
        "run_id": s.run_id,
        "model_config": model_config,
        "train_config": s.train_config,
        "corpus": s.corpus,
        "init_checkpoint": str(resolved.init_checkpoint) if resolved.init_checkpoint else None,
        "expected_checkpoint": str(resolved.expected_checkpoint) if resolved.expected_checkpoint else None,
        "output_dir": str(resolved.output_dir),
        "relation_vocab_signature": relation_vocab_signature,
        "objective_weights": {
            k: v for k, v in (
                ("lm_ratio", s.lm_ratio),
                ("graph_languages", s.graph_languages),
                ("graph_max_sentences_per_language", s.graph_max_sentences_per_language),
            ) if v is not None
        },
    }