"""V4 checkpoint inference and cached joint-tokenized candidate scoring."""
from pathlib import Path

import sentencepiece as spm
import torch

from vimeml.training.data import file_sha
from vimeml.training.infer import JapaneseLM, ROOT
from vimeml.training.model_factory_v4 import FORMAT
from vimeml.training.model_factory import model_from_checkpoint
from vimeml.training.evaluate_ime import common_prefix_length
from vimeml.tokenizer.v4 import read_json


class JapaneseLMV4(JapaneseLM):
    def __init__(self, checkpoint, tokenizer_dir, device="cpu"):
        self.device = torch.device(device)
        saved = torch.load(checkpoint, map_location="cpu", weights_only=True)
        if saved.get("format") != FORMAT or saved.get("architecture") != "tiny_gpt_v4":
            raise ValueError("Explicit V4 checkpoint required")
        token_manifest = ROOT / saved["config"]["token_dir"] / "manifest.json"
        if file_sha(token_manifest) != saved["signatures"]["tokens"]:
            raise ValueError("Token data identity changed")
        tokens = read_json(token_manifest)
        model_path = Path(tokenizer_dir) / "tokenizer.model"
        digest = file_sha(model_path)
        if digest != tokens["tokenizer_model_sha256"]:
            raise ValueError("V4 tokenizer identity changed")
        self.processor = spm.SentencePieceProcessor(model_file=str(model_path))
        self.special = {name: getattr(self.processor, f"{name}_id")() for name in ("pad", "unk", "bos", "eos")}
        identity = {"tokenizer_model_sha256": digest, "vocab_size": 8192, "special_ids": self.special}
        if (self.special != tokens["special_ids"] or self.processor.vocab_size() != 8192
                or saved.get("tokenizer_identity") != identity):
            raise ValueError("Checkpoint tokenizer contract mismatch")
        self.model = model_from_checkpoint(saved).to(self.device).eval()
        self.forbidden = (0, 1, 2, 4, 5, 6)
        self.metadata = {"architecture": "tiny_gpt_v4", "checkpoint_sha256": file_sha(Path(checkpoint)),
                         "checkpoint_step": saved["step"], "tokenizer_sha256": digest,
                         "model": self.model.configuration(), "precision": "fp32",
                         "random_initialization_test_only": saved.get("random_initialization_test_only", False)}

    def prefix_ids(self, prompt):
        ids = super().prefix_ids(prompt)
        if any(i in (4, 5, 6) for i in ids):
            raise ValueError("Reserved conversion markers are not ordinary V4-a text")
        return ids


@torch.inference_mode()
def score_candidates_cached(lm, context, candidates):
    """Identical full-vocabulary suffix logP sum; no final EOS/truncation."""
    if not candidates or len(set(candidates)) != len(candidates) or any(not c for c in candidates):
        raise ValueError("Distinct nonempty candidates required")
    context_ids = lm.prefix_ids(context)
    sequences = [[2, *lm.processor.encode(context + c)] for c in candidates]
    for text, sequence in zip(candidates, sequences, strict=True):
        if (lm.processor.decode(sequence[1:]) != context + text or any(i in range(7) for i in sequence[1:])
                or not 1 <= len(sequence) - 1 <= 128):
            raise ValueError("Invalid candidate or context overflow")
    common = common_prefix_length([context_ids, *sequences])
    if common < 1 or any(len(seq) <= common for seq in sequences):
        raise ValueError("Candidate suffix must contain prediction targets")
    shared = lm.model.new_cache()
    initial = lm.model.forward_cached(torch.tensor([sequences[0][:common]], device=lm.device), shared)[0, -1]
    results = []
    for candidate, sequence in zip(candidates, sequences, strict=True):
        cache, logits = shared.clone(), initial
        values = []
        suffix = sequence[common:]
        for index, token in enumerate(suffix):
            values.append(float(logits.float().log_softmax(-1)[token]))
            if index + 1 < len(suffix):
                logits = lm.model.forward_cached(torch.tensor([[token]], device=lm.device), cache)[0, -1]
        results.append({"text": candidate, "log_probability_sum": sum(values),
            "log_probability_mean": sum(values) / len(values), "scored_tokens": len(values),
            "token_ids": suffix, "token_log_probabilities": values,
            "pieces": [lm.processor.id_to_piece(i) for i in suffix],
            "boundary_retokenized": sequence[:len(context_ids)] != context_ids})
    return {"common_prefix_tokens_including_bos": common,
            "context_tokens_including_bos": len(context_ids), "candidates": results}