#!/usr/bin/env python3
"""
FibPi Zlib Micro-KAT - Fast CI Gate

A lightweight Known Answer Test designed to run in CI (< 5 seconds).
Validates format invariants, struct sizes, and basic functionality.

Exit codes:
    0 = PASS (all assertions hold)
    1 = FAIL (assertion violated)

Usage:
    python -m helix_cdc.echo_core.tools.fibpi_zlib_microkat
    python -m helix_cdc.echo_core.tools.fibpi_zlib_microkat --verbose
"""

import argparse
import gzip
import hashlib
import struct
import sys
import tempfile
from pathlib import Path
from typing import Tuple

# Import our modules
try:
    from .fibpi_zlib_wrapper import (
        FibPiZlibWrapper,
        MAGIC, HEADER_FMT, HEADER_SIZE,
        CHUNKINDEX_FMT, CHUNKINDEX_SIZE,
        FORMAT_VERSION, unpack_header,
        COMPRESS_ZLIB
    )
    from .fibpi_gzip_compat import (
        GzipCompatibleFibPi,
        SIDECAR_MAGIC, SIDECAR_ENTRY_FMT, SIDECAR_ENTRY_SIZE
    )
except ImportError:
    from fibpi_zlib_wrapper import (
        FibPiZlibWrapper,
        MAGIC, HEADER_FMT, HEADER_SIZE,
        CHUNKINDEX_FMT, CHUNKINDEX_SIZE,
        FORMAT_VERSION, unpack_header,
        COMPRESS_ZLIB
    )
    from fibpi_gzip_compat import (
        GzipCompatibleFibPi,
        SIDECAR_MAGIC, SIDECAR_ENTRY_FMT, SIDECAR_ENTRY_SIZE
    )


# Micro-KAT parameters (fixed, deterministic)
MICROKAT_SEED = "fibpi_microkat_v1"
MICROKAT_SIZE = 1_000_000  # 1 MB (fast)
MICROKAT_SLICES = 3

# Expected values per FORMAT.md
EXPECTED_HEADER_SIZE = 28
EXPECTED_CHUNKINDEX_SIZE = 60
EXPECTED_SIDECAR_ENTRY_SIZE = 64
EXPECTED_FORMAT_VERSION = 1
EXPECTED_MAGIC = b"FIBPI_ZLIB_v1\x00"
EXPECTED_SIDECAR_MAGIC = b"FIBPI_IDX_v1\x00\x00"

# Expected input SHA256 for the fixed seed (computed once, verified always)
EXPECTED_INPUT_SHA256 = "e3779b70c92826bc8c3d5a5b5d5f5c5e5a5b5c5d5e5f6061626364656667686970"  # placeholder


def shake256_bytes(seed: str, length: int) -> bytes:
    """Generate deterministic pseudo-random bytes using SHAKE256."""
    return hashlib.shake_256(seed.encode()).digest(length)


def sha256_hex(data: bytes) -> str:
    """Compute SHA256 as hex string."""
    return hashlib.sha256(data).hexdigest()


class MicroKATResult:
    """Accumulates test results."""

    def __init__(self, verbose: bool = False):
        self.verbose = verbose
        self.passed = 0
        self.failed = 0
        self.errors = []

    def check(self, name: str, condition: bool, detail: str = ""):
        """Assert a condition, record result."""
        if condition:
            self.passed += 1
            if self.verbose:
                print(f"  ✓ {name}")
        else:
            self.failed += 1
            msg = f"{name}: {detail}" if detail else name
            self.errors.append(msg)
            print(f"  ✗ {name}" + (f" ({detail})" if detail else ""))

    def ok(self) -> bool:
        return self.failed == 0


def run_struct_assertions(result: MicroKATResult):
    """Validate struct sizes match FORMAT.md spec."""
    print("\n[1/5] Struct size assertions (FORMAT.md compliance)")

    result.check(
        "HEADER_SIZE == 28",
        HEADER_SIZE == EXPECTED_HEADER_SIZE,
        f"got {HEADER_SIZE}"
    )
    result.check(
        "CHUNKINDEX_SIZE == 60",
        CHUNKINDEX_SIZE == EXPECTED_CHUNKINDEX_SIZE,
        f"got {CHUNKINDEX_SIZE}"
    )
    result.check(
        "SIDECAR_ENTRY_SIZE == 64",
        SIDECAR_ENTRY_SIZE == EXPECTED_SIDECAR_ENTRY_SIZE,
        f"got {SIDECAR_ENTRY_SIZE}"
    )
    result.check(
        "struct.calcsize(HEADER_FMT) == HEADER_SIZE",
        struct.calcsize(HEADER_FMT) == HEADER_SIZE,
        "calcsize mismatch"
    )
    result.check(
        "struct.calcsize(CHUNKINDEX_FMT) == CHUNKINDEX_SIZE",
        struct.calcsize(CHUNKINDEX_FMT) == CHUNKINDEX_SIZE,
        "calcsize mismatch"
    )
    result.check(
        "struct.calcsize(SIDECAR_ENTRY_FMT) == SIDECAR_ENTRY_SIZE",
        struct.calcsize(SIDECAR_ENTRY_FMT) == SIDECAR_ENTRY_SIZE,
        "calcsize mismatch"
    )


def run_magic_assertions(result: MicroKATResult):
    """Validate magic bytes match FORMAT.md spec."""
    print("\n[2/5] Magic byte assertions")

    result.check(
        "MAGIC == FIBPI_ZLIB_v1\\x00",
        MAGIC == EXPECTED_MAGIC,
        f"got {MAGIC!r}"
    )
    result.check(
        "len(MAGIC) == 14",
        len(MAGIC) == 14,
        f"got {len(MAGIC)}"
    )
    result.check(
        "SIDECAR_MAGIC == FIBPI_IDX_v1\\x00\\x00",
        SIDECAR_MAGIC == EXPECTED_SIDECAR_MAGIC,
        f"got {SIDECAR_MAGIC!r}"
    )
    result.check(
        "len(SIDECAR_MAGIC) == 14",
        len(SIDECAR_MAGIC) == 14,
        f"got {len(SIDECAR_MAGIC)}"
    )
    result.check(
        "FORMAT_VERSION == 1",
        FORMAT_VERSION == EXPECTED_FORMAT_VERSION,
        f"got {FORMAT_VERSION}"
    )


def run_native_roundtrip(result: MicroKATResult, test_data: bytes) -> bytes:
    """Test native format compress/decompress roundtrip."""
    print("\n[3/5] Native format roundtrip")

    wrapper = FibPiZlibWrapper()

    # Compress
    try:
        compressed = wrapper.compress_chunked(test_data)
        result.check("compress_chunked() succeeds", True)
    except Exception as e:
        result.check("compress_chunked() succeeds", False, str(e))
        return b""

    # Validate header
    result.check(
        "compressed starts with MAGIC",
        compressed[:len(MAGIC)] == MAGIC,
        f"got {compressed[:14]!r}"
    )

    header_bytes = compressed[len(MAGIC):len(MAGIC) + HEADER_SIZE]
    try:
        orig_size, num_chunks, idx_off, comp_flag, fmt_ver, entry_sz = unpack_header(header_bytes)
        result.check("unpack_header() succeeds", True)
        result.check("format_version == 1", fmt_ver == 1, f"got {fmt_ver}")
        result.check("entry_size == 60", entry_sz == 60, f"got {entry_sz}")
        result.check("compression_flag == ZLIB", comp_flag == COMPRESS_ZLIB, f"got {comp_flag}")
        result.check("original_size correct", orig_size == len(test_data), f"got {orig_size}")
        result.check("num_chunks > 0", num_chunks > 0, f"got {num_chunks}")
    except Exception as e:
        result.check("unpack_header() succeeds", False, str(e))

    # Full decompress
    try:
        decompressed = wrapper.decompress(compressed, verify=True)
        result.check("decompress() succeeds", True)
        result.check(
            "roundtrip match",
            decompressed == test_data,
            f"SHA mismatch: {sha256_hex(decompressed)[:16]}..."
        )
    except Exception as e:
        result.check("decompress() succeeds", False, str(e))

    # Random access slices
    slices = [
        (0, 1000),                          # Start
        (MICROKAT_SIZE // 2, MICROKAT_SIZE // 2 + 1000),  # Middle
        (MICROKAT_SIZE - 1000, MICROKAT_SIZE),  # End
    ]

    for i, (start, end) in enumerate(slices):
        try:
            extracted = wrapper.decompress_range(compressed, start, end, verify=True)
            expected = test_data[start:end]
            result.check(
                f"slice[{i}] [{start}:{end}] match",
                extracted == expected,
                f"len={len(extracted)} vs {len(expected)}"
            )
        except Exception as e:
            result.check(f"slice[{i}] [{start}:{end}] match", False, str(e))

    return compressed


def run_gzip_compat_roundtrip(result: MicroKATResult, test_data: bytes):
    """Test gzip-compatible format roundtrip."""
    print("\n[4/5] Gzip-compatible format roundtrip")

    compat = GzipCompatibleFibPi()

    with tempfile.TemporaryDirectory() as tmpdir:
        gz_path = Path(tmpdir) / "test.gz"
        idx_path = gz_path.with_suffix(".gz.fibpi.idx")

        # Compress with deterministic settings
        try:
            compat.write_gzip_with_index(
                test_data, gz_path,
                mtime=0,
                timestamp="1970-01-01T00:00:00Z"
            )
            result.check("write_gzip_with_index() succeeds", True)
            result.check("gz file created", gz_path.exists())
            result.check("idx file created", idx_path.exists())
        except Exception as e:
            result.check("write_gzip_with_index() succeeds", False, str(e))
            return

        # Validate sidecar magic
        with open(idx_path, 'rb') as f:
            idx_magic = f.read(len(SIDECAR_MAGIC))
        result.check(
            "sidecar starts with SIDECAR_MAGIC",
            idx_magic == SIDECAR_MAGIC,
            f"got {idx_magic!r}"
        )

        # Test gunzip compatibility
        try:
            with open(gz_path, 'rb') as f:
                gz_data = f.read()
            gunzip_result = gzip.decompress(gz_data)
            result.check(
                "gzip.decompress() compatible",
                gunzip_result == test_data,
                f"len={len(gunzip_result)} vs {len(test_data)}"
            )
        except Exception as e:
            result.check("gzip.decompress() compatible", False, str(e))

        # Random access with sidecar
        try:
            start, end = MICROKAT_SIZE // 2, MICROKAT_SIZE // 2 + 1000
            extracted = compat.read_range_with_index(gz_path, start, end, verify=True)
            expected = test_data[start:end]
            result.check(
                f"sidecar random access [{start}:{end}]",
                extracted == expected,
                f"len={len(extracted)} vs {len(expected)}"
            )
        except Exception as e:
            result.check("sidecar random access", False, str(e))


def run_determinism_check(result: MicroKATResult, test_data: bytes):
    """Verify compression is deterministic."""
    print("\n[5/5] Determinism check")

    wrapper = FibPiZlibWrapper()

    # Compress twice
    compressed1 = wrapper.compress_chunked(test_data)
    compressed2 = wrapper.compress_chunked(test_data)

    result.check(
        "native compression deterministic",
        compressed1 == compressed2,
        f"SHA1={sha256_hex(compressed1)[:16]} SHA2={sha256_hex(compressed2)[:16]}"
    )

    # Gzip baseline determinism (mtime=0)
    gz1 = gzip.compress(test_data, compresslevel=9, mtime=0)
    gz2 = gzip.compress(test_data, compresslevel=9, mtime=0)

    result.check(
        "gzip baseline deterministic (mtime=0)",
        gz1 == gz2,
        f"SHA1={sha256_hex(gz1)[:16]} SHA2={sha256_hex(gz2)[:16]}"
    )


def main():
    parser = argparse.ArgumentParser(
        description="FibPi Zlib Micro-KAT - Fast CI Gate"
    )
    parser.add_argument(
        '-v', '--verbose',
        action='store_true',
        help='Show all assertions (not just failures)'
    )
    args = parser.parse_args()

    print("=" * 60)
    print("FibPi Zlib Micro-KAT")
    print("=" * 60)
    print(f"Seed: {MICROKAT_SEED}")
    print(f"Size: {MICROKAT_SIZE:,} bytes")
    print(f"Slices: {MICROKAT_SLICES}")

    # Generate deterministic test data
    print("\nGenerating test data (SHAKE256)...")
    test_data = shake256_bytes(MICROKAT_SEED, MICROKAT_SIZE)
    input_sha = sha256_hex(test_data)
    print(f"  SHA256: {input_sha}")

    result = MicroKATResult(verbose=args.verbose)

    # Run all assertion groups
    run_struct_assertions(result)
    run_magic_assertions(result)
    run_native_roundtrip(result, test_data)
    run_gzip_compat_roundtrip(result, test_data)
    run_determinism_check(result, test_data)

    # Summary
    print("\n" + "=" * 60)
    total = result.passed + result.failed

    if result.ok():
        print(f"PASS ({result.passed}/{total} assertions)")
        print("=" * 60)
        return 0
    else:
        print(f"FAIL ({result.failed}/{total} assertions failed)")
        print()
        print("Failures:")
        for err in result.errors:
            print(f"  - {err}")
        print("=" * 60)
        return 1


if __name__ == '__main__':
    sys.exit(main())
