"""Tests for converter reliability and tensor fidelity transport."""

from __future__ import annotations

import sys
import types

import pytest
import torch

from model_osmosis.format.converters import (
    HuggingFaceConverter,
    PyTorchConverter,
    TENSOR_FIDELITY_AUX_KEY,
)
from model_osmosis.format.osmo_format import (
    ArchitectureSpec,
    FrameworkOrigin,
    LayerSpec,
    LayerType,
    ModelCapabilities,
    ModelModality,
    OsmoManifest,
    OsmoModel,
)


class _TinyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(4, 8)


def _minimal_osmo_with_hf_textgen_keys() -> OsmoModel:
    manifest = OsmoManifest(
        name="hf-textgen",
        modality=ModelModality.TEXT,
        framework_origin=FrameworkOrigin.UNKNOWN,
        original_repo="fake/repo",
        capabilities=ModelCapabilities(tasks=["text-generation"]),
    )
    architecture = ArchitectureSpec(
        layers=[
            LayerSpec(
                name="model.embed_tokens",
                layer_type=LayerType.EMBEDDING,
                shape=[4, 4],
                dtype="float32",
            )
        ],
        hidden_size=4,
        num_layers=1,
        num_attention_heads=1,
        vocab_size=4,
        architecture_type="transformer",
    )
    params = {
        "model.embed_tokens.weight": torch.randn(4, 4).numpy(),
        "lm_head.weight": torch.randn(4, 4).numpy(),
    }
    return OsmoModel(manifest=manifest, architecture=architecture, parameters=params)


def test_pytorch_converter_preserves_bfloat16_with_raw_transport():
    converter = PyTorchConverter()
    state_dict = {
        "layer.weight": torch.randn(8, 4, dtype=torch.bfloat16),
        "layer.bias": torch.randn(8, dtype=torch.bfloat16),
    }

    osmo = converter.to_osmo(state_dict, name="bf16-test")
    fidelity = osmo.auxiliary_files[TENSOR_FIDELITY_AUX_KEY]
    tensors = fidelity["tensors"]

    assert tensors["layer.weight"]["encoding_mode"] == "raw_bytes"
    assert tensors["layer.bias"]["encoding_mode"] == "raw_bytes"
    assert osmo.parameters["layer.weight"].dtype.name == "uint8"
    assert osmo.parameters["layer.bias"].dtype.name == "uint8"

    restored = converter.from_osmo(osmo, strict_load=True)
    assert restored["layer.weight"].dtype == torch.bfloat16
    assert restored["layer.bias"].dtype == torch.bfloat16

    # Bit-level fidelity: compare int16 storage bits before/after transport.
    orig_w_bits = state_dict["layer.weight"].view(torch.int16)
    out_w_bits = restored["layer.weight"].view(torch.int16)
    orig_b_bits = state_dict["layer.bias"].view(torch.int16)
    out_b_bits = restored["layer.bias"].view(torch.int16)
    assert torch.equal(orig_w_bits, out_w_bits)
    assert torch.equal(orig_b_bits, out_b_bits)


def test_pytorch_converter_lossy_cast_requires_opt_in():
    converter = PyTorchConverter()
    state_dict = {
        "layer.weight": torch.randn(8, 4, dtype=torch.bfloat16),
        "layer.bias": torch.randn(8, dtype=torch.bfloat16),
    }
    osmo = converter.to_osmo(
        state_dict,
        name="bf16-lossy",
        allow_lossy_cast=True,
    )

    with pytest.raises(ValueError, match="allow_lossy_cast=True"):
        converter.from_osmo(osmo, strict_load=True, allow_lossy_cast=False)

    restored = converter.from_osmo(osmo, strict_load=True, allow_lossy_cast=True)
    assert restored["layer.weight"].dtype == torch.float32
    assert restored["layer.bias"].dtype == torch.float32
    assert "layer.weight" in converter.last_report.lossy_casts


def test_pytorch_strict_load_reports_key_mismatches():
    converter = PyTorchConverter()
    model = _TinyModule()
    osmo = converter.to_osmo(model.state_dict(), name="strict-load")

    # Remove one expected key and add an unexpected one to force mismatch.
    osmo.parameters.pop("layer.bias")
    osmo.parameters["unexpected.param"] = torch.randn(1).numpy()

    with pytest.raises(ValueError, match="Strict PyTorch load failed"):
        converter.from_osmo(
            osmo,
            model_class=_TinyModule,
            strict_load=True,
            allow_lossy_cast=False,
        )

    restored_model = converter.from_osmo(
        osmo,
        model_class=_TinyModule,
        strict_load=False,
        allow_lossy_cast=False,
    )
    assert isinstance(restored_model, _TinyModule)
    assert "layer.bias" in converter.last_report.missing_keys
    assert "unexpected.param" in converter.last_report.unexpected_keys


def test_huggingface_converter_uses_causal_lm_for_text_generation(monkeypatch):
    class _FakeConfig:
        pass

    class _FakeAutoConfig:
        @staticmethod
        def from_pretrained(model_id, trust_remote_code=True):
            assert model_id == "fake/repo"
            return _FakeConfig()

    class _LoadResult:
        def __init__(self, missing_keys=None, unexpected_keys=None):
            self.missing_keys = list(missing_keys or [])
            self.unexpected_keys = list(unexpected_keys or [])

    class _BaseModelInstance:
        def load_state_dict(self, state_dict, strict=False):
            if "model.embed_tokens.weight" in state_dict:
                return _LoadResult(
                    missing_keys=["embed_tokens.weight"],
                    unexpected_keys=["model.embed_tokens.weight"],
                )
            return _LoadResult()

    class _CausalModelInstance:
        def load_state_dict(self, state_dict, strict=False):
            missing = []
            if "model.embed_tokens.weight" not in state_dict:
                missing.append("model.embed_tokens.weight")
            return _LoadResult(missing_keys=missing, unexpected_keys=[])

    class _FakeAutoModel:
        @staticmethod
        def from_config(config, trust_remote_code=True):
            return _BaseModelInstance()

    class _FakeAutoModelForCausalLM:
        @staticmethod
        def from_config(config, trust_remote_code=True):
            return _CausalModelInstance()

    fake_tf = types.ModuleType("transformers")
    fake_tf.AutoConfig = _FakeAutoConfig
    fake_tf.AutoModel = _FakeAutoModel
    fake_tf.AutoModelForCausalLM = _FakeAutoModelForCausalLM
    monkeypatch.setitem(sys.modules, "transformers", fake_tf)

    osmo = _minimal_osmo_with_hf_textgen_keys()
    converter = HuggingFaceConverter()
    model = converter.from_osmo(osmo, strict_load=True)

    assert isinstance(model, _CausalModelInstance)
    assert converter.last_report.missing_keys == []
    assert converter.last_report.unexpected_keys == []
