#!/usr/bin/env python3
"""
test_hxz_roundtrip.py
=====================
Real, LOSSLESS round-trip test:

  1. Create known tensor
  2. Encode to .hxz via hex_lossless_v1
  3. Decode back from .hxz via matching decoder
  4. Verify SHA256 and element-wise equality
"""

import hashlib
import sys
from pathlib import Path

import numpy as np

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

from tools.helix_encode_lossless import _write_hxz_lossless
from tools.helix_decode_lossless import decode_hxz_to_tensor


def sha256_bytes(b: bytes) -> str:
    return hashlib.sha256(b).hexdigest()


def main() -> None:
    print("=" * 80)
    print("ROUND-TRIP TEST - LOSSLESS ENCODE → DECODE → VERIFY")
    print("=" * 80)

    # 1. Create deterministic test tensor
    print("\n1. Creating test tensor (256, 256) float32...")
    rng = np.random.default_rng(1337)
    tensor = rng.standard_normal(size=(256, 256), dtype=np.float32)
    raw = tensor.tobytes()
    sha_orig = sha256_bytes(raw)
    print(f"   ✅ Created: {len(raw):,} bytes")
    print(f"   ✅ SHA256 (original): {sha_orig}")

    # 2. Encode to .hxz
    tmp_dir = Path("tests/tmp")
    tmp_dir.mkdir(parents=True, exist_ok=True)
    capsule_path = tmp_dir / "roundtrip_test.hxz"

    print("\n2. Encoding to .hxz capsule (hex_lossless_v1)...")
    _write_hxz_lossless(tensor, capsule_path)
    print(f"   ✅ Wrote capsule: {capsule_path} ({capsule_path.stat().st_size:,} bytes)")

    # 3. Decode from .hxz
    print("\n3. Decoding from .hxz capsule...")
    tensor_regen = decode_hxz_to_tensor(capsule_path)
    raw_regen = tensor_regen.tobytes()
    sha_regen = sha256_bytes(raw_regen)
    print(f"   ✅ Decoded tensor shape: {tensor_regen.shape}")
    print(f"   ✅ SHA256 (regen): {sha_regen}")

    # 4. Verify equality
    print("\n4. Verifying round-trip equality...")
    if sha_orig != sha_regen:
        raise SystemExit(
            f"❌ SHA256 mismatch:\n   original={sha_orig}\n   regen   ={sha_regen}"
        )

    if not np.array_equal(tensor, tensor_regen):
        raise SystemExit("❌ Tensor contents mismatch despite SHA match (should not happen)")

    print("   ✅ SHA256 match")
    print("   ✅ Element-wise equality verified")
    print("\n🎉 SUCCESS: Lossless Helix .hxz codec round-trip is proven.\n")


if __name__ == "__main__":
    main()
