#!/usr/bin/env python3
"""
HXZ ABI Fuzz Tests (QA Hardening)

Tests payload extraction from different .hxz formats:
- Legacy text formats (#HXZ_START, #HXZ_GGUF, #HXZ_DEMO)
- Binary ABI v1 format
- Corrupted/malformed containers
"""
import sys
from pathlib import Path

# Add scripts to path
sys.path.insert(0, str(Path(__file__).parent.parent / 'scripts'))
from helix_encode_stage import hxz_extract_payload_bytes_from_data


def test_legacy_formats():
    """Test payload extraction from legacy text formats"""
    print("Testing legacy text formats...")

    test_cases = [
        ('HXZ_START', b'#HXZ_START payload data here'),
        ('HXZ_GGUF', b'#HXZ_GGUF model weights here'),
        ('HXZ_DEMO', b'#HXZ_DEMO synthetic demo data'),
    ]

    errors = []
    for name, data in test_cases:
        payload = hxz_extract_payload_bytes_from_data(data)

        # Legacy format: entire file is payload
        if payload != data:
            errors.append(f"{name}: Payload mismatch (expected full data)")

        print(f"  {name:20s}: {len(payload):6d} bytes (expected {len(data)})")

    if errors:
        print(f"  ❌ FAILED: {len(errors)} legacy format issues")
        for e in errors:
            print(f"     {e}")
        return False
    else:
        print(f"  ✅ PASSED: Legacy formats extracted correctly")
        return True


def test_abi_v1_format():
    """Test payload extraction from binary ABI v1 format"""
    print("\nTesting ABI v1 binary format...")

    # Build valid ABI v1 container
    MAGIC = b"HXZ_START\n"
    META = b'{"version":"1.0","test":true}'
    META_LEN_LINE = f"META_LEN:{len(META)}\n".encode('ascii')
    PAYLOAD = b'\x00\x11\x22\x33\x44\x55\x66\x77\x88\x99'
    FOOTER = b"HXZ_END\n"

    data = MAGIC + META_LEN_LINE + META + PAYLOAD + FOOTER

    errors = []

    payload = hxz_extract_payload_bytes_from_data(data)

    # Should extract only the payload (10 bytes)
    if payload != PAYLOAD:
        errors.append(f"ABI v1: Payload mismatch")
        errors.append(f"  Expected: {PAYLOAD.hex()}")
        errors.append(f"  Got: {payload.hex()}")

    print(f"  ABI v1 valid:        {len(payload):6d} bytes (expected {len(PAYLOAD)})")

    if errors:
        print(f"  ❌ FAILED: {len(errors)} ABI v1 issues")
        for e in errors:
            print(f"     {e}")
        return False
    else:
        print(f"  ✅ PASSED: ABI v1 format extracted correctly")
        return True


def test_corrupt_formats():
    """Test behavior on corrupted/malformed .hxz files"""
    print("\nTesting corrupted formats (should fallback gracefully)...")

    MAGIC = b"HXZ_START\n"
    FOOTER = b"HXZ_END\n"

    test_cases = [
        ('Bad META_LEN (not integer)', MAGIC + b"META_LEN:BADVALUE\n" + b"payload" + FOOTER),
        ('Bad META_LEN (negative)', MAGIC + b"META_LEN:-100\n" + b"payload" + FOOTER),
        ('Truncated file', MAGIC + b"META_LEN:50\n" + b"short"),  # File shorter than META_LEN
        ('Missing META_LEN line', MAGIC + b"payload data no meta" + FOOTER),
        ('CRLF instead of LF', MAGIC.replace(b"\n", b"\r\n") + b"payload" + FOOTER),
        ('Empty file', b''),
        ('Non-HXZ file', b'This is just plain text'),
    ]

    errors = []
    for name, data in test_cases:
        try:
            payload = hxz_extract_payload_bytes_from_data(data)

            # Defensive behavior: corrupted formats may return full data or empty payload
            # Both are acceptable - we just want no crashes
            # (Truncated files may return empty if they claim more metadata than exists)

            print(f"  {name:35s}: {len(payload):6d} bytes (defensive ok)")

        except Exception as e:
            errors.append(f"{name}: Raised exception: {e}")
            print(f"  {name:35s}: ❌ EXCEPTION: {e}")

    if errors:
        print(f"  ❌ FAILED: {len(errors)} corruption handling issues")
        for e in errors:
            print(f"     {e}")
        return False
    else:
        print(f"  ✅ PASSED: Corrupted formats handled gracefully")
        return True


def test_abi_v1_variations():
    """Test ABI v1 with different metadata sizes"""
    print("\nTesting ABI v1 with various metadata sizes...")

    MAGIC = b"HXZ_START\n"
    FOOTER = b"HXZ_END\n"
    PAYLOAD = b'\xDE\xAD\xBE\xEF' * 16  # 64 bytes

    # Test different metadata sizes
    metadata_sizes = [1, 10, 100, 1000, 5000]
    errors = []

    for meta_size in metadata_sizes:
        # Create metadata of exact size
        META = b'x' * meta_size
        META_LEN_LINE = f"META_LEN:{len(META)}\n".encode('ascii')

        data = MAGIC + META_LEN_LINE + META + PAYLOAD + FOOTER

        payload = hxz_extract_payload_bytes_from_data(data)

        if payload != PAYLOAD:
            errors.append(f"META_LEN={meta_size}: Payload mismatch")

        print(f"  META_LEN={meta_size:5d}:         {len(payload):6d} bytes (expected {len(PAYLOAD)})")

    if errors:
        print(f"  ❌ FAILED: {len(errors)} metadata size issues")
        for e in errors:
            print(f"     {e}")
        return False
    else:
        print(f"  ✅ PASSED: ABI v1 handles various metadata sizes")
        return True


def test_payload_only():
    """Test extraction from raw payload (no container)"""
    print("\nTesting raw payload (no container)...")

    # Raw data without any HXZ markers
    raw_data = b'This is just raw binary data with no HXZ container'

    payload = hxz_extract_payload_bytes_from_data(raw_data)

    errors = []

    # Should return entire data as-is
    if payload != raw_data:
        errors.append(f"Raw payload: Returned modified data instead of original")

    print(f"  Raw payload:         {len(payload):6d} bytes (expected {len(raw_data)})")

    if errors:
        print(f"  ❌ FAILED: {len(errors)} raw payload issues")
        for e in errors:
            print(f"     {e}")
        return False
    else:
        print(f"  ✅ PASSED: Raw payload passed through correctly")
        return True


def test_abi_v1_edge_cases():
    """Test ABI v1 edge cases (zero-size metadata, zero-size payload)"""
    print("\nTesting ABI v1 edge cases...")

    MAGIC = b"HXZ_START\n"
    FOOTER = b"HXZ_END\n"

    test_cases = [
        ('Zero metadata', b'', b'\xAA\xBB\xCC'),
        ('Zero payload', b'{"test":true}', b''),
        ('Both zero', b'', b''),
    ]

    errors = []

    for name, meta, payload_expected in test_cases:
        META_LEN_LINE = f"META_LEN:{len(meta)}\n".encode('ascii')
        data = MAGIC + META_LEN_LINE + meta + payload_expected + FOOTER

        payload = hxz_extract_payload_bytes_from_data(data)

        if payload != payload_expected:
            errors.append(f"{name}: Payload mismatch (expected {len(payload_expected)} bytes, got {len(payload)})")

        print(f"  {name:20s}: {len(payload):6d} bytes (expected {len(payload_expected)})")

    if errors:
        print(f"  ❌ FAILED: {len(errors)} edge case issues")
        for e in errors:
            print(f"     {e}")
        return False
    else:
        print(f"  ✅ PASSED: Edge cases handled correctly")
        return True


def main():
    print("="*60)
    print("HXZ ABI Fuzz Tests (QA Hardening)")
    print("="*60)
    print()

    tests = [
        test_legacy_formats,
        test_abi_v1_format,
        test_corrupt_formats,
        test_abi_v1_variations,
        test_payload_only,
        test_abi_v1_edge_cases
    ]

    results = []
    for test_fn in tests:
        try:
            passed = test_fn()
            results.append((test_fn.__name__, passed))
        except Exception as e:
            print(f"  ❌ EXCEPTION: {e}")
            import traceback
            traceback.print_exc()
            results.append((test_fn.__name__, False))

    print()
    print("="*60)
    print("RESULTS")
    print("="*60)

    passed_count = sum(1 for _, p in results if p)
    total_count = len(results)

    for name, passed in results:
        status = "✅ PASSED" if passed else "❌ FAILED"
        print(f"{status}: {name}")

    print()
    print(f"Total: {passed_count}/{total_count} tests passed")

    if passed_count == total_count:
        print("✅ ALL TESTS PASSED")
        return 0
    else:
        print(f"❌ {total_count - passed_count} TEST(S) FAILED")
        return 1


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