#!/usr/bin/env python3
"""
test_cdna_kernel.py - Unit Tests for CDNA Kernel and FibPi Agent

Tests:
1. Kernel primitives (address, fetch, verify)
2. Kernel compute (matmul correctness)
3. Agent caching behavior
4. Agent verify policies

Work Order: WO-CDNA-KERNEL-TESTS-V1
"""

import json
import os
import sys
import tempfile
from pathlib import Path

import numpy as np
import pytest

# Add parent to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from helix_cdc.kernel.cdna_kernel import CDNAKernel, KernelReceipt
from helix_cdc.agents.fibpi_chunk_agent import (
    FibPiChunkAgent,
    AccessReceipt,
    LRUCache,
)
from helix_cdc.regrow.cdna_stream_v2 import CDNAv2Writer


# ========== Fixtures ==========

@pytest.fixture
def sample_cdna_file(tmp_path):
    """Create a small CDNAv2 file for testing."""
    # Create synthetic tensor (64 rows x 128 cols)
    rows, cols = 64, 128
    indices = np.random.randint(0, 256, size=(rows, cols), dtype=np.uint8)
    codebook = np.random.randn(256).astype(np.float32) * 0.1

    # Write CDNAv2 file
    output_path = tmp_path / "test_tensor.cdna2.hxz"
    writer = CDNAv2Writer()
    receipt = writer.write(
        indices=indices,
        codebook=codebook,
        output_path=output_path,
        tensor_name="test_weight",
        block_rows=16,
    )

    return {
        "path": output_path,
        "indices": indices,
        "codebook": codebook,
        "rows": rows,
        "cols": cols,
        "block_rows": 16,
        "num_blocks": 4,  # 64 / 16 = 4
    }


@pytest.fixture
def real_cdna_file():
    """Use a real CDNAv2 file if available."""
    test_file = Path("/home/voidstr3m33/helix-cdc/seeds/cdna_v2_sweep/blk.0.attn_k.weight.br16.cdna2.hxz")
    if not test_file.exists():
        pytest.skip("Real CDNA file not available")
    return test_file


# ========== Kernel Tests ==========

class TestCDNAKernel:
    """Tests for CDNAKernel."""

    def test_kernel_init(self, sample_cdna_file):
        """Kernel initializes correctly."""
        k = CDNAKernel(sample_cdna_file["path"])

        assert k.shape == (sample_cdna_file["rows"], sample_cdna_file["cols"])
        assert k.rows == sample_cdna_file["rows"]
        assert k.cols == sample_cdna_file["cols"]
        assert k.num_blocks == sample_cdna_file["num_blocks"]
        assert k.block_rows == sample_cdna_file["block_rows"]
        assert len(k.codebook) == 256

    def test_kernel_address(self, sample_cdna_file):
        """address() maps row ranges to block tuples."""
        k = CDNAKernel(sample_cdna_file["path"])

        # Test single block
        blocks = k.address(0, 10)
        assert len(blocks) == 1
        assert blocks[0] == (0, 0, 10)

        # Test across blocks
        blocks = k.address(10, 40)
        assert len(blocks) == 3
        # Block 0: rows 0-15, need 10-15 = local 10-16
        assert blocks[0] == (0, 10, 16)
        # Block 1: rows 16-31, need all = local 0-16
        assert blocks[1] == (1, 0, 16)
        # Block 2: rows 32-47, need 32-39 = local 0-8
        assert blocks[2] == (2, 0, 8)

    def test_kernel_address_edge_cases(self, sample_cdna_file):
        """address() handles edge cases correctly."""
        k = CDNAKernel(sample_cdna_file["path"])

        # Full tensor
        blocks = k.address(0, k.rows)
        assert len(blocks) == k.num_blocks

        # Invalid ranges
        with pytest.raises(ValueError):
            k.address(-1, 10)
        with pytest.raises(ValueError):
            k.address(10, 5)  # start > end
        with pytest.raises(ValueError):
            k.address(0, k.rows + 1)  # beyond tensor

    def test_kernel_fetch_block(self, sample_cdna_file):
        """fetch_block() returns decompressed indices."""
        k = CDNAKernel(sample_cdna_file["path"])

        indices = k.fetch_block(0)
        expected = sample_cdna_file["indices"][0:16]

        assert indices.shape == (16, sample_cdna_file["cols"])
        assert indices.dtype == np.uint8
        assert np.array_equal(indices, expected)

    def test_kernel_fetch_block_invalid(self, sample_cdna_file):
        """fetch_block() rejects invalid block IDs."""
        k = CDNAKernel(sample_cdna_file["path"])

        with pytest.raises(ValueError):
            k.fetch_block(-1)
        with pytest.raises(ValueError):
            k.fetch_block(k.num_blocks)

    def test_kernel_verify_block(self, sample_cdna_file):
        """verify_block() returns True for valid blocks."""
        k = CDNAKernel(sample_cdna_file["path"])

        for block_id in range(k.num_blocks):
            assert k.verify_block(block_id) is True

    def test_kernel_get_rows(self, sample_cdna_file):
        """get_rows() returns correct rows with verification."""
        k = CDNAKernel(sample_cdna_file["path"])

        # Test single block range
        rows = k.get_rows(5, 15, verify=True)
        expected = sample_cdna_file["indices"][5:15]
        assert np.array_equal(rows, expected)

        # Test cross-block range
        rows = k.get_rows(10, 40, verify=True)
        expected = sample_cdna_file["indices"][10:40]
        assert np.array_equal(rows, expected)

    def test_kernel_matmul_correctness(self, sample_cdna_file):
        """matmul via kernel matches dense matmul."""
        k = CDNAKernel(sample_cdna_file["path"])

        x = np.random.randn(k.cols).astype(np.float32)

        # Kernel matmul
        y_kernel, receipt = k.stream_matmul(x, verify=False)

        # Dense matmul
        indices = k.get_rows(0, k.rows, verify=False)
        W_dense = k.codebook[indices.flatten()].reshape(k.shape)
        y_dense = W_dense @ x

        # Compare (should be exact or very close)
        cosine = np.dot(y_kernel, y_dense) / (np.linalg.norm(y_kernel) * np.linalg.norm(y_dense))
        assert cosine > 0.9999, f"Cosine similarity {cosine} too low"

    def test_kernel_matmul_block(self, sample_cdna_file):
        """matmul_block() computes correct block result."""
        k = CDNAKernel(sample_cdna_file["path"])

        x = np.random.randn(k.cols).astype(np.float32)
        y_block = k.matmul_block(0, x)

        # Verify against dense
        indices = k.fetch_block(0)
        W_block = k.codebook[indices.flatten()].reshape(indices.shape)
        y_expected = W_block @ x

        assert np.allclose(y_block, y_expected, rtol=1e-5)

    def test_kernel_iter_blocks(self, sample_cdna_file):
        """iter_blocks() yields all blocks with receipts."""
        k = CDNAKernel(sample_cdna_file["path"])

        blocks_seen = []
        for indices, start_row, receipt in k.iter_blocks(verify=True):
            blocks_seen.append((indices.shape[0], start_row))
            assert receipt.verified is True

        assert len(blocks_seen) == k.num_blocks

    def test_kernel_cache_control(self, sample_cdna_file):
        """Cache enable/disable works correctly."""
        k = CDNAKernel(sample_cdna_file["path"])

        assert k.cache_info["enabled"] is False
        assert k.cache_info["blocks_cached"] == 0

        # Fetch without cache
        k.fetch_block(0)
        assert k.cache_info["blocks_cached"] == 0

        # Enable cache
        k.enable_cache()
        assert k.cache_info["enabled"] is True

        # Fetch with cache
        k.fetch_block(0)
        assert k.cache_info["blocks_cached"] == 1

        # Clear cache
        k.clear_cache()
        assert k.cache_info["blocks_cached"] == 0

        # Disable cache
        k.disable_cache()
        assert k.cache_info["enabled"] is False


# ========== Agent Tests ==========

class TestFibPiChunkAgent:
    """Tests for FibPiChunkAgent."""

    def test_agent_init(self):
        """Agent initializes with correct defaults."""
        agent = FibPiChunkAgent()
        assert agent.verify_policy == "trust_cached"

        agent = FibPiChunkAgent(verify_policy="always")
        assert agent.verify_policy == "always"

    def test_agent_open_capsule(self, sample_cdna_file):
        """Agent opens capsule and returns kernel."""
        agent = FibPiChunkAgent()
        kernel = agent.open_capsule(sample_cdna_file["path"])

        assert isinstance(kernel, CDNAKernel)
        assert kernel.shape == (sample_cdna_file["rows"], sample_cdna_file["cols"])

    def test_agent_cache_hits(self, sample_cdna_file):
        """Agent caches blocks and reports hits."""
        agent = FibPiChunkAgent(verify_policy="trust_cached")
        kernel = agent.open_capsule(sample_cdna_file["path"])

        # First access (miss)
        _, r1 = agent.fetch_block(kernel, 0)
        assert r1.cache_hit is False

        # Second access (hit)
        _, r2 = agent.fetch_block(kernel, 0)
        assert r2.cache_hit is True

        # Stats should reflect
        stats = agent.cache_stats()
        assert stats["cache_hits"] >= 1
        assert stats["hit_rate"] > 0

    def test_agent_verify_policy_always(self, sample_cdna_file):
        """'always' policy verifies every access."""
        agent = FibPiChunkAgent(verify_policy="always")
        kernel = agent.open_capsule(sample_cdna_file["path"])

        # Both accesses should verify
        _, r1 = agent.fetch_block(kernel, 0)
        _, r2 = agent.fetch_block(kernel, 0)

        stats = agent.cache_stats()
        assert stats["blocks_verified"] == 2

    def test_agent_verify_policy_trust_cached(self, sample_cdna_file):
        """'trust_cached' policy only verifies first access."""
        agent = FibPiChunkAgent(verify_policy="trust_cached")
        kernel = agent.open_capsule(sample_cdna_file["path"])

        # First access (should verify)
        _, r1 = agent.fetch_block(kernel, 0)
        assert r1.verified is True

        # Second access (should trust cache)
        _, r2 = agent.fetch_block(kernel, 0)
        assert r2.cache_hit is True

        stats = agent.cache_stats()
        assert stats["blocks_verified"] == 1  # Only first access verified

    def test_agent_verify_policy_sampled(self, sample_cdna_file):
        """'sampled' policy verifies subset of accesses."""
        # Set seed for reproducibility
        import random
        random.seed(42)

        agent = FibPiChunkAgent(verify_policy="sampled", sample_rate=0.5)
        kernel = agent.open_capsule(sample_cdna_file["path"])

        # Access all blocks
        for block_id in range(kernel.num_blocks):
            agent.fetch_block(kernel, block_id)

        stats = agent.cache_stats()
        # With 50% sample rate, expect roughly half verified
        # (but due to randomness, just check it's not 0 or all)
        assert 0 <= stats["blocks_verified"] <= kernel.num_blocks

    def test_agent_fetch_rows(self, sample_cdna_file):
        """Agent fetch_rows returns correct data."""
        agent = FibPiChunkAgent()
        kernel = agent.open_capsule(sample_cdna_file["path"])

        rows, receipt = agent.fetch_rows(kernel, 10, 40)

        expected = sample_cdna_file["indices"][10:40]
        assert np.array_equal(rows, expected)
        assert receipt.blocks_accessed == 3  # Spans 3 blocks

    def test_agent_stream_matmul(self, sample_cdna_file):
        """Agent stream_matmul computes correctly."""
        agent = FibPiChunkAgent()
        kernel = agent.open_capsule(sample_cdna_file["path"])

        x = np.random.randn(kernel.cols).astype(np.float32)
        y, receipt = agent.stream_matmul(kernel, x)

        # Verify against kernel matmul
        y_kernel, _ = kernel.stream_matmul(x, verify=False)

        assert np.allclose(y, y_kernel, rtol=1e-5)
        assert receipt.blocks_accessed == kernel.num_blocks

    def test_agent_clear_cache(self, sample_cdna_file):
        """Agent clear_cache works correctly."""
        agent = FibPiChunkAgent()
        kernel = agent.open_capsule(sample_cdna_file["path"])

        # Populate cache
        agent.fetch_block(kernel, 0)
        assert agent.cache_stats()["cache"]["entries"] == 1

        # Clear
        agent.clear_cache()
        assert agent.cache_stats()["cache"]["entries"] == 0


# ========== LRU Cache Tests ==========

class TestLRUCache:
    """Tests for LRU cache implementation."""

    def test_lru_basic(self):
        """Basic put/get works."""
        cache = LRUCache(max_size_mb=1.0)

        data = np.zeros((100, 100), dtype=np.uint8)  # 10KB
        cache.put("/path/a", 0, "abc123", data, True)

        entry = cache.get("/path/a", 0)
        assert entry is not None
        assert np.array_equal(entry.data, data)

    def test_lru_miss(self):
        """Cache miss returns None."""
        cache = LRUCache(max_size_mb=1.0)

        entry = cache.get("/path/a", 0)
        assert entry is None

    def test_lru_eviction(self):
        """LRU eviction works correctly."""
        cache = LRUCache(max_size_mb=0.001)  # 1KB limit

        # Add entries that exceed limit
        data1 = np.zeros((100, 10), dtype=np.uint8)  # 1KB
        data2 = np.zeros((100, 10), dtype=np.uint8)  # 1KB

        cache.put("/path/a", 0, "abc", data1, True)
        cache.put("/path/a", 1, "def", data2, True)

        # First entry should be evicted
        assert cache.get("/path/a", 0) is None
        assert cache.get("/path/a", 1) is not None

    def test_lru_access_order(self):
        """Recent access moves entry to end."""
        cache = LRUCache(max_size_mb=0.002)  # 2KB limit

        data1 = np.zeros((100, 10), dtype=np.uint8)  # 1KB
        data2 = np.zeros((100, 10), dtype=np.uint8)  # 1KB

        cache.put("/path/a", 0, "abc", data1, True)
        cache.put("/path/a", 1, "def", data2, True)

        # Access first entry (moves to end)
        cache.get("/path/a", 0)

        # Add third entry (should evict entry 1, not 0)
        data3 = np.zeros((100, 10), dtype=np.uint8)
        cache.put("/path/a", 2, "ghi", data3, True)

        assert cache.get("/path/a", 0) is not None  # Still there
        assert cache.get("/path/a", 1) is None  # Evicted
        assert cache.get("/path/a", 2) is not None  # Still there


# ========== Integration Tests ==========

class TestIntegration:
    """Integration tests with real CDNA files."""

    def test_real_file_kernel(self, real_cdna_file):
        """Kernel works with real CDNA file."""
        k = CDNAKernel(real_cdna_file)

        print(f"\nReal file: {k}")
        print(f"Shape: {k.shape}")
        print(f"Blocks: {k.num_blocks}")

        # Verify all blocks
        ok, fail, failed = k.verify_all_blocks()
        assert fail == 0, f"Verification failures: {failed}"

        # Test matmul
        x = np.random.randn(k.cols).astype(np.float32)
        y, receipt = k.stream_matmul(x, verify=True)
        assert y.shape[0] == k.rows

    def test_real_file_agent(self, real_cdna_file):
        """Agent works with real CDNA file."""
        agent = FibPiChunkAgent(verify_policy="trust_cached")
        kernel = agent.open_capsule(real_cdna_file)

        # Fetch rows
        rows, r1 = agent.fetch_rows(kernel, 0, min(100, kernel.rows))
        assert rows.shape[0] == min(100, kernel.rows)

        # Second fetch should hit cache
        _, r2 = agent.fetch_rows(kernel, 0, min(100, kernel.rows))
        assert r2.cache_hit is True

        print(f"\nCache stats: {agent.cache_stats()}")


# ========== Receipt Tests ==========

class TestReceipts:
    """Tests for receipt generation."""

    def test_kernel_receipt_to_dict(self):
        """KernelReceipt serializes correctly."""
        receipt = KernelReceipt(
            operation="test_op",
            duration_ms=1.23,
            block_id=5,
            verified=True,
        )

        d = receipt.to_dict()
        assert d["operation"] == "test_op"
        assert d["duration_ms"] == 1.23
        assert d["block_id"] == 5
        assert d["verified"] is True
        assert "timestamp" in d

    def test_access_receipt_to_dict(self):
        """AccessReceipt serializes correctly."""
        receipt = AccessReceipt(
            operation="fetch_block",
            duration_ms=0.5,
            cache_hit=True,
            verified=False,
            blocks_accessed=1,
        )

        d = receipt.to_dict()
        assert d["operation"] == "fetch_block"
        assert d["cache_hit"] is True
        assert d["verified"] is False


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
