"""Cross-framework reliability and determinism checks."""

from __future__ import annotations

import hashlib

import numpy as np
import pytest
import torch

from model_osmosis.format.converters import UniversalConverter
from model_osmosis.format.parity import compare_tensor_maps
from model_osmosis.utils.determinism import set_global_seed


def _make_tiny_torch(seed: int = 42) -> torch.nn.Module:
    set_global_seed(seed)
    return torch.nn.Sequential(
        torch.nn.Linear(8, 16),
        torch.nn.ReLU(),
        torch.nn.Linear(16, 4),
    ).eval()


def _tensor_map_fingerprint(tensors: dict[str, torch.Tensor]) -> str:
    digest = hashlib.sha256()
    for key in sorted(tensors.keys()):
        arr = tensors[key].detach().cpu().numpy()
        order = "F" if arr.flags.f_contiguous and not arr.flags.c_contiguous else "C"
        digest.update(key.encode("utf-8"))
        digest.update(arr.dtype.str.encode("utf-8"))
        digest.update(str(tuple(arr.shape)).encode("utf-8"))
        digest.update(order.encode("utf-8"))
        digest.update(arr.tobytes(order=order))
    return digest.hexdigest()


@pytest.mark.crossfw_core
def test_pytorch_to_osmo_to_pytorch_is_deterministic():
    set_global_seed(42)
    model_a = _make_tiny_torch(seed=42)
    model_b = _make_tiny_torch(seed=42)
    conv_a = UniversalConverter()
    conv_b = UniversalConverter()

    osmo_a = conv_a.to_osmo(model_a, name="det-a")
    osmo_b = conv_b.to_osmo(model_b, name="det-b")

    assert osmo_a.fingerprint == osmo_b.fingerprint
    roundtrip_a = conv_a.from_osmo(osmo_a, target_framework="pytorch")
    roundtrip_b = conv_b.from_osmo(osmo_b, target_framework="pytorch")

    result = compare_tensor_maps(
        {k: v.numpy() for k, v in roundtrip_a.items()},
        {k: v.numpy() for k, v in roundtrip_b.items()},
    )
    assert result.passed is True


@pytest.mark.crossfw_core
def test_hf_strict_roundtrip_parity():
    import tempfile

    transformers = pytest.importorskip("transformers")
    GPT2Config = transformers.GPT2Config
    GPT2Model = transformers.GPT2Model

    config = GPT2Config(n_layer=2, n_head=2, n_embd=16, n_positions=32, n_ctx=32, vocab_size=100)
    model = GPT2Model(config).eval()

    converter = UniversalConverter()
    osmo = converter.to_osmo(model, name="hf-strict")
    with tempfile.TemporaryDirectory() as tmpdir:
        model.config.save_pretrained(tmpdir)
        restored = converter.from_osmo(
            osmo,
            target_framework="huggingface",
            strict_load=True,
            model_id=tmpdir,
        ).eval()

        input_ids = torch.randint(0, config.vocab_size, (1, 8))
        with torch.no_grad():
            out_a = model(input_ids).last_hidden_state.detach().cpu().numpy()
            out_b = restored(input_ids).last_hidden_state.detach().cpu().numpy()
        assert np.allclose(out_a, out_b, rtol=1e-5, atol=1e-6)


@pytest.mark.crossfw_core
def test_onnx_adapter_parity_with_torch():
    pytest.importorskip("onnxruntime")
    from model_osmosis.adapters.onnx import ONNXAdapter

    model = _make_tiny_torch(seed=7)
    adapter = ONNXAdapter(verbose=False, providers=["CPUExecutionProvider"])
    onnx_expert = adapter.from_pytorch(model, input_shape=(1, 8))
    assert adapter.validate(onnx_expert, model, input_shape=(1, 8), rtol=1e-3, atol=1e-5)


@pytest.mark.crossfw_core
def test_repeat_run_seed_stability_for_parity_metrics():
    set_global_seed(42)
    model = _make_tiny_torch(seed=42)
    converter = UniversalConverter()
    osmo = converter.to_osmo(model, name="repeat-seed-stability")

    restored_a = converter.from_osmo(osmo, target_framework="pytorch", strict_load=True)
    report_a = compare_tensor_maps(
        {k: v.detach().cpu().numpy() for k, v in model.state_dict().items()},
        {k: v.detach().cpu().numpy() for k, v in restored_a.items()},
    )
    fp_a = _tensor_map_fingerprint(restored_a)

    set_global_seed(42)
    restored_b = converter.from_osmo(osmo, target_framework="pytorch", strict_load=True)
    report_b = compare_tensor_maps(
        {k: v.detach().cpu().numpy() for k, v in model.state_dict().items()},
        {k: v.detach().cpu().numpy() for k, v in restored_b.items()},
    )
    fp_b = _tensor_map_fingerprint(restored_b)

    assert report_a.to_dict() == report_b.to_dict()
    assert fp_a == fp_b


@pytest.mark.crossfw_extended
def test_tensorflow_roundtrip_smoke():
    tf = pytest.importorskip("tensorflow")
    from model_osmosis.format.converters import TensorFlowConverter

    model = tf.keras.Sequential(
        [
            tf.keras.layers.Input(shape=(8,)),
            tf.keras.layers.Dense(16, activation="relu"),
            tf.keras.layers.Dense(4),
        ]
    )
    converter = TensorFlowConverter()
    osmo = converter.to_osmo(model, name="tf-smoke")
    restored = converter.from_osmo(osmo, strict_load=True)
    assert isinstance(restored, dict)
    assert len(restored) > 0


@pytest.mark.crossfw_extended
def test_jax_roundtrip_smoke():
    pytest.importorskip("jax")
    from model_osmosis.format.converters import JAXConverter

    params = {
        "layer0": {"w": np.ones((8, 16), dtype=np.float32), "b": np.zeros((16,), dtype=np.float32)},
        "layer1": {"w": np.ones((16, 4), dtype=np.float32), "b": np.zeros((4,), dtype=np.float32)},
    }
    converter = JAXConverter()
    osmo = converter.to_osmo(params, name="jax-smoke")
    restored = converter.from_osmo(osmo, strict_load=True)
    assert "layer0" in restored
    assert "layer1" in restored
