"""``nida train``: config + latent dataset directory -> a (resumable) training run.

Workspace layout under ``--out``::

    checkpoints/ckpt-<step>.pt   atomic, newest ``keep_last`` kept
    metrics.jsonl                train/eval records (replayed steps are removed on resume)
    run_info.json                config, seed, git commit, data fingerprint, load reports,
                                 ``data_dirs`` (all dataset directories, in load order),
                                 ``dataset_meta`` (the single directory's ``meta.json``, null
                                 for several directories) and ``dataset_metas`` (one
                                 ``meta.json`` or null per directory)
                                 (``resume_info-<step>.json`` when a run is resumed)
    data_rejections.jsonl        every item rejected at load time, with the reason
    latent_stats.json            per-channel mean/std the model was trained on (inference
                                 needs it to de-normalise generated latents)
    config.yaml                  verbatim copy of the training config file
"""

from __future__ import annotations

import json
import logging
import subprocess
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path

import torch

from nida.train.checkpoint import list_checkpoints
from nida.train.config import TrainConfig, load_train_config
from nida.train.data import (
    LANGUAGES,
    DataFormatError,
    LatentDataset,
    LoadReport,
    load_latent_datasets,
)
from nida.train.loop import CHECKPOINT_DIR, ResumeError, Trainer, TrainResult
from nida.utils.io import atomic_write_json, atomic_write_text

__all__ = ["TrainOutcome", "check_languages", "render_report", "resolve_git_commit", "run_training"]

log = logging.getLogger(__name__)


@dataclass
class TrainOutcome:
    """What ``run_training`` returns: the raw result plus the JSON/markdown summaries."""

    result: TrainResult
    metrics: dict[str, object]
    report_md: str


def resolve_git_commit(explicit: str | None, repo_dir: Path) -> str | None:
    """``explicit`` if given, else ``git rev-parse HEAD`` in ``repo_dir`` (None if unavailable)."""
    if explicit is not None:
        return explicit
    try:
        out = subprocess.run(
            ["git", "rev-parse", "HEAD"],
            cwd=repo_dir,
            capture_output=True,
            text=True,
            check=True,
            timeout=10,
        )
    except (OSError, subprocess.SubprocessError):
        return None
    return out.stdout.strip() or None


def _rejection_lines(reports: list[LoadReport]) -> str:
    lines = [
        json.dumps(
            {"split": r.split, "segment_id": x.segment_id, "reason": x.reason, "detail": x.detail},
            ensure_ascii=False,
            sort_keys=True,
        )
        for r in reports
        for x in r.rejections
    ]
    return "".join(f"{line}\n" for line in lines)


def _read_dataset_meta(data_dir: Path) -> object:
    """The directory's ``meta.json`` content, or None when it has none."""
    meta_path = data_dir / "meta.json"
    if not meta_path.is_file():
        return None
    return json.loads(meta_path.read_text(encoding="utf-8"))


def _pick_device(name: str | None) -> torch.device:
    if name is None or name == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return torch.device(name)


def check_languages(n_languages: int, datasets: Sequence[LatentDataset]) -> None:
    """Refuse data whose languages the model cannot represent (ADR-0007), before training.

    ``n_languages == 0`` (no language embedding) accepts only Turkish (id 0); otherwise every
    segment's language id must be below ``n_languages``.

    Raises:
        DataFormatError: naming the first offending language and how many segments use it.
    """
    allowed = max(n_languages, 1)
    counts: dict[str, int] = {}
    for data in datasets:
        for seg in data.segments:
            if LANGUAGES.index(seg.language) >= allowed:
                counts[seg.language] = counts.get(seg.language, 0) + 1
    if counts:
        raise DataFormatError(
            f"model n_languages={n_languages} cannot represent languages {counts} "
            f"(order {LANGUAGES}); set model.overrides.n_languages"
        )


def run_training(
    config_path: Path,
    data_dirs: Path | Sequence[Path],
    out_dir: Path,
    *,
    resume: bool,
    device: str | None = None,
    git_commit: str | None = None,
    repo_dir: Path | None = None,
    handle_sigterm: bool = True,
    init_from: Path | None = None,
) -> TrainOutcome:
    """Load config and data, train, and summarise.

    ``data_dirs`` is one latent dataset directory or several, trained on as their union
    (:func:`nida.train.data.load_latent_datasets`); latent stats come from the merged
    train split.

    ``resume=True`` continues from the newest checkpoint in ``out_dir`` (starting fresh,
    with a warning, when there is none -- so a restart loop can always pass ``--resume``).
    ``resume=False`` refuses to run when checkpoints already exist.

    ``init_from`` (a checkpoint of another run with the same model config, or the parent of
    a ``dur_cond`` "none" -> "emotion" step, see :meth:`Trainer.init_from`) starts a fresh
    run -- a fine-tune -- from its weights and latent stats (:meth:`Trainer.init_from`). It
    applies only when no checkpoint of this run was resumed, so ``--resume --init-from``
    restarts safely: the first launch inherits the weights, later launches continue the
    fine-tune's own checkpoints.

    Raises:
        ConfigError, DataFormatError, BudgetError, ResumeError, CheckpointError: see the
            respective modules; nothing is swallowed.
    """
    cfg = load_train_config(config_path)
    dirs = [data_dirs] if isinstance(data_dirs, Path) else list(data_dirs)
    dirs_label = ", ".join(str(d) for d in dirs)
    # E11: the aligned text path also needs each directory's timings.jsonl.
    timings = cfg.model_config().text_path == "aligned"
    train_data, train_report = load_latent_datasets(dirs, "train", timings=timings)
    try:
        val_data, val_report = load_latent_datasets(dirs, "val", timings=timings)
    except Exception:
        train_data.close()
        raise
    try:
        if len(train_data) == 0:
            raise DataFormatError(
                f"no usable training items in {dirs_label} "
                f"(rejections: {train_report.rejection_counts})"
            )
        check_languages(cfg.model_config().n_languages, [train_data, val_data])
        if len(val_data) == 0 and cfg.eval_every is not None:
            log.warning(
                "no usable validation items in %s: evaluation (eval_every=%s) is skipped",
                dirs_label,
                cfg.eval_every,
            )
        checkpoint_dir = out_dir / CHECKPOINT_DIR
        existing = list_checkpoints(checkpoint_dir)
        if existing and not resume:
            raise ResumeError(
                f"{checkpoint_dir} already holds checkpoints (latest step {existing[-1][0]}); "
                "pass --resume to continue or choose a fresh --out"
            )
        commit = resolve_git_commit(git_commit, repo_dir or Path.cwd())
        trainer = Trainer(
            cfg,
            train_data,
            out_dir,
            val_data=val_data,
            device=_pick_device(device),
            git_commit=commit,
        )
        out_dir.mkdir(parents=True, exist_ok=True)
        atomic_write_text(
            out_dir / "data_rejections.jsonl", _rejection_lines([train_report, val_report])
        )
        atomic_write_text(out_dir / "config.yaml", config_path.read_text(encoding="utf-8"))
        dataset_metas = [_read_dataset_meta(d) for d in dirs]
        resumed_path = trainer.resume_latest() if resume else None
        if resume and resumed_path is None:
            log.warning("--resume given but %s has no checkpoint; starting fresh", checkpoint_dir)
        if resumed_path is None and init_from is not None:
            trainer.init_from(init_from)
        elif init_from is not None and trainer.init_from_path != str(init_from.resolve()):
            log.warning(
                "--init-from %s ignored: resumed %s, whose run started from %s",
                init_from,
                resumed_path,
                trainer.init_from_path,
            )
        atomic_write_json(out_dir / "latent_stats.json", trainer.latent_stats.to_dict())
        info_name = (
            "run_info.json" if resumed_path is None else f"resume_info-{trainer.step:09d}.json"
        )
        atomic_write_json(
            out_dir / info_name,
            {
                "config": cfg.to_dict(),
                "seed": cfg.seed,
                "git_commit": commit,
                "device": str(trainer.device),
                "dataset_fingerprint": train_data.fingerprint(),
                "resumed_from": None if resumed_path is None else str(resumed_path),
                "init_from": trainer.init_from_path,
                "init_from_model_diff": trainer.init_from_model_diff,
                "latent_stats_inherited": trainer.latent_stats_inherited,
                "data": {"train": train_report.to_dict(), "val": val_report.to_dict()},
                "data_dirs": [str(d) for d in dirs],
                "dataset_meta": dataset_metas[0] if len(dirs) == 1 else None,
                "dataset_metas": dataset_metas,
                "prompt_pairing": {
                    **cfg.prompt_pairing.to_dict(),
                    **trainer.speaker_dirs.to_dict(),
                },
            },
        )
        result = trainer.run(handle_sigterm=handle_sigterm)
        num_params = sum(p.numel() for p in trainer.model.parameters())
        metrics = _metrics(cfg, result, train_report, val_report, trainer, commit, num_params)
    finally:
        _close(train_data, val_data)
    return TrainOutcome(result, metrics, render_report(metrics))


def _close(*datasets: LatentDataset) -> None:
    for dataset in datasets:
        dataset.close()


def _metrics(
    cfg: TrainConfig,
    result: TrainResult,
    train_report: LoadReport,
    val_report: LoadReport,
    trainer: Trainer,
    commit: str | None,
    num_params: int,
) -> dict[str, object]:
    return {
        "completed": result.completed,
        "stopped_by_signal": result.stopped_by_signal,
        "steps_done": result.steps_done,
        "total_steps": cfg.steps,
        "last_loss": result.last_loss,
        "mean_frames_per_s": result.mean_frames_per_s,
        "resumed_from_step": result.resumed_from_step,
        "evals": result.evals,
        "checkpoints": result.checkpoints,
        "best_checkpoint": result.best_checkpoint,
        "best_val_loss": trainer.best_val_loss,
        "seed": cfg.seed,
        "git_commit": commit,
        "device": str(trainer.device),
        "model": cfg.model.name,
        "num_params": num_params,
        "overfit_one_batch": cfg.overfit_one_batch,
        # E12: the pairing config, the matched speaker directories and the whole-run totals
        # (cross-segment prompts given, and why selected items kept the in-clip prompt).
        "prompt_pairing": {
            **cfg.prompt_pairing.to_dict(),
            **trainer.speaker_dirs.to_dict(),
            "totals": trainer.pairing_totals.to_dict(),
        },
        "dataset_fingerprint": trainer.train_data.fingerprint(),
        "out_dir": str(trainer.out_dir),
        "data": {"train": train_report.to_dict(), "val": val_report.to_dict()},
    }


def render_report(metrics: dict[str, object]) -> str:
    """Markdown run report from the ``metrics`` dict (a human-readable summary)."""

    def show(value: object) -> str:
        return "-" if value is None else str(value)

    state = (
        "completed"
        if metrics["completed"]
        else (
            "INTERRUPTED by SIGTERM (checkpoint saved, resume with --resume)"
            if metrics["stopped_by_signal"]
            else "stopped early"
        )
    )
    pairing = metrics["prompt_pairing"]
    assert isinstance(pairing, dict)
    totals = pairing["totals"]
    pairing_line = (
        "off (in-clip prompt)"
        if pairing["cross_p"] == 0
        else f"cross_p={pairing['cross_p']}, prompt {pairing['min_prompt_frames']}.."
        f"{pairing['max_prompt_frames']} frames, speaker dirs {pairing['matched_dirs']}; "
        f"{totals['paired']} of {totals['items']} items cross-segment, "
        f"fallbacks {totals['fallbacks']}"
    )
    lines = [
        "# train",
        "",
        f"- state: {state}",
        f"- steps: {metrics['steps_done']} / {metrics['total_steps']}",
        f"- resumed from step: {show(metrics['resumed_from_step'])}",
        f"- last loss: {show(metrics['last_loss'])}",
        f"- mean frames/s (this session): {show(metrics['mean_frames_per_s'])}",
        f"- model: {metrics['model']} ({metrics['num_params']} parameters)",
        f"- device: {metrics['device']}",
        f"- seed: {metrics['seed']}",
        f"- git commit: {show(metrics['git_commit'])}",
        f"- overfit_one_batch: {metrics['overfit_one_batch']}",
        f"- prompt pairing: {pairing_line}",
        f"- dataset fingerprint: {metrics['dataset_fingerprint']}",
        f"- workspace: {metrics['out_dir']}",
        "",
        "## Data",
        "",
    ]
    data = metrics["data"]
    assert isinstance(data, dict)
    for split in ("train", "val"):
        rep = data[split]
        lines.append(
            f"- {split}: {rep['num_accepted']} accepted of {rep['num_in_split']} "
            f"({rep['accepted_frames']} frames), rejected {rep['num_rejected']} "
            f"{rep['rejection_counts']}"
        )
    evals = metrics["evals"]
    assert isinstance(evals, list)
    if evals:
        lines += [
            "",
            "## Validation",
            "",
            "| step | val_loss | val_loss_ema | val_loss_ema (no text) |",
            "|---|---|---|---|",
        ]
        lines += [
            f"| {int(e['step'])} | {e['val_loss']:.5f} | {e['val_loss_ema']:.5f} | "
            f"{e.get('val_loss_ema_notext', float('nan')):.5f} |"
            for e in evals
        ]
    lines += [
        "",
        f"- best val_loss_ema: {show(metrics['best_val_loss'])} "
        f"({show(metrics['best_checkpoint'])})",
    ]
    checkpoints = metrics["checkpoints"]
    assert isinstance(checkpoints, list)
    if checkpoints:
        lines += ["", "## Checkpoints written this session", ""]
        lines += [f"- {c}" for c in checkpoints]
    return "\n".join(lines) + "\n"