#!/usr/bin/env python3
"""
Test suite for CDNA v2 format.

Tests:
    1. test_write_read_roundtrip - Write → read → compare indices/codebook
    2. test_random_access - Seek to block N, verify contents
    3. test_block_sha256_verification - Corrupt a block, verify detection
    4. test_v1_to_v2_conversion - Convert → verify equivalence
    5. test_streaming_decode - Iterate blocks, verify order
    6. test_corrupted_file_handling - Truncated, bad magic, etc.
    7. test_matmul_equivalence - v2 matmul matches v1 output

Run with:
    pytest tests/test_cdna_v2.py -v
"""

import hashlib
import json
import struct
import tempfile
from pathlib import Path

import numpy as np
import pytest

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

from helix_cdc.regrow.cdna_stream_v2 import (
    CDNA2_MAGIC,
    CDNA2_VERSION,
    CODEBOOK_SIZE,
    BLOCK_INDEX_ENTRY_SIZE,
    BlockIndexEntry,
    CDNAv2Writer,
    CDNAv2Reader,
    convert_v1_to_v2,
    cdna2_matmul_row,
    cdna2_matmul_block,
    stream_cdna2_matmul,
    load_cdna_auto,
)

# Check if brotli is available
try:
    import brotli
    HAS_BROTLI = True
except ImportError:
    HAS_BROTLI = False


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

@pytest.fixture
def sample_indices():
    """Create sample uint8 indices [128, 256]."""
    np.random.seed(42)
    return np.random.randint(0, 256, size=(128, 256), dtype=np.uint8)


@pytest.fixture
def sample_codebook():
    """Create sample float32 codebook [256]."""
    np.random.seed(42)
    return np.random.randn(CODEBOOK_SIZE).astype(np.float32) * 0.1


@pytest.fixture
def temp_dir():
    """Create temporary directory for test files."""
    with tempfile.TemporaryDirectory() as tmpdir:
        yield Path(tmpdir)


@pytest.fixture
def sample_v2_file(sample_indices, sample_codebook, temp_dir):
    """Create a sample CDNAv2 file and return its path."""
    if not HAS_BROTLI:
        pytest.skip("brotli not installed")

    output_path = temp_dir / "test_tensor.cdna2.hxz"
    writer = CDNAv2Writer()
    writer.write(
        indices=sample_indices,
        codebook=sample_codebook,
        output_path=output_path,
        tensor_name="test.weight",
        block_rows=32,
    )
    return output_path, sample_indices, sample_codebook


# ============================================================================
# Test 1: Write/Read Roundtrip
# ============================================================================

@pytest.mark.skipif(not HAS_BROTLI, reason="brotli not installed")
def test_write_read_roundtrip(sample_indices, sample_codebook, temp_dir):
    """Write CDNAv2 file, read it back, verify indices and codebook match."""
    output_path = temp_dir / "roundtrip.cdna2.hxz"

    # Write
    writer = CDNAv2Writer(compression_quality=11)
    receipt = writer.write(
        indices=sample_indices,
        codebook=sample_codebook,
        output_path=output_path,
        tensor_name="blk.0.attn_k.weight",
        block_rows=32,
    )

    # Verify receipt
    assert receipt["schema"] == "create_cdna2_hxz_v1"
    assert receipt["shape"] == [128, 256]
    assert receipt["num_blocks"] == 4  # 128 rows / 32 block_rows

    # Read back
    reader = CDNAv2Reader(output_path)

    # Verify header
    assert reader.shape == (128, 256)
    assert reader.rows == 128
    assert reader.cols == 256
    assert reader.tensor_name == "blk.0.attn_k.weight"
    assert reader.block_rows == 32
    assert reader.num_blocks == 4

    # Verify codebook
    assert np.allclose(reader.codebook, sample_codebook)

    # Verify indices
    read_indices = reader.get_all_indices()
    assert read_indices.shape == sample_indices.shape
    assert read_indices.dtype == np.uint8
    assert np.array_equal(read_indices, sample_indices)


# ============================================================================
# Test 2: Random Access
# ============================================================================

@pytest.mark.skipif(not HAS_BROTLI, reason="brotli not installed")
def test_random_access(sample_v2_file):
    """Seek to specific blocks and verify contents."""
    path, original_indices, _ = sample_v2_file

    reader = CDNAv2Reader(path)

    # Test random access to various row ranges
    test_ranges = [
        (0, 32),    # First block
        (32, 64),   # Second block
        (50, 100),  # Spans two blocks
        (100, 128), # Last block partial
        (10, 20),   # Within first block
    ]

    for start, end in test_ranges:
        rows = reader.get_rows(start, end)
        expected = original_indices[start:end]

        assert rows.shape == expected.shape, f"Shape mismatch for range [{start}, {end})"
        assert np.array_equal(rows, expected), f"Content mismatch for range [{start}, {end})"


@pytest.mark.skipif(not HAS_BROTLI, reason="brotli not installed")
def test_random_access_single_row(sample_v2_file):
    """Test accessing a single row."""
    path, original_indices, _ = sample_v2_file

    reader = CDNAv2Reader(path)

    # Access single rows from different blocks
    for row in [0, 31, 32, 64, 100, 127]:
        single_row = reader.get_rows(row, row + 1)
        expected = original_indices[row:row + 1]

        assert single_row.shape == (1, 256)
        assert np.array_equal(single_row, expected), f"Mismatch at row {row}"


# ============================================================================
# Test 3: Block SHA256 Verification
# ============================================================================

@pytest.mark.skipif(not HAS_BROTLI, reason="brotli not installed")
def test_block_sha256_verification(sample_v2_file, temp_dir):
    """Verify SHA256 block verification detects corruption."""
    path, _, _ = sample_v2_file

    reader = CDNAv2Reader(path)

    # All blocks should verify initially
    ok_count, fail_count, failed = reader.verify_all_blocks()
    assert ok_count == 4
    assert fail_count == 0
    assert len(failed) == 0

    # Corrupt the file by modifying a block
    data = path.read_bytes()
    data_array = bytearray(data)

    # Find the first block's offset and corrupt some bytes
    first_block_offset = reader.block_index[0].offset
    data_array[first_block_offset + 10] ^= 0xFF  # Flip bits

    # Write corrupted file
    corrupted_path = temp_dir / "corrupted.cdna2.hxz"
    corrupted_path.write_bytes(bytes(data_array))

    # Read corrupted file
    corrupted_reader = CDNAv2Reader(corrupted_path)

    # Block 0 should fail verification
    assert not corrupted_reader.verify_block(0)

    # Other blocks should still verify (if they weren't corrupted)
    for i in range(1, corrupted_reader.num_blocks):
        # May or may not pass depending on corruption extent
        pass  # Just check no exception


@pytest.mark.skipif(not HAS_BROTLI, reason="brotli not installed")
def test_block_verification_all_pass(sample_v2_file):
    """Verify all blocks pass SHA256 verification on valid file."""
    path, _, _ = sample_v2_file
    reader = CDNAv2Reader(path)

    for i in range(reader.num_blocks):
        assert reader.verify_block(i), f"Block {i} verification failed unexpectedly"


# ============================================================================
# Test 4: V1 to V2 Conversion
# ============================================================================

@pytest.mark.skipif(not HAS_BROTLI, reason="brotli not installed")
def test_v1_to_v2_conversion(sample_indices, sample_codebook, temp_dir):
    """Convert v1 file to v2 and verify equivalence."""
    from helix_cdc.regrow.cdna_stream import create_cdna_hxz, CDNAStreamReader

    # Create v1 file
    v1_path = temp_dir / "test.cdna.hxz"
    create_cdna_hxz(
        indices=sample_indices,
        codebook=sample_codebook,
        tensor_name="blk.0.ffn_gate.weight",
        output_path=str(v1_path),
    )

    # Convert to v2
    v2_path = temp_dir / "test.cdna2.hxz"
    receipt = convert_v1_to_v2(v1_path, v2_path, block_rows=32, verify=True)

    # Check conversion receipt
    assert receipt["verification"]["match"] == True
    assert receipt["verification"]["indices_match"] == True
    assert receipt["verification"]["codebook_match"] == True

    # Read both and compare
    v1_reader = CDNAStreamReader(str(v1_path))
    v2_reader = CDNAv2Reader(v2_path)

    # Compare shape
    assert v1_reader.shape == v2_reader.shape

    # Compare codebook
    assert np.allclose(v1_reader.codebook, v2_reader.codebook)

    # Compare indices
    v1_blob = v1_reader._ensure_indices_blob()
    v1_indices = np.frombuffer(v1_blob, dtype=np.uint8).reshape(v1_reader.rows, v1_reader.cols)
    v2_indices = v2_reader.get_all_indices()

    assert np.array_equal(v1_indices, v2_indices)


@pytest.mark.skipif(not HAS_BROTLI, reason="brotli not installed")
def test_load_cdna_auto_v1(sample_indices, sample_codebook, temp_dir):
    """Test auto-detection of v1 format."""
    from helix_cdc.regrow.cdna_stream import create_cdna_hxz, CDNAStreamReader

    # Create v1 file
    v1_path = temp_dir / "auto_v1.cdna.hxz"
    create_cdna_hxz(
        indices=sample_indices,
        codebook=sample_codebook,
        tensor_name="test.weight",
        output_path=str(v1_path),
    )

    # Auto-load should return v1 reader
    reader = load_cdna_auto(v1_path)
    assert isinstance(reader, CDNAStreamReader)


@pytest.mark.skipif(not HAS_BROTLI, reason="brotli not installed")
def test_load_cdna_auto_v2(sample_v2_file):
    """Test auto-detection of v2 format."""
    path, _, _ = sample_v2_file

    reader = load_cdna_auto(path)
    assert isinstance(reader, CDNAv2Reader)


# ============================================================================
# Test 5: Streaming Decode
# ============================================================================

@pytest.mark.skipif(not HAS_BROTLI, reason="brotli not installed")
def test_streaming_decode(sample_v2_file):
    """Stream all blocks and verify they arrive in order."""
    path, original_indices, _ = sample_v2_file

    reader = CDNAv2Reader(path)

    all_indices = []
    prev_start_row = -1
    block_count = 0

    for indices_block, start_row in reader.stream_blocks():
        # Verify order
        assert start_row > prev_start_row, f"Out of order: {start_row} <= {prev_start_row}"
        prev_start_row = start_row

        # Verify block content
        block_end = start_row + indices_block.shape[0]
        expected_block = original_indices[start_row:block_end]
        assert np.array_equal(indices_block, expected_block)

        all_indices.append(indices_block)
        block_count += 1

    # Verify all blocks were streamed
    assert block_count == reader.num_blocks

    # Verify concatenated result matches original
    full_result = np.vstack(all_indices)
    assert np.array_equal(full_result, original_indices)


@pytest.mark.skipif(not HAS_BROTLI, reason="brotli not installed")
def test_streaming_partial_range(sample_v2_file):
    """Stream only a subset of blocks."""
    path, original_indices, _ = sample_v2_file

    reader = CDNAv2Reader(path)

    # Stream blocks 1 and 2 only
    blocks = list(reader.stream_blocks(start_block=1, end_block=3))

    assert len(blocks) == 2

    # Verify block 1
    indices1, start1 = blocks[0]
    assert start1 == 32
    assert np.array_equal(indices1, original_indices[32:64])

    # Verify block 2
    indices2, start2 = blocks[1]
    assert start2 == 64
    assert np.array_equal(indices2, original_indices[64:96])


# ============================================================================
# Test 6: Corrupted File Handling
# ============================================================================

def test_bad_magic(temp_dir):
    """Test handling of file with bad magic bytes."""
    bad_file = temp_dir / "bad_magic.cdna2.hxz"
    bad_file.write_bytes(b"BAAD" + b"\x00" * 100)

    with pytest.raises(ValueError, match="Not a CDNAv2 file"):
        CDNAv2Reader(bad_file)


def test_truncated_file(temp_dir):
    """Test handling of truncated file."""
    truncated_file = temp_dir / "truncated.cdna2.hxz"
    # Just magic and partial header
    truncated_file.write_bytes(CDNA2_MAGIC + b"\x00\x01")

    with pytest.raises(Exception):  # Could be struct.error or EOFError
        CDNAv2Reader(truncated_file)


def test_unsupported_version(temp_dir):
    """Test handling of unsupported version."""
    bad_version_file = temp_dir / "bad_version.cdna2.hxz"

    # Magic + version 99 + minimal structure
    data = CDNA2_MAGIC + struct.pack("<H", 99) + b"\x00" * 100
    bad_version_file.write_bytes(data)

    with pytest.raises(ValueError, match="Unsupported CDNAv2 version"):
        CDNAv2Reader(bad_version_file)


@pytest.mark.skipif(not HAS_BROTLI, reason="brotli not installed")
def test_invalid_row_range(sample_v2_file):
    """Test handling of invalid row ranges."""
    path, _, _ = sample_v2_file
    reader = CDNAv2Reader(path)

    # Negative start
    with pytest.raises(ValueError):
        reader.get_rows(-1, 10)

    # End beyond tensor
    with pytest.raises(ValueError):
        reader.get_rows(100, 200)

    # Start >= end
    with pytest.raises(ValueError):
        reader.get_rows(50, 50)

    with pytest.raises(ValueError):
        reader.get_rows(60, 50)


# ============================================================================
# Test 7: Matmul Equivalence
# ============================================================================

@pytest.mark.skipif(not HAS_BROTLI, reason="brotli not installed")
def test_matmul_equivalence(sample_v2_file):
    """Verify v2 matmul produces same output as v1."""
    from helix_cdc.regrow.cdna_stream import (
        cdna_matmul_row as v1_matmul_row,
        cdna_matmul_block as v1_matmul_block,
    )

    path, original_indices, codebook = sample_v2_file

    # Random input vector
    np.random.seed(123)
    x = np.random.randn(256).astype(np.float32)

    # V2 matmul via streaming
    y_v2, stats = stream_cdna2_matmul(path, x, block_rows=32)

    # V1-style matmul (direct computation on original data)
    y_v1 = np.zeros(128, dtype=np.float32)
    for i in range(128):
        y_v1[i] = v1_matmul_row(original_indices[i], codebook, x)

    # Compare
    assert y_v2.shape == y_v1.shape
    assert np.allclose(y_v2, y_v1, rtol=1e-5, atol=1e-6)


@pytest.mark.skipif(not HAS_BROTLI, reason="brotli not installed")
def test_matmul_block_equivalence(sample_indices, sample_codebook):
    """Test block matmul produces correct results."""
    np.random.seed(456)
    x = np.random.randn(256).astype(np.float32)

    # Compute block matmul
    block = sample_indices[:32]
    y_block = cdna2_matmul_block(block, sample_codebook, x)

    # Verify shape
    assert y_block.shape == (32,)

    # Verify against row-by-row
    for i in range(32):
        y_row = cdna2_matmul_row(block[i], sample_codebook, x)
        assert np.isclose(y_block[i], y_row, rtol=1e-5)


# ============================================================================
# Additional Tests
# ============================================================================

@pytest.mark.skipif(not HAS_BROTLI, reason="brotli not installed")
def test_1d_indices_handling(sample_codebook, temp_dir):
    """Test handling of 1D indices (treated as single row)."""
    np.random.seed(789)
    indices_1d = np.random.randint(0, 256, size=256, dtype=np.uint8)

    output_path = temp_dir / "1d.cdna2.hxz"
    writer = CDNAv2Writer()
    receipt = writer.write(
        indices=indices_1d,
        codebook=sample_codebook,
        output_path=output_path,
        tensor_name="1d_tensor",
    )

    assert receipt["shape"] == [1, 256]

    reader = CDNAv2Reader(output_path)
    assert reader.shape == (1, 256)

    read_indices = reader.get_all_indices()
    assert read_indices.shape == (1, 256)
    assert np.array_equal(read_indices.flatten(), indices_1d)


@pytest.mark.skipif(not HAS_BROTLI, reason="brotli not installed")
def test_memory_stats(sample_v2_file):
    """Test memory statistics reporting."""
    path, _, _ = sample_v2_file

    reader = CDNAv2Reader(path)
    stats = reader.memory_stats

    assert stats["shape"] == [128, 256]
    assert stats["rows"] == 128
    assert stats["cols"] == 256
    assert stats["block_rows"] == 32
    assert stats["num_blocks"] == 4
    assert stats["full_tensor_bytes"] == 128 * 256 * 4  # float32
    assert stats["codebook_bytes"] == 256 * 4
    assert stats["codebook_entries"] == 256


@pytest.mark.skipif(not HAS_BROTLI, reason="brotli not installed")
def test_block_index_entry_serialization():
    """Test BlockIndexEntry round-trip serialization."""
    original = BlockIndexEntry(
        offset=12345678,
        comp_len=1024,
        decomp_len=32768,
        sha256=hashlib.sha256(b"test").digest(),
    )

    serialized = original.to_bytes()
    assert len(serialized) == BLOCK_INDEX_ENTRY_SIZE

    restored = BlockIndexEntry.from_bytes(serialized)

    assert restored.offset == original.offset
    assert restored.comp_len == original.comp_len
    assert restored.decomp_len == original.decomp_len
    assert restored.sha256 == original.sha256


@pytest.mark.skipif(not HAS_BROTLI, reason="brotli not installed")
def test_extra_metadata(sample_indices, sample_codebook, temp_dir):
    """Test writing and reading extra metadata."""
    output_path = temp_dir / "with_meta.cdna2.hxz"

    extra_meta = {
        "source_gguf": "mistral-7b.gguf",
        "layer_idx": 5,
        "custom_field": "test_value",
    }

    writer = CDNAv2Writer()
    receipt = writer.write(
        indices=sample_indices,
        codebook=sample_codebook,
        output_path=output_path,
        tensor_name="test.weight",
        extra_meta=extra_meta,
    )

    reader = CDNAv2Reader(output_path)

    assert reader.header.get("source_gguf") == "mistral-7b.gguf"
    assert reader.header.get("layer_idx") == 5
    assert reader.header.get("custom_field") == "test_value"


@pytest.mark.skipif(not HAS_BROTLI, reason="brotli not installed")
def test_different_block_sizes(sample_indices, sample_codebook, temp_dir):
    """Test various block sizes."""
    for block_rows in [8, 16, 32, 64, 128]:
        output_path = temp_dir / f"block_{block_rows}.cdna2.hxz"

        writer = CDNAv2Writer()
        writer.write(
            indices=sample_indices,
            codebook=sample_codebook,
            output_path=output_path,
            tensor_name="test",
            block_rows=block_rows,
        )

        reader = CDNAv2Reader(output_path)

        expected_blocks = (128 + block_rows - 1) // block_rows
        assert reader.num_blocks == expected_blocks
        assert reader.block_rows == block_rows

        # Verify data integrity
        read_indices = reader.get_all_indices()
        assert np.array_equal(read_indices, sample_indices)


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