import json
from pathlib import Path

import torch
from safetensors.torch import load_model
from transformers import LlamaConfig, LlamaForCausalLM, LlamaModel, PreTrainedTokenizerFast

from decision_lab.model_config import ScratchConfig
from decision_lab.models.decision_head import HeadDecisionModel
from decision_lab.models.scratch_transformer import ScratchDecisionModel
from decision_lab.models.token_model import TokenDecisionModel


def load_release(directory: Path, device: str = "cpu"):
    metadata = json.loads((directory / "config.json").read_text())
    architecture = metadata["architecture"]
    if architecture == "scratch":
        model = ScratchDecisionModel(ScratchConfig.model_validate(metadata["scratch"]))
    elif architecture == "token":
        model = TokenDecisionModel(
            LlamaForCausalLM(LlamaConfig(**metadata["backbone"])), metadata["answer_ids"]
        )
    elif architecture == "head":
        config = LlamaConfig(**metadata["backbone"])
        model = HeadDecisionModel(LlamaModel(config), config.hidden_size, 3)
    else:
        raise ValueError("Unsupported release architecture")
    load_model(model, directory / "model.safetensors", strict=True)
    model.to(device).eval()
    tokenizer = PreTrainedTokenizerFast.from_pretrained(directory)
    tokenizer.padding_side = "right"
    return model, tokenizer, metadata


def predict_release(directory: Path, prompts: list[str], device: str = "cpu"):
    model, tokenizer, metadata = load_release(directory, device)
    encoded = tokenizer(prompts, padding=True, truncation=False, return_tensors="pt").to(device)
    if encoded.input_ids.shape[1] > metadata["max_length"]:
        raise ValueError("Release inference input exceeds the length limit")
    with torch.no_grad():
        logits = model.class_logits(model(encoded.input_ids, encoded.attention_mask)).float().cpu()
    temperature = json.loads((directory / "calibration.json").read_text())["temperature"]
    probabilities = (logits / temperature).softmax(-1)
    return {
        "class_labels": metadata["class_labels"],
        "logits": logits.tolist(),
        "probabilities": probabilities.tolist(),
        "decisions": [metadata["class_labels"][index] for index in logits.argmax(-1).tolist()],
        "precision": "FP32 artifact verification inference",
    }


def verify_release(directory: Path, device: str = "cpu"):
    from decision_lab.training.checkpoint import file_hash

    inventory = json.loads((directory / "artifact-manifest.json").read_text())
    for name, digest in inventory["files"].items():
        path = directory / name
        if not path.resolve().is_relative_to(directory.resolve()) or file_hash(path) != digest:
            raise ValueError(f"Release file checksum mismatch: {name}")
    reference = json.loads((directory / "verification.json").read_text())
    actual = predict_release(directory, reference["prompts"], device)
    torch.testing.assert_close(
        torch.tensor(actual["logits"]),
        torch.tensor(reference["logits"]),
        atol=reference["logit_atol"],
        rtol=reference["logit_rtol"],
    )
    if actual["decisions"] != reference["decisions"]:
        raise ValueError("Artifact reload changed a fixed verification decision")
    return {
        "verified": True,
        "inputs": len(reference["prompts"]),
        "device": device,
        "artifact_identity": inventory["protocol_identity"],
    }