"""Local teacher generation and answer-verified sequence distillation."""

from __future__ import annotations

import argparse
from collections import Counter
import json
from pathlib import Path
import re
from typing import Any

from .training import (_keys, _positive, _text, chat_template_kwargs, local_model_hashes, model_kwargs,
                       model_spec, read_json_snapshot, read_jsonl_snapshot, require_versions,
                       sha256, versions, write_manifest)


def final_answer(completion: Any) -> str | None:
    """Accept a whole answer or exactly one nonempty final <answer> block.

    No number guessing, code execution, symbolic algebra, or judge model. An
    answer match verifies only the result, never the teacher's explanation.
    """
    if isinstance(completion, list):
        if len(completion) != 1 or not isinstance(completion[0], dict) or completion[0].get("role") != "assistant":
            return None
        completion = completion[0].get("content")
    if not isinstance(completion, str) or not completion.strip():
        return None
    text = completion.strip()
    if "<answer" in text or "</answer" in text:
        if text.count("<answer") != 1 or text.count("</answer") != 1:
            return None
        match = re.search(r"<answer>([^<>]+)</answer>\Z", text)
        if match is None:
            return None
        answer = match.group(1).strip()
        return answer or None
    return text


def verified(completion: Any, solution: str) -> bool:
    answer = final_answer(completion)
    return isinstance(solution, str) and bool(solution.strip()) and answer == solution.strip()


def exact_reward(completions: list[Any], solution: list[str], **_: Any) -> list[float]:
    """TRL custom reward signature; additional dataset columns are supplied by TRL."""
    if len(completions) != len(solution):
        raise ValueError("Completion and reference batches must have the same length")
    return [float(verified(completion, reference)) for completion, reference in zip(completions, solution)]


def _write_rows(path: Path, rows: list[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    # Exclusive creation prevents accidental replacement of curated training data.
    with path.open("x", encoding="utf-8") as stream:
        for row in rows:
            stream.write(json.dumps(row, ensure_ascii=False) + "\n")


def prepare(input_path: str | Path, output_path: str | Path) -> dict[str, Any]:
    """Filter pre-generated teacher sequences, retaining one correct row per prompt."""
    source, output = Path(input_path).resolve(), Path(output_path).resolve()
    if source == output:
        raise ValueError("Input and output must differ")
    manifest_path = output.with_suffix(output.suffix + ".manifest.json")
    if output.exists() or manifest_path.exists():
        raise ValueError(f"Output or manifest already exists: {output}")
    rows, source_digest = read_jsonl_snapshot(source)
    accepted, rejected = [], Counter()
    seen: set[str] = set()
    references: dict[str, str] = {}
    for index, row in enumerate(rows, 1):
        prompt = _text(row.get("prompt"), f"row {index} prompt")
        completion = row.get("completion")
        if not isinstance(completion, str):
            raise ValueError(f"row {index} completion must be a string")
        solution = _text(row.get("solution"), f"row {index} solution")
        normalized = prompt.strip()
        if normalized in references and references[normalized] != solution.strip():
            raise ValueError(f"Conflicting references for the same prompt at row {index}")
        references[normalized] = solution.strip()
        if not verified(completion, solution):
            rejected["incorrect_or_malformed_final_answer"] += 1
        elif normalized in seen:
            rejected["duplicate_prompt"] += 1
        else:
            seen.add(normalized)
            accepted.append({"id": row.get("id", f"teacher-{index}"), "prompt": prompt,
                             "completion": completion, "solution": solution,
                             "provenance": {"source_sha256": source_digest, "source_row": index,
                                            "verifier": "strict-final-answer-v1",
                                            "teacher": row.get("teacher")}})
    if not accepted:
        raise ValueError("No teacher sequences passed verification; no output was written")
    _write_rows(output, accepted)
    report = {"schema_version": 1, "method": "verified_sequence_distillation",
              "input": str(source), "input_sha256": source_digest, "output": str(output),
              "output_sha256": sha256(output), "total": len(rows), "accepted": len(accepted),
              "rejected": dict(rejected), "verifier": "strict-final-answer-v1",
              "verification_scope": "final answer only; explanations are not verified",
              "dependencies": versions()}
    write_manifest(manifest_path, report)
    return report


def generation_plan(config_path: str | Path) -> dict[str, Any]:
    path = Path(config_path).resolve()
    config, config_digest = read_json_snapshot(path)
    _keys(config, {"schema_version", "model", "dataset", "output", "max_new_tokens", "seed", "device", "chat_template_kwargs"}, "generation")
    if type(config.get("schema_version")) is not int or config["schema_version"] != 1:
        raise ValueError("schema_version must be 1")
    dataset = (path.parent / _text(config.get("dataset"), "dataset")).resolve()
    output = (path.parent / _text(config.get("output"), "output")).resolve()
    if output == dataset:
        raise ValueError("Generation input and output must differ")
    model = model_spec(config.get("model"), path.parent)
    template_kwargs = chat_template_kwargs(config.get("chat_template_kwargs", {}))
    maximum, seed = config.get("max_new_tokens", 128), config.get("seed", 42)
    _positive(maximum, "max_new_tokens", integer=True)
    _positive(seed, "seed", integer=True, zero=True)
    device = config.get("device", "cpu")
    if not isinstance(device, str) or device not in {"cpu", "cuda"}:
        raise ValueError("device must be cpu or cuda")
    rows, dataset_digest = read_jsonl_snapshot(dataset)
    for index, row in enumerate(rows, 1):
        _text(row.get("prompt"), f"row {index} prompt")
        _text(row.get("solution"), f"row {index} solution")
    return {"schema_version": 1, "status": "planned", "model": model,
            "config_path": str(path), "config_sha256": config_digest,
            "input": str(dataset), "input_sha256": dataset_digest,
            "output": str(output), "rows": len(rows), "max_new_tokens": maximum,
            "seed": seed, "device": device, "do_sample": False,
            "chat_template_kwargs": template_kwargs}


def generate(config_path: str | Path) -> dict[str, Any]:
    """Generate sequences with a locally loaded Transformers teacher; never an API."""
    report = generation_plan(config_path)
    if sha256(report["config_path"]) != report["config_sha256"]:
        raise ValueError("Configuration changed after validation")
    rows, dataset_digest = read_jsonl_snapshot(report["input"])
    if dataset_digest != report["input_sha256"]:
        raise ValueError("Dataset changed after validation")
    output = Path(report["output"])
    manifest_path = output.with_suffix(output.suffix + ".manifest.json")
    if output.exists() or manifest_path.exists():
        raise ValueError(f"Output or manifest already exists: {output}")
    require_versions(("transformers", "torch"))
    from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
    import torch

    set_seed(report["seed"])
    spec = report["model"]
    if spec["source"] == "local":
        report["local_model_sha256"] = local_model_hashes(spec)
    tokenizer = AutoTokenizer.from_pretrained(spec["name"], **model_kwargs(spec))
    if not tokenizer.chat_template:
        raise ValueError("Teacher tokenizer requires a chat_template")
    model = AutoModelForCausalLM.from_pretrained(spec["name"], **model_kwargs(spec), dtype=torch.float32)
    if spec["source"] == "local" and local_model_hashes(spec) != report["local_model_sha256"]:
        raise ValueError("Local teacher model or tokenizer changed while loading")
    model.to(report["device"])
    model.eval()
    generated = []
    for index, row in enumerate(rows, 1):
        text = tokenizer.apply_chat_template([{"role": "user", "content": row["prompt"]}],
                                             tokenize=False, add_generation_prompt=True,
                                             **report["chat_template_kwargs"])
        inputs = tokenizer(text, return_tensors="pt", add_special_tokens=False).to(report["device"])
        with torch.inference_mode():
            tokens = model.generate(**inputs, max_new_tokens=report["max_new_tokens"],
                                    do_sample=False, pad_token_id=tokenizer.eos_token_id)
        completion = tokenizer.decode(tokens[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True)
        generated.append({"id": row.get("id", f"teacher-{index}"), "prompt": row["prompt"],
                          "solution": row["solution"], "completion": completion, "teacher": spec})
    _write_rows(output, generated)
    report.update(status="completed", output_sha256=sha256(output), dependencies=versions())
    write_manifest(manifest_path, report)
    return report


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    sub = parser.add_subparsers(dest="command", required=True)
    filtering = sub.add_parser("prepare")
    filtering.add_argument("input", type=Path)
    filtering.add_argument("output", type=Path)
    generation = sub.add_parser("generate")
    generation.add_argument("config", type=Path)
    generation.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()
    try:
        if args.command == "prepare":
            result = prepare(args.input, args.output)
        else:
            result = generation_plan(args.config) if args.dry_run else generate(args.config)
        print(json.dumps(result, indent=2))
    except (ValueError, OSError, RuntimeError, ImportError) as error:
        parser.exit(2, f"error: {error}\n")


if __name__ == "__main__":
    main()