#!/usr/bin/env python3
"""
FibPi Zlib KAT (Known Answer Test)

Produces deterministic, reproducible test results with receipts.
This is the bridge from "cool utility" to "provable substrate module."

Features:
- Deterministic pseudo-random input (SHAKE256 seeded)
- Tests both native FibPi and gzip-compatible modes
- Compares against standard gzip baseline
- Verifies K random slices against original
- Measures timing and memory metrics
- Writes receipt JSON + separate artifact files
- Reproducible across runs (same seed = same test)

Output structure:
    receipts/fibpi_zlib_kat/20260122_194500/
    ├── receipt.json          # Main receipt with all results
    ├── bench.json            # Benchmark comparison data
    ├── input.sha256          # SHA256 of test input
    ├── fibpi_native.sha256   # SHA256 of native compressed
    ├── fibpi_compat.sha256   # SHA256 of gzip-compat .gz
    ├── compat_sidecar.sha256 # SHA256 of sidecar index
    ├── gzip_baseline.sha256  # SHA256 of standard gzip
    └── slice_proofs.json     # Detailed slice verification

Usage:
    python -m helix_cdc.echo_core.tools.fibpi_zlib_kat
    python -m helix_cdc.echo_core.tools.fibpi_zlib_kat --seed "custom_seed" --size 10000000
"""

import argparse
import gzip
import hashlib
import json
import platform
import resource
import subprocess
import sys
import time
import zlib
from datetime import datetime
from pathlib import Path
from typing import List, Tuple, Dict, Any, Optional

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


# Default KAT parameters
DEFAULT_SEED = "fibpi_zlib_kat:v1"  # Stable seed per spec
DEFAULT_SIZE = 10_000_000  # 10MB per spec
DEFAULT_NUM_SLICES = 10
DEFAULT_RECEIPTS_DIR = Path(__file__).parent.parent.parent.parent / "receipts" / "fibpi_zlib_kat"


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()


def get_peak_rss_mb() -> float:
    """Get peak RSS memory usage in MB (Linux: ru_maxrss is KB)."""
    usage = resource.getrusage(resource.RUSAGE_SELF)
    # On Linux, ru_maxrss is in KB; on macOS it's bytes
    if platform.system() == 'Darwin':
        return usage.ru_maxrss / (1024 * 1024)
    return usage.ru_maxrss / 1024


def measure_operation(func, *args, **kwargs) -> Tuple[Any, float, float]:
    """
    Measure time and peak RSS for an operation.

    Returns (result, elapsed_seconds, peak_rss_mb)
    """
    # Get baseline RSS
    rss_before = get_peak_rss_mb()

    start = time.perf_counter()
    result = func(*args, **kwargs)
    elapsed = time.perf_counter() - start

    rss_after = get_peak_rss_mb()

    return result, elapsed, rss_after


def generate_deterministic_slices(
    data_len: int,
    num_slices: int,
    seed: str
) -> List[Tuple[int, int]]:
    """
    Generate deterministic slice ranges for testing.

    Uses SHAKE256 to generate reproducible (start, end) pairs.
    """
    random_bytes = shake256_bytes(f"{seed}_slices", num_slices * 8)

    slices = []
    for i in range(num_slices):
        offset = i * 8
        start_raw = int.from_bytes(random_bytes[offset:offset+4], 'big')
        len_raw = int.from_bytes(random_bytes[offset+4:offset+8], 'big')

        start = start_raw % data_len
        max_len = min(1024 * 1024, data_len - start)  # Max 1MB slices
        length = (len_raw % max_len) + 1
        end = start + length

        slices.append((start, end))

    return slices


def run_gzip_baseline(test_data: bytes) -> Dict[str, Any]:
    """Run standard gzip compression as baseline (deterministic mtime=0)."""
    # Compress with mtime=0 for determinism (Python 3.8+)
    # Without mtime=0, gzip embeds current timestamp → nondeterministic SHA256
    compressed, compress_time, compress_rss = measure_operation(
        gzip.compress, test_data, compresslevel=9, mtime=0
    )

    # Decompress
    decompressed, decompress_time, decompress_rss = measure_operation(
        gzip.decompress, compressed
    )

    return {
        "format": "gzip_baseline",
        "original_size": len(test_data),
        "compressed_size": len(compressed),
        "compression_ratio": len(compressed) / len(test_data),
        "original_sha256": sha256_hex(test_data),
        "compressed_sha256": sha256_hex(compressed),
        "decompressed_sha256": sha256_hex(decompressed),
        "full_roundtrip_match": decompressed == test_data,
        "timing": {
            "compress_seconds": compress_time,
            "decompress_seconds": decompress_time,
        },
        "memory": {
            "compress_peak_rss_mb": compress_rss,
            "decompress_peak_rss_mb": decompress_rss,
        }
    }


def run_native_kat(
    test_data: bytes,
    slices: List[Tuple[int, int]]
) -> Dict[str, Any]:
    """Run KAT on native FibPi format with timing/memory metrics."""
    wrapper = FibPiZlibWrapper()

    # Compress with timing
    compressed, compress_time, compress_rss = measure_operation(
        wrapper.compress_chunked, test_data
    )

    # Full decompress with timing
    decompressed, decompress_time, decompress_rss = measure_operation(
        wrapper.decompress, compressed, verify=True
    )
    full_match = decompressed == test_data

    # Slice tests with timing
    slice_results = []
    total_slice_time = 0.0

    for i, (start, end) in enumerate(slices):
        extracted, slice_time, _ = measure_operation(
            wrapper.decompress_range, compressed, start, end, verify=True
        )
        total_slice_time += slice_time

        expected = test_data[start:end]
        match = extracted == expected

        slice_results.append({
            "slice_id": i,
            "start": start,
            "end": end,
            "length": end - start,
            "expected_sha256": sha256_hex(expected),
            "extracted_sha256": sha256_hex(extracted),
            "match": match,
            "extract_seconds": slice_time
        })

    # Random access benchmark: 1MB at deterministic offset
    mid_offset = len(test_data) // 2
    range_end = min(mid_offset + 1024 * 1024, len(test_data))
    _, random_access_time, random_access_rss = measure_operation(
        wrapper.decompress_range, compressed, mid_offset, range_end, verify=False
    )

    stats = wrapper.get_stats()

    return {
        "format": "fibpi_native",
        "original_size": len(test_data),
        "compressed_size": len(compressed),
        "compression_ratio": len(compressed) / len(test_data),
        "original_sha256": sha256_hex(test_data),
        "compressed_sha256": sha256_hex(compressed),
        "decompressed_sha256": sha256_hex(decompressed),
        "full_roundtrip_match": full_match,
        "num_chunks": stats.get("chunks_compressed", 0),
        "timing": {
            "compress_seconds": compress_time,
            "decompress_seconds": decompress_time,
            "total_slice_seconds": total_slice_time,
            "random_access_1mb_seconds": random_access_time,
        },
        "memory": {
            "compress_peak_rss_mb": compress_rss,
            "decompress_peak_rss_mb": decompress_rss,
            "random_access_peak_rss_mb": random_access_rss,
        },
        "slice_tests": slice_results,
        "all_slices_pass": all(s["match"] for s in slice_results)
    }


def run_compat_kat(
    test_data: bytes,
    slices: List[Tuple[int, int]],
    temp_dir: Path
) -> Dict[str, Any]:
    """Run KAT on gzip-compatible format with timing/memory metrics."""
    compat = GzipCompatibleFibPi()

    gz_path = temp_dir / "kat_test.gz"
    idx_path = temp_dir / "kat_test.gz.fibpi.idx"

    # Compress with timing (mtime=0, fixed timestamp for determinism)
    def do_compress():
        compat.write_gzip_with_index(
            test_data, gz_path,
            mtime=0,
            timestamp="1970-01-01T00:00:00Z"  # Fixed epoch for reproducibility
        )
    _, compress_time, compress_rss = measure_operation(do_compress)

    gz_size = gz_path.stat().st_size
    idx_size = idx_path.stat().st_size

    with open(gz_path, 'rb') as f:
        gz_data = f.read()
    with open(idx_path, 'rb') as f:
        idx_data = f.read()

    # Full decompress with timing
    decompressed, decompress_time, decompress_rss = measure_operation(
        compat.full_decompress_with_index, gz_path, verify=True
    )
    full_match = decompressed == test_data

    # Test standard gunzip compatibility
    result = subprocess.run(
        ['gunzip', '-c', str(gz_path)],
        capture_output=True
    )
    gunzip_data = result.stdout
    gunzip_match = gunzip_data == test_data

    # Slice tests with timing
    slice_results = []
    total_slice_time = 0.0

    for i, (start, end) in enumerate(slices):
        extracted, slice_time, _ = measure_operation(
            compat.read_range_with_index, gz_path, start, end, verify=True
        )
        total_slice_time += slice_time

        expected = test_data[start:end]
        match = extracted == expected

        slice_results.append({
            "slice_id": i,
            "start": start,
            "end": end,
            "length": end - start,
            "expected_sha256": sha256_hex(expected),
            "extracted_sha256": sha256_hex(extracted),
            "match": match,
            "extract_seconds": slice_time
        })

    # Random access benchmark: 1MB at deterministic offset
    mid_offset = len(test_data) // 2
    range_end = min(mid_offset + 1024 * 1024, len(test_data))
    _, random_access_time, random_access_rss = measure_operation(
        compat.read_range_with_index, gz_path, mid_offset, range_end, verify=False
    )

    with open(idx_path, 'rb') as f:
        sidecar = SidecarIndex.from_bytes(f.read())

    # Cleanup
    gz_path.unlink()
    idx_path.unlink()

    return {
        "format": "fibpi_gzip_compat",
        "original_size": len(test_data),
        "gz_size": gz_size,
        "idx_size": idx_size,
        "total_size": gz_size + idx_size,
        "compression_ratio": gz_size / len(test_data),
        "original_sha256": sha256_hex(test_data),
        "gz_sha256": sha256_hex(gz_data),
        "idx_sha256": sha256_hex(idx_data),
        "decompressed_sha256": sha256_hex(decompressed),
        "full_roundtrip_match": full_match,
        "gunzip_compatible": gunzip_match,
        "num_chunks": sidecar.num_chunks,
        "timing": {
            "compress_seconds": compress_time,
            "decompress_seconds": decompress_time,
            "total_slice_seconds": total_slice_time,
            "random_access_1mb_seconds": random_access_time,
        },
        "memory": {
            "compress_peak_rss_mb": compress_rss,
            "decompress_peak_rss_mb": decompress_rss,
            "random_access_peak_rss_mb": random_access_rss,
        },
        "slice_tests": slice_results,
        "all_slices_pass": all(s["match"] for s in slice_results)
    }


def write_sha256_file(path: Path, sha256: str, label: str = ""):
    """Write a .sha256 file with the hash."""
    with open(path, 'w') as f:
        if label:
            f.write(f"{sha256}  {label}\n")
        else:
            f.write(f"{sha256}\n")


def run_kat(
    seed: str = DEFAULT_SEED,
    size: int = DEFAULT_SIZE,
    num_slices: int = DEFAULT_NUM_SLICES,
    receipts_dir: Path = DEFAULT_RECEIPTS_DIR,
    verbose: bool = True
) -> Dict[str, Any]:
    """
    Run full KAT suite and produce receipt with separate artifact files.

    Args:
        seed: Deterministic seed for test data generation
        size: Size of test data in bytes
        num_slices: Number of random slices to test
        receipts_dir: Directory for receipt output
        verbose: Print progress

    Returns:
        Complete receipt dict
    """
    timestamp = datetime.utcnow().isoformat() + "Z"
    safe_timestamp = timestamp.replace(":", "").replace("-", "").replace(".", "_")[:15]

    # Create timestamped output directory
    output_dir = Path(receipts_dir) / safe_timestamp
    output_dir.mkdir(parents=True, exist_ok=True)

    if verbose:
        print(f"FibPi Zlib KAT v2.0")
        print(f"=" * 60)
        print(f"Seed: {seed}")
        print(f"Size: {size:,} bytes")
        print(f"Slices: {num_slices}")
        print(f"Output: {output_dir}")
        print()

    # Generate deterministic test data
    if verbose:
        print("Generating deterministic test data (SHAKE256)...")
    test_data = shake256_bytes(seed, size)
    test_sha256 = sha256_hex(test_data)

    if verbose:
        print(f"  SHA256: {test_sha256}")
        print()

    # Write input.sha256
    write_sha256_file(output_dir / "input.sha256", test_sha256, "input")

    # Generate deterministic slice ranges
    slices = generate_deterministic_slices(size, num_slices, seed)

    if verbose:
        print(f"Test slices: {len(slices)}")
        for i, (s, e) in enumerate(slices[:3]):
            print(f"  [{i}] {s:,} - {e:,} ({e-s:,} bytes)")
        if len(slices) > 3:
            print(f"  ... and {len(slices) - 3} more")
        print()

    # Run gzip baseline
    if verbose:
        print("Running gzip baseline...")
    baseline_result = run_gzip_baseline(test_data)
    write_sha256_file(output_dir / "gzip_baseline.sha256",
                      baseline_result["compressed_sha256"], "gzip_baseline")

    if verbose:
        print(f"  Compressed: {baseline_result['compressed_size']:,} bytes ({baseline_result['compression_ratio']:.2%})")
        print(f"  Compress: {baseline_result['timing']['compress_seconds']:.3f}s")
        print(f"  Decompress: {baseline_result['timing']['decompress_seconds']:.3f}s")
        print()

    # Run native format KAT
    if verbose:
        print("Testing native FibPi format...")
    native_result = run_native_kat(test_data, slices)
    write_sha256_file(output_dir / "fibpi_native.sha256",
                      native_result["compressed_sha256"], "fibpi_native")

    if verbose:
        print(f"  Compressed: {native_result['compressed_size']:,} bytes ({native_result['compression_ratio']:.2%})")
        print(f"  Chunks: {native_result['num_chunks']}")
        print(f"  Compress: {native_result['timing']['compress_seconds']:.3f}s")
        print(f"  Decompress: {native_result['timing']['decompress_seconds']:.3f}s")
        print(f"  Random access 1MB: {native_result['timing']['random_access_1mb_seconds']:.3f}s")
        print(f"  Roundtrip: {'PASS' if native_result['full_roundtrip_match'] else 'FAIL'}")
        print(f"  Slices: {'PASS' if native_result['all_slices_pass'] else 'FAIL'}")
        print()

    # Run gzip-compatible KAT
    if verbose:
        print("Testing gzip-compatible format...")

    import tempfile
    with tempfile.TemporaryDirectory() as tmpdir:
        compat_result = run_compat_kat(test_data, slices, Path(tmpdir))

    write_sha256_file(output_dir / "fibpi_compat.sha256",
                      compat_result["gz_sha256"], "fibpi_compat.gz")
    write_sha256_file(output_dir / "compat_sidecar.sha256",
                      compat_result["idx_sha256"], "fibpi_compat.gz.fibpi.idx")

    if verbose:
        print(f"  .gz: {compat_result['gz_size']:,} bytes ({compat_result['compression_ratio']:.2%})")
        print(f"  .idx: {compat_result['idx_size']:,} bytes")
        print(f"  Compress: {compat_result['timing']['compress_seconds']:.3f}s")
        print(f"  Decompress: {compat_result['timing']['decompress_seconds']:.3f}s")
        print(f"  Random access 1MB: {compat_result['timing']['random_access_1mb_seconds']:.3f}s")
        print(f"  Roundtrip: {'PASS' if compat_result['full_roundtrip_match'] else 'FAIL'}")
        print(f"  Gunzip: {'COMPATIBLE' if compat_result['gunzip_compatible'] else 'INCOMPATIBLE'}")
        print(f"  Slices: {'PASS' if compat_result['all_slices_pass'] else 'FAIL'}")
        print()

    # Build benchmark comparison
    bench = {
        "test_size_bytes": size,
        "formats": {
            "gzip_baseline": {
                "compressed_size": baseline_result["compressed_size"],
                "compression_ratio": baseline_result["compression_ratio"],
                "compress_seconds": baseline_result["timing"]["compress_seconds"],
                "decompress_seconds": baseline_result["timing"]["decompress_seconds"],
                "random_access_supported": False,
            },
            "fibpi_native": {
                "compressed_size": native_result["compressed_size"],
                "compression_ratio": native_result["compression_ratio"],
                "compress_seconds": native_result["timing"]["compress_seconds"],
                "decompress_seconds": native_result["timing"]["decompress_seconds"],
                "random_access_1mb_seconds": native_result["timing"]["random_access_1mb_seconds"],
                "random_access_supported": True,
                "num_chunks": native_result["num_chunks"],
            },
            "fibpi_gzip_compat": {
                "gz_size": compat_result["gz_size"],
                "idx_size": compat_result["idx_size"],
                "total_size": compat_result["total_size"],
                "compression_ratio": compat_result["compression_ratio"],
                "compress_seconds": compat_result["timing"]["compress_seconds"],
                "decompress_seconds": compat_result["timing"]["decompress_seconds"],
                "random_access_1mb_seconds": compat_result["timing"]["random_access_1mb_seconds"],
                "random_access_supported": True,
                "gunzip_compatible": True,
                "num_chunks": compat_result["num_chunks"],
            },
        },
        "comparison": {
            "native_vs_gzip_compress_overhead": (
                native_result["timing"]["compress_seconds"] /
                baseline_result["timing"]["compress_seconds"]
            ) if baseline_result["timing"]["compress_seconds"] > 0 else None,
            "native_random_access_speedup": (
                baseline_result["timing"]["decompress_seconds"] /
                native_result["timing"]["random_access_1mb_seconds"]
            ) if native_result["timing"]["random_access_1mb_seconds"] > 0 else None,
            "compat_random_access_speedup": (
                baseline_result["timing"]["decompress_seconds"] /
                compat_result["timing"]["random_access_1mb_seconds"]
            ) if compat_result["timing"]["random_access_1mb_seconds"] > 0 else None,
        }
    }

    # Write bench.json
    with open(output_dir / "bench.json", 'w') as f:
        json.dump(bench, f, indent=2)

    # Write slice_proofs.json
    slice_proofs = {
        "seed": seed,
        "slices": [{"start": s, "end": e, "length": e-s} for s, e in slices],
        "native_proofs": native_result["slice_tests"],
        "compat_proofs": compat_result["slice_tests"],
    }
    with open(output_dir / "slice_proofs.json", 'w') as f:
        json.dump(slice_proofs, f, indent=2)

    # Build main receipt
    receipt = {
        "kat_version": "2.0",
        "timestamp": timestamp,
        "output_dir": str(output_dir),
        "parameters": {
            "seed": seed,
            "size": size,
            "num_slices": num_slices
        },
        "test_data": {
            "sha256": test_sha256,
            "size": size,
            "generator": "SHAKE256(seed)"
        },
        "build_info": {
            "python_version": sys.version.split()[0],
            "platform": platform.platform(),
            "zlib_version": zlib.ZLIB_VERSION,
            "zlib_runtime_version": getattr(zlib, 'ZLIB_RUNTIME_VERSION', 'unknown'),
            "compression_settings": {
                "level": 9,                  # zlib compression level (1-9)
                "gzip_mtime": 0,             # Fixed for determinism
                "zlib_wbits": 15,            # Default zlib window bits (no override)
                "zlib_strategy": "Z_DEFAULT_STRATEGY"  # Default strategy
            },
            "struct_sizes": {
                "CHUNKINDEX_SIZE": CHUNKINDEX_SIZE,
                "HEADER_SIZE": HEADER_SIZE,
                "SIDECAR_ENTRY_SIZE": SIDECAR_ENTRY_SIZE
            },
            "format_magic": {
                "FIBPI_ZLIB": MAGIC.hex(),
                "SIDECAR": SIDECAR_MAGIC.hex(),
            }
        },
        "gzip_baseline": {
            "compressed_sha256": baseline_result["compressed_sha256"],
            "compression_ratio": baseline_result["compression_ratio"],
            "roundtrip_match": baseline_result["full_roundtrip_match"],
        },
        "fibpi_native": {
            "compressed_sha256": native_result["compressed_sha256"],
            "compression_ratio": native_result["compression_ratio"],
            "num_chunks": native_result["num_chunks"],
            "roundtrip_match": native_result["full_roundtrip_match"],
            "all_slices_pass": native_result["all_slices_pass"],
            "timing": native_result["timing"],
            "memory": native_result["memory"],
        },
        "fibpi_gzip_compat": {
            "gz_sha256": compat_result["gz_sha256"],
            "idx_sha256": compat_result["idx_sha256"],
            "compression_ratio": compat_result["compression_ratio"],
            "num_chunks": compat_result["num_chunks"],
            "roundtrip_match": compat_result["full_roundtrip_match"],
            "gunzip_compatible": compat_result["gunzip_compatible"],
            "all_slices_pass": compat_result["all_slices_pass"],
            "timing": compat_result["timing"],
            "memory": compat_result["memory"],
        },
        "benchmark_summary": bench["comparison"],
        "overall_status": "PASS" if (
            baseline_result["full_roundtrip_match"] and
            native_result["full_roundtrip_match"] and
            native_result["all_slices_pass"] and
            compat_result["full_roundtrip_match"] and
            compat_result["gunzip_compatible"] and
            compat_result["all_slices_pass"]
        ) else "FAIL"
    }

    # Write receipt.json
    with open(output_dir / "receipt.json", 'w') as f:
        json.dump(receipt, f, indent=2)

    if verbose:
        print(f"=" * 60)
        print(f"Overall: {receipt['overall_status']}")
        print()
        print("Benchmark Summary:")
        if bench["comparison"]["native_random_access_speedup"]:
            print(f"  Native random access: {bench['comparison']['native_random_access_speedup']:.1f}x faster than full gzip decompress")
        if bench["comparison"]["compat_random_access_speedup"]:
            print(f"  Compat random access: {bench['comparison']['compat_random_access_speedup']:.1f}x faster than full gzip decompress")
        print()
        print(f"Artifacts written to: {output_dir}")
        print(f"  receipt.json       - Main receipt")
        print(f"  bench.json         - Benchmark data")
        print(f"  slice_proofs.json  - Slice verification")
        print(f"  *.sha256           - Artifact hashes")

    return receipt


def main():
    parser = argparse.ArgumentParser(
        description='FibPi Zlib KAT (Known Answer Test) v2.0'
    )
    parser.add_argument('--seed', default=DEFAULT_SEED,
                        help=f'Deterministic seed (default: {DEFAULT_SEED})')
    parser.add_argument('--size', type=int, default=DEFAULT_SIZE,
                        help=f'Test data size in bytes (default: {DEFAULT_SIZE:,})')
    parser.add_argument('--slices', type=int, default=DEFAULT_NUM_SLICES,
                        help=f'Number of random slices (default: {DEFAULT_NUM_SLICES})')
    parser.add_argument('--receipts-dir', type=Path, default=DEFAULT_RECEIPTS_DIR,
                        help='Directory for receipt output')
    parser.add_argument('-q', '--quiet', action='store_true',
                        help='Suppress output')

    args = parser.parse_args()

    receipt = run_kat(
        seed=args.seed,
        size=args.size,
        num_slices=args.slices,
        receipts_dir=args.receipts_dir,
        verbose=not args.quiet
    )

    return 0 if receipt["overall_status"] == "PASS" else 1


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