import json
from pathlib import Path

import pytest
import torch
from safetensors import safe_open

from inference import DecisionModel, format_candidate


ROOT = Path(__file__).resolve().parent


def test_required_release_files_exist():
    required = {
        "README.md", "adapter_model.safetensors", "adapter_config.json",
        "decision_head.safetensors", "tokenizer.json", "tokenizer_config.json",
        "model.py", "inference.py", "requirements.txt", "example.py",
        "calibration.json", "evaluation_report.json",
    }
    assert not required.difference(path.name for path in ROOT.iterdir())


def test_release_configs_are_portable():
    adapter = json.loads((ROOT / "adapter_config.json").read_text(encoding="utf-8"))
    calibration = json.loads((ROOT / "calibration.json").read_text(encoding="utf-8"))
    assert adapter["base_model_name_or_path"] == "openbmb/MiniCPM5-2B-Base"
    assert calibration["use_temperature"] is False
    assert calibration["default_prediction"] == "raw"


def test_safetensors_are_readable_and_finite():
    for name in ("adapter_model.safetensors", "decision_head.safetensors"):
        with safe_open(ROOT / name, framework="pt", device="cpu") as handle:
            keys = list(handle.keys())
            assert keys
            for key in keys:
                assert torch.isfinite(handle.get_tensor(key)).all(), key


def test_probability_sum_with_stubbed_runtime(monkeypatch):
    model = DecisionModel.__new__(DecisionModel)
    model.device = torch.device("cpu")
    model.max_length = 512

    class Encoded(dict):
        def to(self, _device):
            return self

    class Tokenizer:
        def __call__(self, texts, **_kwargs):
            n = len(texts)
            return Encoded(input_ids=torch.ones((n, 2), dtype=torch.long), attention_mask=torch.ones((n, 2), dtype=torch.long))

    class Scorer:
        def __call__(self, input_ids, attention_mask):
            del attention_mask
            return torch.arange(input_ids.shape[0], dtype=torch.float32)

    model.tokenizer = Tokenizer()
    model.model = Scorer()
    result = model.decide(state="s", question="q", options=["a", "b", "c"])
    assert sum(result["probabilities"]) == pytest.approx(1.0, abs=1e-7)
    assert result["choice"] == "c"


def test_prompt_format_is_stable():
    text = format_candidate("state", "choice", "question", "option")
    assert text.endswith("How well does this candidate answer the question?")