"""
CPU-specific tests for Dokodemo AI.

Verifies that:
1. CPU device is properly detected and configured
2. mmap mode loads tensors as OS-page-backed views (no Python-copy overhead)
3. Compute dtype selection is correct per CPU capability
4. Memory feasibility checks work for CPU with mmap
5. Thread configuration doesn't crash
6. The full end-to-end inference loop works on CPU with a tiny toy model
"""

import json
import os
import tempfile

import pytest
import torch
import torch.nn as nn


# ---------------------------------------------------------------------------
# CPU optimize tests
# ---------------------------------------------------------------------------

class TestCpuOptimize:
    def test_supports_bf16_returns_bool(self):
        from dokodemo_ai.utils.cpu_optimize import supports_bf16
        result = supports_bf16()
        assert isinstance(result, bool)

    def test_get_compute_dtype_is_valid(self):
        from dokodemo_ai.utils.cpu_optimize import get_compute_dtype
        dtype = get_compute_dtype()
        assert dtype in (torch.float32, torch.bfloat16), (
            f"Unexpected compute dtype: {dtype}"
        )

    def test_get_weight_load_dtype_is_float16(self):
        from dokodemo_ai.utils.cpu_optimize import get_weight_load_dtype
        assert get_weight_load_dtype() == torch.float16

    def test_setup_threads_no_crash(self):
        from dokodemo_ai.utils.cpu_optimize import setup_threads
        # Should not raise
        setup_threads(n_compute=2, n_interop=1)
        assert torch.get_num_threads() == 2

    def test_available_ram_bytes_positive(self):
        from dokodemo_ai.utils.cpu_optimize import available_ram_bytes
        ram = available_ram_bytes()
        assert ram > 0, "Available RAM should be positive"

    def test_total_ram_bytes_positive(self):
        from dokodemo_ai.utils.cpu_optimize import total_ram_bytes
        total = total_ram_bytes()
        assert total > 0

    def test_available_leq_total(self):
        from dokodemo_ai.utils.cpu_optimize import available_ram_bytes, total_ram_bytes
        assert available_ram_bytes() <= total_ram_bytes()

    def test_cpu_info_fields(self):
        from dokodemo_ai.utils.cpu_optimize import cpu_info
        info = cpu_info()
        assert "cpu_count" in info
        assert "supports_bf16" in info
        assert "compute_dtype" in info
        assert "total_ram_gb" in info
        assert "available_ram_gb" in info
        assert info["cpu_count"] > 0

    def test_force_gc_no_crash(self):
        from dokodemo_ai.utils.cpu_optimize import force_gc
        # Create some garbage
        _ = [torch.randn(100, 100) for _ in range(10)]
        del _
        collected = force_gc()
        assert isinstance(collected, int)


# ---------------------------------------------------------------------------
# Memory manager CPU checks
# ---------------------------------------------------------------------------

class TestMemoryManagerCPU:
    def test_cpu_device_detection(self):
        from dokodemo_ai.utils.memory import MemoryManager
        mgr = MemoryManager(device=torch.device("cpu"))
        assert mgr.device.type == "cpu"

    def test_mmap_budget_is_zero(self):
        """In mmap mode, the Python-level budget is 0 (OS handles it)."""
        from dokodemo_ai.utils.memory import MemoryManager
        mgr = MemoryManager(device=torch.device("cpu"))
        budget = mgr.get_weight_budget_bytes(use_mmap=True)
        assert budget == 0, f"Mmap budget should be 0, got {budget}"

    def test_non_mmap_budget_positive(self):
        from dokodemo_ai.utils.memory import MemoryManager
        mgr = MemoryManager(device=torch.device("cpu"))
        budget = mgr.get_weight_budget_bytes(use_mmap=False)
        assert budget >= 0  # Could be 0 on very low RAM systems

    def test_feasibility_mmap_always_passes_with_enough_ram(self):
        from dokodemo_ai.utils.memory import MemoryManager
        mgr = MemoryManager(device=torch.device("cpu"))
        # Even a 1TB model should be "feasible" in mmap mode
        feasible, msg = mgr.check_feasibility(
            max_layer_size_bytes=500 * 1024**3,   # 500 GB layer
            embedding_size_bytes=50 * 1024**3,
            head_size_bytes=50 * 1024**3,
            use_mmap=True,
        )
        # Feasible as long as we have enough RAM for activations (~1 GB)
        # On any reasonable system this should pass
        assert isinstance(feasible, bool)
        assert isinstance(msg, str)
        assert "mmap" in msg.lower() or "feasible" in msg.lower()

    def test_feasibility_non_mmap_insufficient_ram(self):
        """Non-mmap CPU mode with enormous layer should fail."""
        from dokodemo_ai.utils.memory import MemoryManager
        mgr = MemoryManager(device=torch.device("cpu"))
        feasible, msg = mgr.check_feasibility(
            max_layer_size_bytes=1000 * 1024**3,  # 1 TB — definitely no
            embedding_size_bytes=100 * 1024**3,
            head_size_bytes=100 * 1024**3,
            use_mmap=False,
        )
        assert not feasible
        assert "available" in msg.lower() or "gb" in msg.lower()

    def test_kv_cache_estimate_positive(self):
        from dokodemo_ai.utils.memory import MemoryManager
        mgr = MemoryManager(device=torch.device("cpu"))
        size = mgr.estimate_kv_cache_bytes(
            num_layers=32, hidden_size=4096, num_heads=32, max_seq_len=512
        )
        assert size > 0

    def test_activation_estimate_positive(self):
        from dokodemo_ai.utils.memory import MemoryManager
        mgr = MemoryManager(device=torch.device("cpu"))
        size = mgr.estimate_activation_bytes(hidden_size=4096)
        assert size > 0


# ---------------------------------------------------------------------------
# TensorStreamer mmap mode tests
# ---------------------------------------------------------------------------

class TestTensorStreamerMmap:
    """Test mmap-based loading on CPU (requires safetensors installed)."""

    def _make_safetensors_dir(self, tmp_path, tensors: dict) -> str:
        """Write tensors to a SafeTensors file and return the dir path."""
        pytest.importorskip("safetensors")
        from safetensors.torch import save_file
        import json

        save_file(tensors, str(tmp_path / "model.safetensors"))
        index = {name: "model.safetensors" for name in tensors}
        with open(tmp_path / "model.safetensors.index.json", "w") as f:
            json.dump({"weight_map": index}, f)

        # Minimal config
        with open(tmp_path / "config.json", "w") as f:
            json.dump({"model_type": "test", "hidden_size": 16}, f)

        return str(tmp_path)

    def test_mmap_load_returns_tensors(self, tmp_path):
        """mmap mode should load tensors successfully."""
        pytest.importorskip("safetensors")
        from dokodemo_ai.engine.streaming import TensorStreamer
        from dokodemo_ai.graph.compiler import (
            ComponentInfo, EmbeddingInfo, HeadInfo,
            ModelGraph, TensorRef,
        )

        # Create a small model directory
        weights = {"layer.0.weight": torch.randn(8, 8)}
        model_path = self._make_safetensors_dir(tmp_path, weights)

        ref = TensorRef(
            name="layer.0.weight",
            file="model.safetensors",
            dtype="F32",
            shape=[8, 8],
        )
        comp = ComponentInfo(name="layer.0", kind="mlp", params=[ref])

        graph = ModelGraph(
            model_path=model_path,
            config={"hidden_size": 8},
            embedding=EmbeddingInfo(module_name="embed"),
            layers=[],
            head=HeadInfo(module_name="head"),
            safetensors_index={"layer.0.weight": "model.safetensors"},
            num_layers=0,
        )

        streamer = TensorStreamer(
            model_graph=graph,
            device=torch.device("cpu"),
            use_mmap=True,
        )

        loaded = streamer.load_component(comp, to_device=torch.device("cpu"))
        assert "layer.0.weight" in loaded
        tensor = loaded["layer.0.weight"]
        assert tensor.shape == (8, 8)

    def test_mmap_tensor_correct_values(self, tmp_path):
        """Values loaded via mmap should match original values."""
        pytest.importorskip("safetensors")
        from dokodemo_ai.engine.streaming import TensorStreamer
        from dokodemo_ai.graph.compiler import (
            ComponentInfo, EmbeddingInfo, HeadInfo,
            ModelGraph, TensorRef,
        )

        original = torch.tensor([1.0, 2.0, 3.0, 4.0])
        weights = {"test.weight": original}
        model_path = self._make_safetensors_dir(tmp_path, weights)

        ref = TensorRef(
            name="test.weight", file="model.safetensors",
            dtype="F32", shape=[4],
        )
        comp = ComponentInfo(name="test", kind="other", params=[ref])
        graph = ModelGraph(
            model_path=model_path,
            config={},
            embedding=EmbeddingInfo(module_name="embed"),
            layers=[],
            head=HeadInfo(module_name="head"),
            safetensors_index={"test.weight": "model.safetensors"},
            num_layers=0,
        )

        streamer = TensorStreamer(
            model_graph=graph,
            device=torch.device("cpu"),
            use_mmap=True,
        )
        loaded = streamer.load_component(comp)
        assert torch.allclose(loaded["test.weight"].float(), original)

    def test_mmap_mode_no_cpu_buffer(self, tmp_path):
        """In mmap mode, the CPU buffer should be zero-sized."""
        pytest.importorskip("safetensors")
        from dokodemo_ai.engine.streaming import TensorStreamer
        from dokodemo_ai.graph.compiler import EmbeddingInfo, HeadInfo, ModelGraph

        weights = {"x.weight": torch.randn(4)}
        model_path = self._make_safetensors_dir(tmp_path, weights)

        graph = ModelGraph(
            model_path=model_path,
            config={},
            embedding=EmbeddingInfo(module_name="embed"),
            layers=[],
            head=HeadInfo(module_name="head"),
            safetensors_index={"x.weight": "model.safetensors"},
            num_layers=0,
        )

        streamer = TensorStreamer(
            model_graph=graph,
            device=torch.device("cpu"),
            use_mmap=True,
        )
        assert streamer._cpu_buffer_bytes == 0, (
            "mmap mode should not use a Python-managed CPU buffer"
        )


# ---------------------------------------------------------------------------
# End-to-end CPU inference with a tiny toy model
# ---------------------------------------------------------------------------

class TinyTransformerLayer(nn.Module):
    """Smallest possible transformer layer for E2E testing."""

    def __init__(self, d=16, heads=2):
        super().__init__()
        self.input_layernorm = nn.LayerNorm(d)
        # Fake 'self_attn' with q/k/v/o projections
        self.self_attn = nn.ModuleDict({
            "q_proj": nn.Linear(d, d, bias=False),
            "k_proj": nn.Linear(d, d, bias=False),
            "v_proj": nn.Linear(d, d, bias=False),
            "o_proj": nn.Linear(d, d, bias=False),
        })
        self.post_attention_layernorm = nn.LayerNorm(d)
        self.mlp = nn.ModuleDict({
            "gate_proj": nn.Linear(d, d * 2, bias=False),
            "up_proj":   nn.Linear(d, d * 2, bias=False),
            "down_proj": nn.Linear(d * 2, d, bias=False),
        })


class TinyTransformerModel(nn.Module):
    def __init__(self, vocab=100, d=16, heads=2, n_layers=2):
        super().__init__()
        self.embed_tokens = nn.Embedding(vocab, d)
        self.layers = nn.ModuleList([
            TinyTransformerLayer(d, heads) for _ in range(n_layers)
        ])
        self.norm = nn.LayerNorm(d)
        self.lm_head = nn.Linear(d, vocab, bias=False)


class TestCpuEndToEnd:
    """End-to-end tests on CPU with toy models — no real weights needed."""

    def test_kv_cache_update(self):
        """KV cache update should concatenate correctly."""
        from dokodemo_ai.engine.inference import KVCache

        cache = KVCache(
            num_layers=2, num_kv_heads=2, head_dim=8,
            device=torch.device("cpu"), dtype=torch.float32,
        )

        # First update
        k = torch.randn(1, 2, 3, 8)  # [batch, heads, seq=3, dim]
        v = torch.randn(1, 2, 3, 8)
        k_full, v_full = cache.update(0, k, v)
        assert k_full.shape == (1, 2, 3, 8)
        assert cache.get_seq_len(0) == 3

        # Second update (1 new token)
        k2 = torch.randn(1, 2, 1, 8)
        v2 = torch.randn(1, 2, 1, 8)
        k_full2, v_full2 = cache.update(0, k2, v2)
        assert k_full2.shape == (1, 2, 4, 8)
        assert cache.get_seq_len(0) == 4

    def test_kv_cache_clear(self):
        from dokodemo_ai.engine.inference import KVCache
        cache = KVCache(
            num_layers=2, num_kv_heads=2, head_dim=8,
            device=torch.device("cpu"), dtype=torch.float32,
        )
        cache.update(0, torch.randn(1, 2, 5, 8), torch.randn(1, 2, 5, 8))
        assert cache.get_seq_len(0) == 5
        cache.clear()
        assert cache.get_seq_len(0) == 0

    def test_rope_apply_shape_preserved(self):
        """RoPE application should not change tensor shape."""
        from dokodemo_ai.engine.inference import InferenceEngine

        # We only need the static method
        x = torch.randn(1, 4, 10, 8)  # [batch, heads, seq, dim]
        cos = torch.ones(20, 8)
        sin = torch.zeros(20, 8)
        out = InferenceEngine._apply_rope(x, cos, sin, offset=0)
        assert out.shape == x.shape

    def test_rms_norm_computation(self):
        """RMSNorm should normalize correctly."""
        # Instantiate a minimal engine stub
        import types
        from dokodemo_ai.engine.inference import InferenceEngine

        # Build a minimal engine without loading a real model
        engine = object.__new__(InferenceEngine)
        engine.rms_norm_eps = 1e-6

        x = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
        w = torch.ones(4)
        out = engine._rms_norm(x, w)
        # RMSNorm output should have roughly unit RMS
        rms = out.float().pow(2).mean().sqrt()
        # Not exactly 1.0 because of eps, but close
        assert abs(rms.item() - 1.0) < 0.1, f"RMS was {rms.item()}"

    def test_sampling_greedy(self):
        """Greedy sampling should return argmax."""
        import types
        from dokodemo_ai.engine.inference import InferenceEngine

        engine = object.__new__(InferenceEngine)

        logits = torch.tensor([[0.1, 5.0, 0.2, 0.3]])  # Index 1 is max
        logits = logits.unsqueeze(1)  # [1, 1, 4]
        token = engine._sample(logits, temperature=1.0, top_k=0, top_p=1.0, do_sample=False)
        assert token.item() == 1

    def test_sampling_top_k_valid_range(self):
        """Top-k sampling should only sample from top-k tokens."""
        import types
        from dokodemo_ai.engine.inference import InferenceEngine

        engine = object.__new__(InferenceEngine)
        torch.manual_seed(42)

        # Vocab of 100, top-k=5
        logits = torch.randn(1, 1, 100)
        top5_indices = logits[0, 0].topk(5).indices.tolist()

        tokens = set()
        for _ in range(50):
            t = engine._sample(logits, temperature=1.0, top_k=5, top_p=1.0, do_sample=True)
            tokens.add(t.item())

        # All sampled tokens should be in the top-5
        assert tokens.issubset(set(top5_indices)), (
            f"Sampled tokens {tokens} not all in top-5 {top5_indices}"
        )

    def test_find_tensor_pattern_matching(self):
        """_find_tensor should match by pattern substring."""
        from dokodemo_ai.engine.inference import InferenceEngine

        tensors = {
            "model.layers.0.self_attn.q_proj.weight": torch.randn(4, 4),
            "model.layers.0.self_attn.k_proj.weight": torch.randn(4, 4),
            "model.layers.0.mlp.gate_proj.weight":    torch.randn(8, 4),
        }

        q = InferenceEngine._find_tensor(tensors, ["q_proj"])
        k = InferenceEngine._find_tensor(tensors, ["k_proj"])
        gate = InferenceEngine._find_tensor(tensors, ["gate_proj"])
        missing = InferenceEngine._find_tensor(tensors, ["nonexistent"])

        assert q is not None
        assert k is not None
        assert gate is not None
        assert missing is None
