"""Small, inspectable adapters to pinned TRL SFT and GRPO trainers.

Planning uses only the Python standard library. Training is an explicit action;
models remain local unless the recipe explicitly permits a pinned download.
"""

from __future__ import annotations

import argparse
from datetime import datetime, timezone
import hashlib
import importlib.metadata
import json
import math
from pathlib import Path
import re
from typing import Any


PINS = {
    "trl": "0.29.0", "transformers": "5.3.0", "datasets": "4.6.1",
    "accelerate": "1.12.0", "peft": "0.18.1", "torch": "2.10.0",
}
_COMMON = {
    "max_steps": 10, "per_device_train_batch_size": 1,
    "gradient_accumulation_steps": 1, "learning_rate": 2e-5,
    "logging_steps": 1, "gradient_checkpointing": False,
    "bf16": False, "use_cpu": False,
}
_SFT = {"max_length": 512}
_GRPO = {"max_completion_length": 128, "num_generations": 2,
         "temperature": 1.0, "beta": 0.0, "loss_type": "grpo"}


def sha256(path: str | Path) -> str:
    digest = hashlib.sha256()
    with Path(path).open("rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def _unique_pairs(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
    result: dict[str, Any] = {}
    for key, value in pairs:
        if key in result:
            raise ValueError(f"Duplicate JSON key: {key}")
        result[key] = value
    return result


def read_json_snapshot(path: str | Path) -> tuple[dict[str, Any], str]:
    data = Path(path).read_bytes()
    value = json.loads(data.decode("utf-8"), object_pairs_hook=_unique_pairs)
    if not isinstance(value, dict):
        raise ValueError("Configuration must be a JSON object")
    return value, hashlib.sha256(data).hexdigest()


def read_json(path: str | Path) -> dict[str, Any]:
    return read_json_snapshot(path)[0]


def read_jsonl(path: str | Path) -> list[dict[str, Any]]:
    return read_jsonl_snapshot(path)[0]


def read_jsonl_snapshot(path: str | Path) -> tuple[list[dict[str, Any]], str]:
    """Parse and fingerprint the same byte snapshot, rejecting duplicate keys."""
    data = Path(path).read_bytes()
    rows = []
    for line_number, line in enumerate(data.decode("utf-8").splitlines(), 1):
        if not line.strip():
            continue
        try:
            row = json.loads(line, object_pairs_hook=_unique_pairs)
        except ValueError as error:
            raise ValueError(f"{path}:{line_number}: invalid JSON: {error}") from error
        if not isinstance(row, dict):
            raise ValueError(f"{path}:{line_number}: expected an object")
        rows.append(row)
    if not rows:
        raise ValueError(f"Dataset is empty: {path}")
    return rows, hashlib.sha256(data).hexdigest()


def _keys(value: dict[str, Any], allowed: set[str], label: str) -> None:
    if unknown := set(value) - allowed:
        raise ValueError(f"Unknown {label} fields: {', '.join(sorted(unknown))}")


def _text(value: Any, label: str) -> str:
    if not isinstance(value, str) or not value.strip():
        raise ValueError(f"{label} must be a nonempty string")
    return value


def _positive(value: Any, label: str, *, integer: bool = False, zero: bool = False) -> None:
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise ValueError(f"{label} must be numeric")
    if not math.isfinite(value) or (value < 0 if zero else value <= 0):
        raise ValueError(f"{label} must be finite and {'nonnegative' if zero else 'positive'}")
    if integer and not isinstance(value, int):
        raise ValueError(f"{label} must be an integer")


def model_spec(value: Any, base: Path) -> dict[str, Any]:
    if not isinstance(value, dict):
        raise ValueError("model must be an object")
    _keys(value, {"name", "revision", "allow_download"}, "model")
    name = _text(value.get("name"), "model.name")
    allow = value.get("allow_download", False)
    if not isinstance(allow, bool):
        raise ValueError("model.allow_download must be boolean")
    candidate = Path(name).expanduser()
    candidate = candidate if candidate.is_absolute() else base / candidate
    if candidate.is_dir():
        if allow:
            raise ValueError("allow_download is only meaningful for a Hub model")
        return {"name": str(candidate.resolve()), "revision": None, "allow_download": False,
                "source": "local"}
    if name.startswith((".", "/", "~")):
        raise ValueError(f"Local model directory does not exist: {candidate}")
    revision = value.get("revision")
    if not isinstance(revision, str) or re.fullmatch(r"[a-fA-F0-9]{40}", revision) is None:
        raise ValueError("Hub models require an immutable 40-character commit revision")
    return {"name": name, "revision": revision, "allow_download": allow, "source": "hub"}


def model_kwargs(spec: dict[str, Any]) -> dict[str, Any]:
    result = {"local_files_only": not spec["allow_download"], "trust_remote_code": False}
    if spec["revision"]:
        result["revision"] = spec["revision"]
    return result


def versions() -> dict[str, str | None]:
    found = {}
    for package in PINS:
        try:
            found[package] = importlib.metadata.version(package)
        except importlib.metadata.PackageNotFoundError:
            found[package] = None
    return found


def require_versions(packages: tuple[str, ...] = tuple(PINS)) -> dict[str, str | None]:
    installed = versions()
    mismatches = [f"{p}=={PINS[p]} (installed {installed[p]})" for p in packages
                  if (installed[p] or "").split("+")[0] != PINS[p]]
    if mismatches:
        raise RuntimeError("Install the pinned training extra in an isolated environment: " + ", ".join(mismatches))
    return installed


def local_model_hashes(spec: dict[str, Any]) -> dict[str, str]:
    """Fingerprint weights and tokenizer files for locally supplied checkpoints."""
    if spec["source"] != "local":
        return {}
    return {str(file.relative_to(spec["name"])): sha256(file)
            for file in sorted(Path(spec["name"]).rglob("*")) if file.is_file()}


def write_manifest(path: Path, data: dict[str, Any]) -> None:
    path.write_text(json.dumps(data, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")


def chat_template_kwargs(value: Any) -> dict[str, bool]:
    """Expose only the explicit Qwen3 thinking-mode switch, never arbitrary kwargs."""
    if not isinstance(value, dict):
        raise ValueError("chat_template_kwargs must be an object")
    _keys(value, {"enable_thinking"}, "chat_template_kwargs")
    if "enable_thinking" in value and not isinstance(value["enable_thinking"], bool):
        raise ValueError("chat_template_kwargs.enable_thinking must be boolean")
    return dict(value)


def validate_thinking_target(completion: str, template_kwargs: dict[str, bool]) -> None:
    if template_kwargs.get("enable_thinking") is not False:
        return
    blocks = re.findall(r"<think>(.*?)</think>", completion, flags=re.DOTALL)
    remainder = re.sub(r"<think>.*?</think>", "", completion, flags=re.DOTALL)
    if any(block.strip() for block in blocks) or "<think" in remainder or "</think" in remainder:
        raise ValueError("Non-thinking mode requires targets without a nonempty or malformed <think> block")


def plan(config_path: str | Path) -> dict[str, Any]:
    """Validate configuration and every training row without importing ML packages."""
    path = Path(config_path).resolve()
    config, config_digest = read_json_snapshot(path)
    _keys(config, {"schema_version", "method", "model", "dataset", "output_dir", "seed", "training", "lora", "chat_template_kwargs"}, "recipe")
    if type(config.get("schema_version")) is not int or config["schema_version"] != 1:
        raise ValueError("schema_version must be 1")
    method = config.get("method")
    if not isinstance(method, str) or method not in {"sft", "grpo"}:
        raise ValueError("method must be sft or grpo")
    model = model_spec(config.get("model"), path.parent)
    template_kwargs = chat_template_kwargs(config.get("chat_template_kwargs", {}))
    dataset = (path.parent / _text(config.get("dataset"), "dataset")).resolve()
    output = (path.parent / _text(config.get("output_dir"), "output_dir")).resolve()
    if output == dataset or output in dataset.parents:
        raise ValueError("output_dir cannot contain the training input")
    seed = config.get("seed", 42)
    _positive(seed, "seed", integer=True, zero=True)
    supplied = config.get("training", {})
    if not isinstance(supplied, dict):
        raise ValueError("training must be an object")
    defaults = {**_COMMON, **(_SFT if method == "sft" else _GRPO)}
    _keys(supplied, set(defaults), "training")
    settings = defaults | supplied
    for key, value in settings.items():
        if key in {"bf16", "use_cpu", "gradient_checkpointing"}:
            if not isinstance(value, bool):
                raise ValueError(f"{key} must be boolean")
        elif key == "loss_type":
            if not isinstance(value, str) or value not in {"grpo", "dapo", "dr_grpo"}:
                raise ValueError("loss_type must be grpo, dapo or dr_grpo")
        else:
            _positive(value, key, integer=key not in {"learning_rate", "temperature", "beta"}, zero=key == "beta")
    if settings["use_cpu"] and settings["bf16"]:
        raise ValueError("The CPU recipe uses float32; set bf16=false")
    if method == "grpo":
        generations = settings["num_generations"]
        if generations < 2:
            raise ValueError("GRPO num_generations must be at least 2")
        # Explicit batch = per-device batch * accumulation * world-size. This
        # adapter fixes steps_per_generation to accumulation for predictable runs.
        import os
        world_size = int(os.environ.get("WORLD_SIZE", "1"))
        batch = settings["per_device_train_batch_size"] * settings["gradient_accumulation_steps"] * world_size
        if batch % generations:
            raise ValueError("GRPO effective generation batch must divide by num_generations; increase batch/accumulation")
    lora = config.get("lora")
    if lora is not None:
        if not isinstance(lora, dict):
            raise ValueError("lora must be an object or null")
        _keys(lora, {"r", "alpha", "dropout", "target_modules"}, "lora")
        _positive(lora.get("r", 8), "lora.r", integer=True)
        _positive(lora.get("alpha", 16), "lora.alpha", integer=True)
        dropout = lora.get("dropout", 0.0)
        _positive(dropout, "lora.dropout", zero=True)
        if dropout >= 1:
            raise ValueError("lora.dropout must be below 1")
        targets = lora.get("target_modules", "all-linear")
        if targets != "all-linear" and (not isinstance(targets, list) or not targets or any(not isinstance(t, str) or not t for t in targets)):
            raise ValueError("lora.target_modules must be all-linear or a nonempty list of names")
    rows, dataset_digest = read_jsonl_snapshot(dataset)
    seen: set[str] = set()
    for index, row in enumerate(rows, 1):
        prompt = _text(row.get("prompt"), f"row {index} prompt")
        _text(row.get("completion" if method == "sft" else "solution"), f"row {index} target")
        if method == "sft":
            validate_thinking_target(row["completion"], template_kwargs)
        if prompt.strip() in seen:
            raise ValueError(f"Duplicate training prompt at row {index}")
        seen.add(prompt.strip())
    return {"schema_version": 1, "status": "planned", "method": method,
            "config_path": str(path), "config_sha256": config_digest, "model": model,
            "dataset": str(dataset), "dataset_sha256": dataset_digest, "rows": len(rows),
            "output_dir": str(output), "seed": seed, "training": settings, "lora": lora,
            "chat_template_kwargs": template_kwargs,
            "dependency_pins": PINS, "report_to": "none", "push_to_hub": False}


def run(config_path: str | Path) -> dict[str, Any]:
    """Run the upstream optimizer/trainer, saving a local checkpoint and manifest."""
    result = plan(config_path)
    if sha256(result["config_path"]) != result["config_sha256"]:
        raise ValueError("Configuration changed after validation")
    rows, dataset_digest = read_jsonl_snapshot(result["dataset"])
    if dataset_digest != result["dataset_sha256"]:
        raise ValueError("Dataset changed after validation")
    installed = require_versions()
    import torch
    from datasets import Dataset
    from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
    from trl import GRPOConfig, GRPOTrainer, SFTConfig, SFTTrainer
    from .distillation import exact_reward

    output = Path(result["output_dir"])
    if output.exists() and any(output.iterdir()):
        raise ValueError(f"output_dir must be empty to avoid overwriting a run: {output}")
    output.mkdir(parents=True, exist_ok=True)
    result.update(status="running", dependencies=installed,
                  started_at=datetime.now(timezone.utc).isoformat())
    if result["model"]["source"] == "local":
        result["local_model_sha256"] = local_model_hashes(result["model"])
    manifest = output / "recipe-manifest.json"
    write_manifest(manifest, result)
    try:
        set_seed(result["seed"])
        settings = dict(result["training"])
        tokenizer = AutoTokenizer.from_pretrained(result["model"]["name"], **model_kwargs(result["model"]))
        if not tokenizer.chat_template:
            raise ValueError("This conversational recipe requires a tokenizer with a chat_template")
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        tokenizer.padding_side = "left" if result["method"] == "grpo" else "right"
        model = AutoModelForCausalLM.from_pretrained(
            result["model"]["name"], **model_kwargs(result["model"]),
            dtype=torch.bfloat16 if settings["bf16"] else torch.float32,
        )
        if result["model"]["source"] == "local" and local_model_hashes(result["model"]) != result["local_model_sha256"]:
            raise ValueError("Local model or tokenizer changed while loading")
        records = []
        for row in rows:
            record = {"prompt": [{"role": "user", "content": row["prompt"]}]}
            if result["method"] == "sft":
                record["completion"] = [{"role": "assistant", "content": row["completion"]}]
                record["chat_template_kwargs"] = result["chat_template_kwargs"]
            else:
                record["solution"] = row["solution"]
            records.append(record)
        common = dict(output_dir=str(output), seed=result["seed"], data_seed=result["seed"],
                      report_to="none", push_to_hub=False, optim="adamw_torch",
                      save_strategy="no", dataloader_pin_memory=False, **settings)
        peft_config = None
        if result["lora"] is not None:
            from peft import LoraConfig
            lora = result["lora"]
            peft_config = LoraConfig(r=lora.get("r", 8), lora_alpha=lora.get("alpha", 16),
                                     lora_dropout=lora.get("dropout", 0.0),
                                     target_modules=lora.get("target_modules", "all-linear"), task_type="CAUSAL_LM")
        trainer_options = dict(model=model, processing_class=tokenizer,
                               train_dataset=Dataset.from_list(records), peft_config=peft_config)
        if result["method"] == "sft":
            trainer = SFTTrainer(args=SFTConfig(completion_only_loss=True, packing=False, **common), **trainer_options)
        else:
            trainer = GRPOTrainer(
                args=GRPOConfig(use_vllm=False, chat_template_kwargs=result["chat_template_kwargs"],
                                steps_per_generation=settings["gradient_accumulation_steps"], **common),
                reward_funcs=exact_reward, **trainer_options,
            )
        trained = trainer.train()
        trainer.save_model(str(output / "checkpoint"))
        tokenizer.save_pretrained(str(output / "checkpoint"))
        result.update(status="completed", metrics=trained.metrics, train_log_history=trainer.state.log_history,
                      checkpoint=str(output / "checkpoint"), completed_at=datetime.now(timezone.utc).isoformat())
        write_manifest(manifest, result)
        return result
    except Exception as error:
        result.update(status="failed", error_type=type(error).__name__)
        write_manifest(manifest, result)
        raise


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("config", type=Path)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()
    try:
        print(json.dumps(plan(args.config) if args.dry_run else run(args.config), indent=2))
    except (ValueError, OSError, RuntimeError) as error:
        parser.exit(2, f"error: {error}\n")


if __name__ == "__main__":
    main()