from __future__ import annotations

import importlib.util
from pathlib import Path

import pytest

pytest.importorskip("torch")
pytest.importorskip("onnxruntime")
pytest.importorskip("transformers")

from vons.models import ModelConfig, build_model

_SPEC = importlib.util.spec_from_file_location(
    "verify_onnx_bundle", Path(__file__).parents[1] / "tools/verify_onnx_bundle.py"
)
assert _SPEC is not None and _SPEC.loader is not None
_MODULE = importlib.util.module_from_spec(_SPEC)
_SPEC.loader.exec_module(_MODULE)
_model_config_from_checkpoint = _MODULE._model_config_from_checkpoint


def test_token_conditioned_checkpoint_config_reconstructs_model_state() -> None:
    expected = ModelConfig(
        hidden_size=256,
        latent_size=64,
        diffusion_steps=64,
        inference_steps=8,
        abstain_threshold=0.55,
        diffusion_conditioning="token_cross_attention",
        attention_heads=4,
    )
    trained_head = build_model("diffusion", option_count=32, config=expected)
    checkpoint = {
        "backend": "diffusion",
        "hidden_size": 256,
        "config": {
            "diffusion_conditioning": "token_cross_attention",
            "latent_size": 64,
            "diffusion_steps": 64,
            "inference_steps": 8,
            "abstain_threshold": 0.55,
            "attention_heads": 4,
        },
    }

    reconstructed = _model_config_from_checkpoint(checkpoint)
    verifier_head = build_model("diffusion", option_count=32, config=reconstructed)
    verifier_head.load_state_dict(trained_head.state_dict())

    assert reconstructed == expected


def test_token_conditioned_checkpoint_uses_training_defaults() -> None:
    checkpoint = {
        "backend": "diffusion",
        "hidden_size": 256,
        "config": {"diffusion_conditioning": "token_cross_attention"},
    }

    reconstructed = _model_config_from_checkpoint(checkpoint)

    assert reconstructed.latent_size == 64
    assert reconstructed.diffusion_steps == 64
    assert reconstructed.inference_steps == 8
    assert reconstructed.diffusion_conditioning == "token_cross_attention"


def test_pooled_checkpoint_without_schedule_fields_uses_export_defaults() -> None:
    expected = ModelConfig(
        hidden_size=256,
        latent_size=32,
        diffusion_steps=32,
        inference_steps=4,
        diffusion_conditioning="pooled",
    )
    trained_head = build_model("diffusion", option_count=32, config=expected)
    checkpoint = {
        "backend": "diffusion",
        "hidden_size": 256,
        "config": {},
    }

    reconstructed = _model_config_from_checkpoint(checkpoint)
    verifier_head = build_model("diffusion", option_count=32, config=reconstructed)
    verifier_head.load_state_dict(trained_head.state_dict())

    assert reconstructed == expected