"""CORTEX training pipeline.

Trains a CORTEX model from scratch on code and instruction data. Every dataset must be
declared in ``datasets/registry.json`` together with its source and license; the trainer
refuses to run on an undeclared or license-incompatible dataset.
"""

from __future__ import annotations

import json
import math
from dataclasses import dataclass
from pathlib import Path

import torch

from model.cortex_model import CortexConfig, count_parameters

__all__ = ["load_registry", "load_jsonl", "CortexTrainer", "train"]

REPO_ROOT = Path(__file__).resolve().parent.parent


def load_registry(path: str | Path | None = None) -> dict:
    path = Path(path) if path else REPO_ROOT / "datasets" / "registry.json"
    return json.loads(path.read_text(encoding="utf-8"))


def load_jsonl(path: str | Path) -> list[dict]:
    records = []
    with open(path, encoding="utf-8") as fh:
        for line in fh:
            line = line.strip()
            if line:
                records.append(json.loads(line))
    return records


def assert_dataset_usable(registry: dict, dataset_id: str) -> dict:
    """Refuse datasets whose license or provenance is not declared.

    This is the gate that keeps unlicensed or untraceable data out of CORTEX.
    """
    entries = {d["id"]: d for d in registry["datasets"]}
    if dataset_id not in entries:
        raise KeyError(f"dataset {dataset_id!r} is not declared in the registry")
    entry = entries[dataset_id]
    for field in ("source", "license", "license_url", "origin", "redistribution"):
        if not entry.get(field):
            raise ValueError(f"dataset {dataset_id!r} has no {field!r}; refusing to train on it")
    if entry["license"].lower() in {"unknown", "none", "proprietary"}:
        raise ValueError(
            f"dataset {dataset_id!r} has license {entry['license']!r}; refusing to train on it"
        )
    return entry


def build_token_stream(records: list[dict], tokenizer, seq_len: int):
    """Turn records into a flat token tensor cut into fixed-length windows."""
    ids: list[int] = []
    eos = tokenizer.token_to_id("<｜place▁holder▁no▁0｜>")
    if eos is None:
        eos = 1
    for r in records:
        text = r.get("text")
        if text is None:
            prompt, completion = r.get("prompt", ""), r.get("completion", "")
            text = f"{prompt}\n{completion}" if (prompt or completion) else ""
        if not text:
            continue
        ids.extend(tokenizer.encode(text, add_special_tokens=False).ids)
        ids.append(eos)
    if len(ids) < seq_len + 1:
        raise ValueError(f"not enough tokens: {len(ids)} for seq_len={seq_len}")
    n_windows = len(ids) // (seq_len + 1)
    tensor = torch.tensor(ids[: n_windows * (seq_len + 1)], dtype=torch.long)
    return tensor.view(n_windows, seq_len + 1)


@dataclass
class TrainState:
    step: int = 0
    tokens_seen: int = 0
    best_loss: float = float("inf")


class CortexTrainer:
    def __init__(self, model, config: CortexConfig, tokenizer, token_data: torch.Tensor,
                 train_cfg: dict, device: str = "cpu"):
        self.model = model.to(device)
        self.config = config
        self.tokenizer = tokenizer
        self.data = token_data
        self.cfg = train_cfg
        self.device = device
        self.state = TrainState()

        decay, no_decay = [], []
        for _, p in model.named_parameters():
            if not p.requires_grad:
                continue
            (no_decay if p.ndim < 2 else decay).append(p)
        self.optimizer = torch.optim.AdamW(
            [
                {"params": decay, "weight_decay": train_cfg["weight_decay"]},
                {"params": no_decay, "weight_decay": 0.0},
            ],
            lr=train_cfg["learning_rate"],
            betas=(train_cfg["beta1"], train_cfg["beta2"]),
        )

    def lr_at(self, step: int) -> float:
        warmup, total = self.cfg["warmup_steps"], self.cfg["max_steps"]
        base, floor = self.cfg["learning_rate"], self.cfg["min_learning_rate"]
        if step < warmup:
            return base * (step + 1) / max(warmup, 1)
        progress = (step - warmup) / max(total - warmup, 1)
        return floor + 0.5 * (base - floor) * (1 + math.cos(math.pi * min(progress, 1.0)))

    def train(self, max_steps: int | None = None, on_step=None, on_save=None,
              save_every: int | None = None, log_every: int | None = None):
        self.model.train()
        max_steps = max_steps or self.cfg["max_steps"]
        save_every = save_every or self.cfg["save_every"]
        log_every = log_every or self.cfg["log_every"]
        bs = self.cfg["micro_batch_size"]
        accum = self.cfg["gradient_accumulation_steps"]
        seq_len = self.cfg["sequence_length"]

        n_windows = self.data.shape[0]
        g = torch.Generator().manual_seed(self.cfg["seed"])
        history = []

        while self.state.step < max_steps:
            self.optimizer.zero_grad(set_to_none=True)
            accum_loss = 0.0
            for _ in range(accum):
                idx = torch.randint(0, n_windows, (bs,), generator=g)
                batch = self.data[idx].to(self.device)
                inputs, labels = batch[:, :-1], batch[:, 1:].clone()
                loss = self.model(inputs, labels=labels)["loss"] / accum
                loss.backward()
                accum_loss += loss.item()
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.cfg["grad_clip"])
            lr = self.lr_at(self.state.step)
            for group in self.optimizer.param_groups:
                group["lr"] = lr
            self.optimizer.step()

            self.state.step += 1
            self.state.tokens_seen += bs * accum * seq_len
            self.state.best_loss = min(self.state.best_loss, accum_loss)

            if log_every and self.state.step % log_every == 0:
                history.append({"step": self.state.step, "loss": accum_loss, "lr": lr})
                if on_step:
                    on_step(self.state.step, accum_loss, lr)
            if save_every and self.state.step % save_every == 0 and on_save:
                on_save(self.state.step, self.state.best_loss)

        self.model.eval()
        return history


def train(config_path, dataset_ids, tokenizer_path=None, out_dir="checkpoints/dev",
          max_steps=None, device="cpu", registry_path=None, seed=None, progress=False,
          save_every=None, log_every=None):
    """End-to-end training entry point used by the CLI."""
    from model.init import init_cortex_model, save_cortex_checkpoint
    from tokenizer.cortex_tokenizer import load_tokenizer, assert_vocab_matches_config

    raw = json.loads(Path(config_path).read_text(encoding="utf-8"))
    config = CortexConfig.from_json(config_path)
    config_seed = seed if seed is not None else raw.get("training", {}).get("seed", 1337)

    registry = load_registry(registry_path)
    records = []
    for dataset_id in dataset_ids:
        entry = assert_dataset_usable(registry, dataset_id)
        ds_path = REPO_ROOT / entry["path"]
        if not ds_path.exists():
            raise FileNotFoundError(
                f"dataset {dataset_id!r} declared at {entry['path']} is not present. "
                f"See datasets/README.md for how it is built."
            )
        records.extend(load_jsonl(ds_path))

    tokenizer = load_tokenizer(tokenizer_path)
    assert_vocab_matches_config(tokenizer, config.vocab_size)

    train_cfg = dict(raw["training"])
    if max_steps:
        train_cfg["max_steps"] = max_steps
    token_data = build_token_stream(records, tokenizer, train_cfg["sequence_length"])

    model = init_cortex_model(config, seed=config_seed)
    trainer = CortexTrainer(model, config, tokenizer, token_data, train_cfg, device=device)

    def dataset_provenance():
        return [
            {"id": d, "license": assert_dataset_usable(registry, d)["license"],
             "source": assert_dataset_usable(registry, d)["source"]}
            for d in dataset_ids
        ]

    def on_step(step, loss, lr):
        if progress:
            print(f"step {step:>6}  loss {loss:.4f}  lr {lr:.2e}", flush=True)

    def on_save(step, best):
        save_cortex_checkpoint(
            model, config, Path(out_dir) / f"step-{step}", step=step,
            extra_provenance={"datasets": dataset_provenance(), "training_config": train_cfg},
        )
        if progress:
            print(f"  saved checkpoint at step {step}", flush=True)

    history = trainer.train(max_steps=max_steps, on_step=on_step, on_save=on_save,
                            save_every=save_every, log_every=log_every)

    final_dir = save_cortex_checkpoint(
        model, config, Path(out_dir) / "final", step=trainer.state.step,
        extra_provenance={
            "datasets": dataset_provenance(),
            "training_config": train_cfg,
            "tokens_seen": trainer.state.tokens_seen,
            "final_loss": history[-1]["loss"] if history else None,
        },
    )
    return {
        "checkpoint": str(final_dir),
        "steps": trainer.state.step,
        "tokens_seen": trainer.state.tokens_seen,
        "final_loss": history[-1]["loss"] if history else None,
        "parameters": count_parameters(model),
        "history": history,
    }