"""Test OsmoImage with mandatory encryption and lossless compression."""

import tempfile
from pathlib import Path

import numpy as np
import pytest
import torch
import torch.nn as nn

from model_osmosis.format.image import OsmoImage, OsmoRuntime
from model_osmosis.format.serialization import save_osmo
from model_osmosis.format.file_format import (
    OSMOIMG_MAGIC,
    OSMOIMG_EXTENSION,
    is_osmoimg,
    write_osmoimg_header,
    read_osmoimg_header,
)
from model_osmosis.format.adaptive_compression import (
    ADAPTIVE_PAYLOAD_CODEC,
    CompressionDependencyError,
)


@pytest.fixture(autouse=True)
def _trust_home_tmp(monkeypatch, tmp_path):
    monkeypatch.setenv("OSMOSIS_TRUST_HOME", str(tmp_path / "trust"))


def _make_model(seed=0):
    torch.manual_seed(seed)
    model = nn.Sequential(
        nn.Linear(6, 16),
        nn.ReLU(),
        nn.Linear(16, 8),
        nn.ReLU(),
        nn.Linear(8, 3),
    )
    model.eval()
    return model


class TestOsmoRuntime:
    def test_roundtrip(self):
        rt = OsmoRuntime(
            source_framework="pytorch",
            min_ram_gb=2.0,
            is_encrypted=True,
        )
        d = rt.to_dict()
        rt2 = OsmoRuntime.from_dict(d)
        assert rt2.source_framework == "pytorch"
        assert rt2.min_ram_gb == 2.0
        assert rt2.is_encrypted is True

    def test_unknown_fields_ignored(self):
        d = {"source_framework": "jax", "unknown_field": 123}
        rt = OsmoRuntime.from_dict(d)
        assert rt.source_framework == "jax"


class TestFileFormat:
    def test_magic_bytes(self):
        assert OSMOIMG_MAGIC == b"OSMOIMG\x00"
        assert len(OSMOIMG_MAGIC) == 8

    def test_extension(self):
        assert OSMOIMG_EXTENSION == ".osmoimg"

    def test_header_roundtrip(self):
        import io
        metadata = {"source_framework": "pytorch", "is_encrypted": True}

        buf = io.BytesIO()
        write_osmoimg_header(buf, metadata)
        buf.seek(0)

        assert buf.read(8) == OSMOIMG_MAGIC
        buf.seek(0)

        result = read_osmoimg_header(buf)
        assert result["source_framework"] == "pytorch"
        assert result["is_encrypted"] is True

    def test_invalid_magic_raises(self):
        import io
        buf = io.BytesIO(b"NOTMAGIC" + b"\x00" * 20)
        with pytest.raises(ValueError, match="Not a valid .osmoimg"):
            read_osmoimg_header(buf)

    def test_is_osmoimg_detection(self):
        with tempfile.NamedTemporaryFile(suffix=".osmoimg", delete=False) as f:
            f.write(OSMOIMG_MAGIC + b"\x00" * 100)
            f.flush()
            assert is_osmoimg(f.name) is True

        with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f:
            f.write(b"not an osmoimg file")
            f.flush()
            assert is_osmoimg(f.name) is False


class TestMandatoryEncryption:
    """All .osmoimg files MUST be encrypted and losslessly compressed."""

    def test_always_encrypted(self):
        """No opt-out — every image is encrypted."""
        model = _make_model()
        image = OsmoImage.build(model, name="enc-mandatory")
        assert image.is_encrypted is True
        assert image.runtime.is_encrypted is True

    def test_always_preserve_quality(self):
        """Encryption is always lossless (preserve_quality=True)."""
        model = _make_model()
        image = OsmoImage.build(model, name="quality-check")
        assert image.runtime.preserve_quality is True

    def test_no_encrypt_parameter(self):
        """build() should NOT accept encrypt/compress/preserve_quality args."""
        import inspect
        sig = inspect.signature(OsmoImage.build)
        params = list(sig.parameters.keys())
        assert "encrypt" not in params
        assert "compress" not in params
        assert "preserve_quality" not in params


class TestOsmoImageBuild:
    def test_build_from_pytorch(self):
        model = _make_model()
        image = OsmoImage.build(model, name="test-model")
        assert image.name == "test-model"
        assert image.parameter_count > 0
        assert image.is_encrypted is True  # enforced!


class TestOsmoImageSaveLoad:
    def test_saves_as_osmoimg(self):
        model = _make_model()
        image = OsmoImage.build(model, name="ext-test")

        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "test"
            saved = image.save(path)
            assert saved.suffix == ".osmoimg"
            assert saved.exists()

    def test_magic_bytes_present(self):
        model = _make_model()
        image = OsmoImage.build(model, name="magic-test")

        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "model.osmoimg"
            saved = image.save(path)

            with open(saved, "rb") as f:
                magic = f.read(8)
            assert magic == OSMOIMG_MAGIC

    def test_is_osmoimg_on_saved_file(self):
        model = _make_model()
        image = OsmoImage.build(model, name="detect-test")

        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "model.osmoimg"
            saved = image.save(path)
            assert is_osmoimg(saved) is True

    def test_save_load_roundtrip(self):
        model = _make_model()
        image = OsmoImage.build(model, name="roundtrip")

        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "test.osmoimg"
            saved_path = image.save(path)
            assert saved_path.exists()

            loaded = OsmoImage.load(saved_path)
            assert loaded.name == "roundtrip"
            assert loaded.parameter_count == image.parameter_count
            assert loaded.is_encrypted is True  # verified on load too

    def test_always_compressed(self):
        """Verify the .osmoimg uses compression (zip payload is valid)."""
        import zipfile, io
        model = _make_model()
        image = OsmoImage.build(model, name="compress-test")

        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "model.osmoimg"
            saved = image.save(path)

            # Read past the magic header to get the zip payload
            with open(saved, "rb") as f:
                f.read(8)  # magic
                version = int.from_bytes(f.read(4), "little")
                header_len = int.from_bytes(f.read(4), "little")
                f.read(header_len)  # compressed header
                zip_bytes = f.read()

            # The payload must be a valid zip archive (= compressed)
            assert version in (1, 3)
            assert zipfile.is_zipfile(io.BytesIO(zip_bytes)), "Payload must be valid zip"

    def test_lossless_max_v2_roundtrip(self):
        model = _make_model()
        image = OsmoImage.build(model, name="v2-roundtrip", compression_profile="lossless-max")

        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "v2.osmoimg"
            saved = image.save(path, compression_profile="lossless-max")

            with open(saved, "rb") as f:
                f.read(8)  # magic
                version = int.from_bytes(f.read(4), "little")
            assert version in (2, 3)

            loaded = OsmoImage.load(saved)
            assert loaded.runtime.compression_profile == "lossless-max"
            assert loaded.runtime.payload_codec == ADAPTIVE_PAYLOAD_CODEC

            for key, arr in image.osmo_model.parameters.items():
                assert key in loaded.osmo_model.parameters
                assert loaded.osmo_model.parameters[key].dtype == arr.dtype
                assert loaded.osmo_model.parameters[key].shape == arr.shape
                assert loaded.osmo_model.parameters[key].tobytes() == arr.tobytes()

    def test_lossless_max_missing_dependency_raises(self, monkeypatch):
        model = _make_model()
        image = OsmoImage.build(model, name="v2-missing-zstd")

        def _raise():
            raise CompressionDependencyError("zstd missing")

        monkeypatch.setattr("model_osmosis.format.image.ensure_zstandard", _raise)
        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "v2.osmoimg"
            with pytest.raises(CompressionDependencyError):
                image.save(path, compression_profile="lossless-max")

    def test_always_embeds_tokenizer_manifest(self):
        model = _make_model()
        image = OsmoImage.build(model, name="tok-manifest")

        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "tok-manifest.osmoimg"
            image.save(path)
            loaded = OsmoImage.load(path)
            assert "tokenizer_manifest.json" in loaded.tokenizer_data
            manifest = loaded.tokenizer_data["tokenizer_manifest.json"]
            assert isinstance(manifest, dict)
            assert manifest.get("embedded") is True

    def test_placeholder_tokenizer_bundle_does_not_break_generic_run(self):
        model = _make_model()
        image = OsmoImage.build(model, name="tok-placeholder")
        image.tokenizer_data = {
            "tokenizer_manifest.json": {
                "schema_version": 1,
                "embedded": True,
                "has_real_tokenizer_assets": False,
                "model_type": "generic",
                "source_framework": "pytorch",
                "hf_model_id": "",
            },
            "tokenizer_placeholder.txt": "placeholder",
        }
        x = np.zeros((1, 6), dtype=np.float32)
        out = image.run(input_data=x)
        assert isinstance(out, np.ndarray)
        assert out.shape[0] == 1

    def test_binary_tokenizer_assets_roundtrip_exact_bytes(self):
        model = _make_model()
        image = OsmoImage.build(model, name="tok-binary-roundtrip")
        spiece_bytes = b"\x00\x01\xffSENTENCEPIECE\x10\x11"
        image.tokenizer_data = {
            "spiece.model": spiece_bytes,
            "tokenizer_config.json": {"model_max_length": 128},
            "vocab.txt": "hello\nworld\n",
        }

        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "tok-binary-roundtrip.osmoimg"
            image.save(path)
            loaded = OsmoImage.load(path)
            assert loaded.tokenizer_data["spiece.model"] == spiece_bytes
            validation = loaded.runtime.compression_stats.get("tokenizer_validation", {})
            assert validation.get("status") == "real_assets"

    def test_load_rejects_corrupt_tokenizer_manifest_schema_v2(self):
        model = _make_model()
        image = OsmoImage.build(model, name="tok-corrupt-manifest")
        osmo = image.osmo_model
        osmo.auxiliary_files["runtime.json"] = image.runtime.to_dict()
        osmo.auxiliary_files["tokenizer/spiece.model"] = b"abc"
        osmo.auxiliary_files["tokenizer/tokenizer_manifest.json"] = {
            "schema_version": 2,
            "embedded": True,
            "has_real_tokenizer_assets": True,
            "model_type": "causal_lm",
            "source_framework": "huggingface",
            "hf_model_id": "dummy/model",
            "files": [
                {
                    "name": "spiece.model",
                    "kind": "binary",
                    "sha256": "0" * 64,
                    "size_bytes": 3,
                }
            ],
        }

        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "corrupt-tokenizer.osmo"
            save_osmo(osmo, path, compress=True, obfuscate=False)
            with pytest.raises(ValueError, match="Tokenizer bundle validation failed"):
                OsmoImage.load(path)


class TestOsmoImageInspect:
    def test_inspect_shows_encrypted(self):
        model = _make_model()
        image = OsmoImage.build(model, name="inspectable")

        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "inspect.osmoimg"
            image.save(path)

            info = OsmoImage.inspect(path)
            assert info["name"] == "inspectable"
            assert info["parameter_count"] > 0
            assert info["runtime"]["is_encrypted"] is True
            assert info["runtime"]["preserve_quality"] is True

    def test_inspect_shows_compression_fields(self):
        model = _make_model()
        image = OsmoImage.build(model, name="inspect-compression")

        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "inspect.osmoimg"
            image.save(path)
            info = OsmoImage.inspect(path)

            assert "compression_profile" in info["runtime"]
            assert "payload_codec" in info["runtime"]
            assert "compression_engine_version" in info["runtime"]
            assert "compression_stats" in info["runtime"]
            assert "compression_ratio_stored" in info["image"]
            assert "effective_compression_ratio" in info["image"]
            file_size_mb = path.stat().st_size / (1024 * 1024)
            expected_effective = (
                info["image"]["original_size_mb"] / file_size_mb
                if file_size_mb > 0 and info["image"]["original_size_mb"] > 0
                else 1.0
            )
            assert info["image"]["effective_compression_ratio"] == pytest.approx(
                expected_effective, rel=0.05
            )


class TestOsmoImageExport:
    def test_export_pytorch(self):
        model = _make_model()
        image = OsmoImage.build(model, name="exportable")

        exported = image.export(target_format="pytorch")
        assert isinstance(exported, dict)
        assert len(exported) > 0


class TestOsmoImageReadme:
    def test_readme_references_osmoimg(self):
        model = _make_model()
        image = OsmoImage.build(model, name="readme-test")
        assert "readme-test" in image.readme
        assert ".osmoimg" in image.readme


class TestHFBuildObfuscationSettings:
    def test_hf_build_disables_name_obfuscation_and_verification(self, monkeypatch):
        from model_osmosis.format.converters import UniversalConverter
        import model_osmosis.obfuscation as obf

        base_model = _make_model()
        fake_osmo = UniversalConverter().to_osmo(base_model, name="hf-mock")
        fake_osmo.architecture.architecture_type = "qwen"

        def fake_build_from_hf(cls, model_id, name, runtime):
            runtime.source_framework = "huggingface"
            runtime.hf_model_id = model_id
            runtime.model_type = "causal_lm"
            return fake_osmo, {"tokenizer.json": "{}"}, runtime

        captured = {}

        class DummyObfuscator:
            def __init__(self, config):
                captured["config"] = config

            def obfuscate(self, model):
                return model

        monkeypatch.setattr(OsmoImage, "_build_from_hf", classmethod(fake_build_from_hf))
        monkeypatch.setattr(obf, "WeightObfuscator", DummyObfuscator)

        image = OsmoImage.build("Qwen/Qwen2.5-Coder-0.5B", name="hf-mock-image")

        assert captured["config"].obfuscate_param_names is False
        assert captured["config"].obfuscate_layer_names is False
        assert captured["config"].verify is False
        assert image.runtime.hf_model_id == "Qwen/Qwen2.5-Coder-0.5B"

    def test_text_build_requires_real_tokenizer_assets(self, monkeypatch):
        from model_osmosis.format.converters import UniversalConverter

        base_model = _make_model()
        fake_osmo = UniversalConverter().to_osmo(base_model, name="hf-no-tokenizer")

        def fake_build_from_hf(cls, model_id, name, runtime):
            runtime.source_framework = "huggingface"
            runtime.hf_model_id = model_id
            runtime.model_type = "causal_lm"
            return fake_osmo, {}, runtime

        monkeypatch.setattr(OsmoImage, "_build_from_hf", classmethod(fake_build_from_hf))

        with pytest.raises(ValueError, match="Tokenizer packaging failed"):
            OsmoImage.build("mock/requires-tokenizer", name="hf-no-tokenizer-image")


class TestRuntimeCachingAndQuantization:
    def test_runtime_cache_hit_avoids_rebuild(self, monkeypatch):
        image = OsmoImage.build(_make_model(), name="runtime-cache-hit")

        calls = {"count": 0}

        def _fake_rebuild(*, runtime_profile, resolved_device, runtime_quantization, cache_key):
            calls["count"] += 1
            return {
                "cache_key": cache_key,
                "model": object(),
                "tokenizer": None,
                "runtime_profile": runtime_profile,
                "resolved_device": resolved_device,
                "runtime_quantization": runtime_quantization,
            }

        monkeypatch.setattr(image, "_rebuild_runtime_bundle", _fake_rebuild)
        bundle_a = image._load_for_inference(runtime_cache=True)
        bundle_b = image._load_for_inference(runtime_cache=True)
        assert calls["count"] == 1
        assert bundle_a is bundle_b

    def test_runtime_cache_invalidates_on_profile_device_quantization(self, monkeypatch):
        image = OsmoImage.build(_make_model(), name="runtime-cache-key")

        calls = {"count": 0}

        def _fake_rebuild(*, runtime_profile, resolved_device, runtime_quantization, cache_key):
            calls["count"] += 1
            return {
                "cache_key": cache_key,
                "model": object(),
                "tokenizer": None,
                "runtime_profile": runtime_profile,
                "resolved_device": resolved_device,
                "runtime_quantization": runtime_quantization,
            }

        monkeypatch.setattr(image, "_rebuild_runtime_bundle", _fake_rebuild)
        image._load_for_inference(runtime_profile="default", device="cpu", runtime_quantization="none")
        image._load_for_inference(runtime_profile="cpu-lowram", device="cpu", runtime_quantization="none")
        image._load_for_inference(runtime_profile="cpu-lowram", device="cpu", runtime_quantization="int8")
        assert calls["count"] == 3

    def test_clear_runtime_cache_forces_rebuild(self, monkeypatch):
        image = OsmoImage.build(_make_model(), name="runtime-cache-clear")

        calls = {"count": 0}

        def _fake_rebuild(*, runtime_profile, resolved_device, runtime_quantization, cache_key):
            calls["count"] += 1
            return {
                "cache_key": cache_key,
                "model": object(),
                "tokenizer": None,
                "runtime_profile": runtime_profile,
                "resolved_device": resolved_device,
                "runtime_quantization": runtime_quantization,
            }

        monkeypatch.setattr(image, "_rebuild_runtime_bundle", _fake_rebuild)
        image._load_for_inference(runtime_cache=True)
        image.clear_runtime_cache()
        image._load_for_inference(runtime_cache=True)
        assert calls["count"] == 2

    def test_runtime_quantization_int8_on_non_cpu_falls_back_to_none(self):
        image = OsmoImage.build(_make_model(), name="runtime-quant-fallback-device")
        model = nn.Linear(4, 4).eval()
        _, effective = image._apply_runtime_optimizations(
            model,
            runtime_profile="gpu-fast",
            resolved_device="cuda",
            runtime_quantization="int8",
        )
        assert effective == "none"

    def test_runtime_quantization_int8_failure_falls_back_to_none(self, monkeypatch):
        image = OsmoImage.build(_make_model(), name="runtime-quant-fallback-error")
        model = nn.Linear(4, 4).eval()

        def _raise(*args, **kwargs):
            raise RuntimeError("quantization unavailable")

        monkeypatch.setattr("torch.ao.quantization.quantize_dynamic", _raise, raising=False)
        monkeypatch.setattr("torch.quantization.quantize_dynamic", _raise, raising=False)
        _, effective = image._apply_runtime_optimizations(
            model,
            runtime_profile="cpu-lowram",
            resolved_device="cpu",
            runtime_quantization="int8",
        )
        assert effective == "none"


class TestLanguageRunDecoding:
    def test_run_returns_only_new_tokens_by_default(self, monkeypatch):
        image = OsmoImage.build(_make_model(), name="decode-new-only")

        class _FakeTokenizer:
            eos_token_id = 0

            def __call__(self, prompt, return_tensors="pt"):
                return {"input_ids": torch.tensor([[11, 12, 13]], dtype=torch.long)}

            def decode(self, ids, skip_special_tokens=True):
                if hasattr(ids, "tolist"):
                    ids = ids.tolist()
                return " ".join(str(int(x)) for x in ids)

        class _FakeModel:
            class config:
                is_encoder_decoder = False

            def generate(self, **kwargs):
                return torch.tensor([[11, 12, 13, 21, 22]], dtype=torch.long)

        monkeypatch.setattr(
            image,
            "_load_for_inference",
            lambda **kwargs: {
                "model": _FakeModel(),
                "tokenizer": _FakeTokenizer(),
                "resolved_device": "cpu",
            },
        )

        out = image.run(prompt="hi", max_tokens=2, temperature=0.0)
        assert out == "21 22"

    def test_run_include_prompt_restores_legacy_output(self, monkeypatch):
        image = OsmoImage.build(_make_model(), name="decode-with-prompt")

        class _FakeTokenizer:
            eos_token_id = 0

            def __call__(self, prompt, return_tensors="pt"):
                return {"input_ids": torch.tensor([[11, 12, 13]], dtype=torch.long)}

            def decode(self, ids, skip_special_tokens=True):
                if hasattr(ids, "tolist"):
                    ids = ids.tolist()
                return " ".join(str(int(x)) for x in ids)

        class _FakeModel:
            class config:
                is_encoder_decoder = False

            def generate(self, **kwargs):
                return torch.tensor([[11, 12, 13, 21, 22]], dtype=torch.long)

        monkeypatch.setattr(
            image,
            "_load_for_inference",
            lambda **kwargs: {
                "model": _FakeModel(),
                "tokenizer": _FakeTokenizer(),
                "resolved_device": "cpu",
            },
        )

        out = image.run(
            prompt="hi",
            max_tokens=2,
            temperature=0.0,
            include_prompt=True,
        )
        assert out == "11 12 13 21 22"
