"""
End-to-end pipeline test: Build -> Obfuscate -> Save -> Load -> Verify.

Tests the FULL pipeline with a local PyTorch model (no HuggingFace download).
Verifies that:
- Obfuscation actually transforms weights
- Obfuscated model produces identical outputs (functional equivalence)
- Save/load roundtrip preserves everything
- Inspect returns correct metadata
- Export reconstructs a working model
"""

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.file_format import OSMOIMG_MAGIC, is_osmoimg
from model_osmosis.format.converters import UniversalConverter
from model_osmosis.obfuscation import (
    WeightObfuscator,
    ObfuscationConfig,
    ObfuscationVerifier,
)
from model_osmosis.obfuscation.graph import LayerGraph


# ---------------------------------------------------------------------------
# Test model — big enough to exercise all obfuscation paths
# ---------------------------------------------------------------------------

class _TransformerLikeModel(nn.Module):
    """Mimics a tiny transformer with named layers that the graph analyzer
    can detect: attention projections (q/k/v/o) and FFN (fc1/fc2)."""

    def __init__(self, hidden=32, intermediate=64):
        super().__init__()
        # Embedding
        self.embed_tokens = nn.Embedding(128, hidden)
        # Layer 0 — attention
        self.layers_0_self_attn_q_proj = nn.Linear(hidden, hidden)
        self.layers_0_self_attn_k_proj = nn.Linear(hidden, hidden)
        self.layers_0_self_attn_v_proj = nn.Linear(hidden, hidden)
        self.layers_0_self_attn_o_proj = nn.Linear(hidden, hidden)
        # Layer 0 — FFN
        self.layers_0_mlp_fc1 = nn.Linear(hidden, intermediate)
        self.layers_0_mlp_fc2 = nn.Linear(intermediate, hidden)
        # Output
        self.lm_head = nn.Linear(hidden, 128, bias=False)

    def forward(self, x):
        h = self.embed_tokens(x)
        # Simplified attention (just projections, no actual attention)
        q = self.layers_0_self_attn_q_proj(h)
        k = self.layers_0_self_attn_k_proj(h)
        v = self.layers_0_self_attn_v_proj(h)
        attn_out = self.layers_0_self_attn_o_proj(q + k + v)
        h = h + attn_out
        # FFN
        ffn = self.layers_0_mlp_fc2(torch.relu(self.layers_0_mlp_fc1(h)))
        h = h + ffn
        return self.lm_head(h)


def _make_transformer(seed=42):
    torch.manual_seed(seed)
    model = _TransformerLikeModel(hidden=32, intermediate=64)
    model.eval()
    return model


def _make_simple_model(seed=0):
    """Simple sequential model for basic tests."""
    torch.manual_seed(seed)
    model = nn.Sequential(
        nn.Linear(8, 32),
        nn.ReLU(),
        nn.Linear(32, 16),
        nn.ReLU(),
        nn.Linear(16, 4),
    )
    model.eval()
    return model


# ---------------------------------------------------------------------------
# Full Pipeline: Build -> Save -> Load -> Inspect -> Export
# ---------------------------------------------------------------------------

class TestFullPipeline:
    """End-to-end: build an OsmoImage, save, load, inspect, export."""

    def test_build_save_load_roundtrip(self):
        model = _make_simple_model()
        image = OsmoImage.build(model, name="e2e-test")

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

            # File must exist with correct magic
            assert saved.exists()
            assert is_osmoimg(saved)
            with open(saved, "rb") as f:
                assert f.read(8) == OSMOIMG_MAGIC

            # Load and verify
            loaded = OsmoImage.load(saved)
            assert loaded.name == "e2e-test"
            assert loaded.parameter_count == image.parameter_count
            assert loaded.runtime.preserve_quality is True

    def test_inspect_metadata(self):
        model = _make_simple_model()
        image = OsmoImage.build(model, name="inspect-e2e")

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

            info = OsmoImage.inspect(path)
            assert info["name"] == "inspect-e2e"
            assert info["parameter_count"] > 0
            assert info["runtime"]["preserve_quality"] is True
            assert isinstance(info["image"]["file_size_mb"], float)
            assert info["image"]["file_size_mb"] > 0

    def test_export_returns_state_dict(self):
        model = _make_simple_model()
        image = OsmoImage.build(model, name="export-e2e")
        exported = image.export(target_format="pytorch")
        assert isinstance(exported, dict)
        assert len(exported) > 0

    def test_compression_reduces_size(self):
        model = _make_simple_model()
        image = OsmoImage.build(model, name="compress-e2e")

        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "model.osmoimg"
            saved = image.save(path)
            file_size = saved.stat().st_size
            # Compressed file should exist and be non-empty
            assert file_size > 0


# ---------------------------------------------------------------------------
# Obfuscation: Weights actually change + functional equivalence
# ---------------------------------------------------------------------------

class TestObfuscationCorrectness:
    """Verify obfuscation transforms weights while preserving behavior."""

    def test_weights_are_different_after_obfuscation(self):
        """Obfuscated weights must NOT be identical to originals."""
        model = _make_simple_model()
        converter = UniversalConverter()
        osmo = converter.to_osmo(model, name="diff-test")

        original_params = {k: v.copy() for k, v in osmo.parameters.items()}

        config = ObfuscationConfig(
            enable_permutation=True,
            enable_scaling=True,
            enable_rotation=False,
            preserve_quality=True,
            seed=42,
            verify=False,  # skip verification for this test
            obfuscate_param_names=False,  # keep names so we can compare
        )
        obfuscator = WeightObfuscator(config)
        obfuscated = obfuscator.obfuscate(osmo)

        # At least some weight matrices must have changed
        any_different = False
        for key in original_params:
            if key in obfuscated.parameters:
                orig = original_params[key]
                obf = obfuscated.parameters[key]
                if orig.shape == obf.shape and not np.allclose(orig, obf):
                    any_different = True
                    break

        assert any_different, "Obfuscation did not change any weights"

    def test_functional_equivalence_simple(self):
        """Obfuscated model produces same outputs as original."""
        model = _make_simple_model()
        converter = UniversalConverter()
        osmo = converter.to_osmo(model, name="equiv-test")

        # Generate sample data
        rng = np.random.default_rng(0)
        sample_data = rng.standard_normal((20, 8)).astype(np.float32)

        # Forward pass on original
        original_out = self._forward(osmo.parameters, sample_data)

        # Obfuscate with permutation + scaling (exact transforms)
        config = ObfuscationConfig(
            enable_permutation=True,
            enable_scaling=True,
            enable_rotation=False,
            preserve_quality=True,
            seed=42,
            verify=False,
            obfuscate_param_names=False,
        )
        obfuscator = WeightObfuscator(config)
        obfuscated = obfuscator.obfuscate(osmo)

        # Forward pass on obfuscated
        obfuscated_out = self._forward(obfuscated.parameters, sample_data)

        # Must be very close (exact transforms = zero error in theory,
        # but floating point gives tiny differences)
        np.testing.assert_allclose(
            original_out, obfuscated_out,
            atol=1e-5, rtol=1e-5,
            err_msg="Obfuscated model output differs from original",
        )

    def _forward(self, params: dict, x: np.ndarray) -> np.ndarray:
        """Manual forward pass through a Sequential(Linear, ReLU, ...) model."""
        # Find weight/bias pairs in order
        layers = []
        i = 0
        while True:
            wkey = f"{i}.weight"
            bkey = f"{i}.bias"
            if wkey in params:
                layers.append((params[wkey], params.get(bkey)))
                i += 1
            else:
                i += 1
                if i > 10:
                    break
                continue

        out = x.astype(np.float64)
        for idx, (w, b) in enumerate(layers):
            out = out @ w.astype(np.float64).T
            if b is not None:
                out = out + b.astype(np.float64)
            # ReLU between layers (not after last)
            if idx < len(layers) - 1:
                out = np.maximum(out, 0)

        return out.astype(np.float32)

    def test_verification_passes(self):
        """Built-in verifier should pass for exact transforms."""
        model = _make_simple_model()
        converter = UniversalConverter()
        osmo = converter.to_osmo(model, name="verify-test")

        config = ObfuscationConfig(
            enable_permutation=True,
            enable_scaling=True,
            enable_rotation=False,
            preserve_quality=True,
            seed=42,
            verify=True,
            verification_tolerance=1e-5,
        )
        obfuscator = WeightObfuscator(config)
        # Should not raise
        obfuscated = obfuscator.obfuscate(osmo)
        assert obfuscated is not None


# ---------------------------------------------------------------------------
# Graph Analysis: Correct layer pairing
# ---------------------------------------------------------------------------

class TestGraphAnalysis:
    """Verify the layer graph correctly identifies transform groups."""

    def test_sequential_model_finds_pairs(self):
        model = _make_simple_model()
        converter = UniversalConverter()
        osmo = converter.to_osmo(model, name="graph-test")

        graph = LayerGraph(osmo.parameters, osmo.architecture)
        groups = graph.build_transform_groups()

        # Sequential(Linear(8,32), ReLU, Linear(32,16), ReLU, Linear(16,4))
        # Should find pairs: 0->2 (dim=32) and 2->4 (dim=16)
        assert len(groups) >= 1, "Should find at least one transform group"

        # Verify dimensions make sense
        for g in groups:
            assert g.dim_size > 0
            w = osmo.parameters[g.weight_key]
            assert w.ndim >= 2

    def test_transformer_finds_attention_and_ffn(self):
        """Transformer-like model should find attention + FFN pairs."""
        model = _make_transformer()
        converter = UniversalConverter()
        osmo = converter.to_osmo(model, name="transformer-graph")

        graph = LayerGraph(osmo.parameters, osmo.architecture)
        groups = graph.build_transform_groups()

        group_names = [g.layer_name for g in groups]

        # Should find FFN pair (fc1 -> fc2)
        has_ffn = any("ffn" in name.lower() for name in group_names)
        # Should find at least some groups
        assert len(groups) >= 1, f"Expected groups, got: {group_names}"

    def test_deduplication(self):
        """No weight key should appear in more than one group."""
        model = _make_transformer()
        converter = UniversalConverter()
        osmo = converter.to_osmo(model, name="dedup-test")

        graph = LayerGraph(osmo.parameters, osmo.architecture)
        groups = graph.build_transform_groups()

        seen_weight = set()
        seen_next = set()
        for g in groups:
            assert g.weight_key not in seen_weight, f"Duplicate weight_key: {g.weight_key}"
            assert g.next_weight_key not in seen_next, f"Duplicate next_weight_key: {g.next_weight_key}"
            seen_weight.add(g.weight_key)
            seen_next.add(g.next_weight_key)


# ---------------------------------------------------------------------------
# Metadata Obfuscation
# ---------------------------------------------------------------------------

class TestMetadataObfuscation:
    """Verify metadata, param names, and architecture are obfuscated."""

    def test_param_names_hashed(self):
        model = _make_simple_model()
        converter = UniversalConverter()
        osmo = converter.to_osmo(model, name="names-test")

        config = ObfuscationConfig(
            enable_permutation=False,
            enable_scaling=False,
            enable_rotation=False,
            obfuscate_param_names=True,
            seed=42,
            verify=False,
        )
        obfuscator = WeightObfuscator(config)
        obfuscated = obfuscator.obfuscate(osmo)

        # All param keys should be P_<hash> format
        for key in obfuscated.parameters:
            assert key.startswith("P_"), f"Expected hashed name, got: {key}"

    def test_param_name_hashing_is_deterministic(self):
        model = _make_simple_model()
        converter = UniversalConverter()

        config = ObfuscationConfig(
            enable_permutation=False,
            enable_scaling=False,
            enable_rotation=False,
            obfuscate_param_names=True,
            seed=42,
            verify=False,
        )

        osmo1 = converter.to_osmo(model, name="det-test")
        obf1 = WeightObfuscator(config).obfuscate(osmo1)

        osmo2 = converter.to_osmo(model, name="det-test")
        obf2 = WeightObfuscator(config).obfuscate(osmo2)

        assert list(obf1.parameters.keys()) == list(obf2.parameters.keys())

    def test_manifest_metadata_stripped(self):
        model = _make_simple_model()
        converter = UniversalConverter()
        osmo = converter.to_osmo(model, name="meta-test")
        osmo.manifest.original_repo = "secret/repo"
        osmo.manifest.training_info.base_model = "secret-base"

        config = ObfuscationConfig(
            enable_permutation=False,
            enable_scaling=False,
            enable_rotation=False,
            randomize_metadata=True,
            seed=42,
            verify=False,
        )
        obfuscator = WeightObfuscator(config)
        obfuscated = obfuscator.obfuscate(osmo)

        assert obfuscated.manifest.original_repo is None
        assert obfuscated.manifest.training_info.base_model is None


# ---------------------------------------------------------------------------
# Progressive Fallback
# ---------------------------------------------------------------------------

class TestProgressiveFallback:
    """OsmoImage.build should gracefully degrade if obfuscation fails."""

    def test_build_always_succeeds(self):
        """Even if transforms fail, build should produce a valid image."""
        model = _make_simple_model()
        image = OsmoImage.build(model, name="fallback-test")
        assert image is not None
        assert image.parameter_count > 0
        assert image.runtime.preserve_quality is True

    def test_build_save_load_cycle(self):
        """Full cycle must work end-to-end."""
        model = _make_simple_model()
        image = OsmoImage.build(model, name="cycle-test")

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

            # Parameter counts must match
            assert loaded.parameter_count == image.parameter_count

            # Runtime metadata must survive
            assert loaded.runtime.source_framework == image.runtime.source_framework
            assert loaded.runtime.preserve_quality == image.runtime.preserve_quality
