"""Freeze the suite as token ids for one model: `crit eval prepare --model DIR --data DIR`.

Engines never see text. Every engine evaluated against the same data directory receives
exactly the same prompt and corpus token ids, rendered once with the model's own chat
template. Requires the `tokenizers` and `jinja2` packages (preparation only).
"""
import hashlib
import json
from pathlib import Path

from . import SUITE_DIR, write_json

PREPARED = "prepared.json"


def file_sha256(path):
    digest = hashlib.sha256()
    with open(path, "rb") as handle:
        for block in iter(lambda: handle.read(1 << 20), b""):
            digest.update(block)
    return digest.hexdigest()


def model_identity(model):
    """Hash of the files that define tokens and shapes, plus weight file sizes.
    Weight bytes are not hashed here: engines and crit run records already do that."""
    model = Path(model)
    parts = {}
    for name in ("config.json", "tokenizer.json", "tokenizer_config.json", "chat_template.jinja",
                 "generation_config.json", "model.safetensors.index.json"):
        if (model / name).is_file():
            parts[name] = file_sha256(model / name)
    parts["weight_files"] = {p.name: p.stat().st_size for p in sorted(model.glob("*.safetensors"))}
    blob = json.dumps(parts, sort_keys=True).encode()
    return {"sha256": hashlib.sha256(blob).hexdigest(), "files": parts}


def suite_identity():
    digest = hashlib.sha256()
    for path in sorted(SUITE_DIR.rglob("*")):
        if path.is_file():
            digest.update(path.relative_to(SUITE_DIR).as_posix().encode() + b"\0" + path.read_bytes())
    return digest.hexdigest()


class Renderer:
    def __init__(self, model):
        from jinja2.sandbox import ImmutableSandboxedEnvironment
        import jinja2
        from tokenizers import Tokenizer
        model = Path(model)
        self.tokenizer = Tokenizer.from_file(str(model / "tokenizer.json"))
        template = model / "chat_template.jinja"
        if template.is_file():
            source = template.read_text(encoding="utf-8")
        else:
            source = json.loads((model / "tokenizer_config.json").read_text(encoding="utf-8"))["chat_template"]
        env = ImmutableSandboxedEnvironment(trim_blocks=True, lstrip_blocks=True)

        def raise_exception(message):
            raise jinja2.exceptions.TemplateError(message)
        env.globals["raise_exception"] = raise_exception
        self.template = env.from_string(source)

    def text(self, messages, generation_prompt):
        # Non-thinking mode for templates that support it; ignored by others.
        return self.template.render(messages=messages, add_generation_prompt=generation_prompt, enable_thinking=False)

    def ids(self, text):
        return self.tokenizer.encode(text, add_special_tokens=False).ids


def stop_ids(model):
    ids = set()
    for name in ("generation_config.json", "config.json"):
        path = Path(model) / name
        if path.is_file():
            data = json.loads(path.read_text(encoding="utf-8"))
            for source in (data, data.get("text_config", {})):
                value = source.get("eos_token_id")
                ids.update(value if isinstance(value, list) else [value] if value is not None else [])
    return sorted(ids)


def speed_case(renderer, case):
    messages = case["messages"]
    if "document" not in case:
        return renderer.ids(renderer.text(messages, True))
    # Exact-length long-context prompt: template head + instruction + document tokens cut to fit + template tail.
    sentinel = "\u0000DOCUMENT\u0000"
    content = messages[-1]["content"].replace("{document}", sentinel)
    rendered = renderer.text(messages[:-1] + [{"role": "user", "content": content}], True)
    head, tail = rendered.split(sentinel)
    head_ids, tail_ids = renderer.ids(head), renderer.ids(tail)
    document = (SUITE_DIR / case["document"]).read_text(encoding="utf-8")
    need = case["prompt_tokens"] - len(head_ids) - len(tail_ids)
    document_ids = renderer.ids(document)
    if need <= 0 or len(document_ids) < need:
        raise ValueError(f"{case['id']}: document has {len(document_ids)} tokens, needs {need}")
    return head_ids + document_ids[:need] + tail_ids


def quality_conversation(renderer, conversation):
    messages = conversation["messages"]
    if messages[-1]["role"] != "assistant":
        raise ValueError(f"{conversation['id']}: the scored turn must be the final assistant message")
    prefix = renderer.ids(renderer.text(messages[:-1], True))
    full = renderer.ids(renderer.text(messages, False).rstrip("\n"))
    if full[:len(prefix)] != prefix:
        raise ValueError(f"{conversation['id']}: the generation prompt is not a token prefix of the conversation")
    return full, len(prefix)


def prepare(model, data, speed_tokens=None):
    protocol = json.loads((SUITE_DIR.parent / "protocol.json").read_text(encoding="utf-8"))
    renderer = Renderer(model)
    speed = json.loads((SUITE_DIR / "speed.json").read_text(encoding="utf-8"))
    quality = json.loads((SUITE_DIR / "quality.json").read_text(encoding="utf-8"))
    prepared = {
        "protocol": protocol["protocol"],
        "suite_sha256": suite_identity(),
        "model": model_identity(model),
        "stop_ids": stop_ids(model),
        # Token ids that exist; head rows beyond this are padding and excluded from scoring.
        "valid_vocab": renderer.tokenizer.get_vocab_size(with_added_tokens=True),
        "speed": [{"id": c["id"], "kind": c["kind"], "prompt_ids": speed_case(renderer, c),
                   "output_tokens": speed_tokens or protocol["speed"]["output_tokens"]} for c in speed["cases"]],
        "quality": [],
    }
    for conversation in quality["conversations"]:
        ids, first = quality_conversation(renderer, conversation)
        prepared["quality"].append({"id": conversation["id"], "ids": ids, "first_target": first})
    Path(data).mkdir(parents=True, exist_ok=True)
    write_json(Path(data) / PREPARED, prepared)
    return prepared
