#!/usr/bin/env python3
"""
WO#SV4.A: Seed+Vault Codec Implementation

Based on the historical ChatGPT codec design extracted in WO#SV1-SV3.
Implements the exact primitives described:
- SHA256 for content hashing
- Base64 → Ternary → FibPi3D mapping
- SHAKE256 with domain tag for regeneration

This codec is self-contained: no external model/cache dependencies.
"""

import hashlib
import base64
import math
import json
import zlib
from typing import Tuple, Dict, Any

# Domain tag from the historical specs
DOMAIN_TAG = b"HELIX\x1b\x00v1\x00REGEN\x00"

# FibPi3D parameters
PHI = 1.618033988749895
GOLDEN_ANGLE = 2.39996323  # ~137.5 degrees in radians


def sha256_hash(data: bytes) -> bytes:
    """Compute SHA256 hash of data."""
    return hashlib.sha256(data).digest()


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


def data_to_base64(data: bytes) -> str:
    """Convert data to base64 string."""
    return base64.b64encode(data).decode('ascii')


def base64_to_data(b64: str) -> bytes:
    """Convert base64 string back to data."""
    return base64.b64decode(b64.encode('ascii'))


def base64_to_ternary(b64: str) -> str:
    """
    Convert base64 to ternary using the DNA composter mapping:
        '0' if c < 'I' else '1' if c < 'a' else '2'

    This is the KEY mapping from the historical ChatGPT specs.
    """
    return ''.join(['0' if c < 'I' else '1' if c < 'a' else '2' for c in b64])


def ternary_to_int(tern: str) -> int:
    """Convert ternary string to integer value."""
    val = 0
    for c in tern:
        val = val * 3 + int(c)
    return val


def ternary_to_fibpi3d(tern: str) -> Tuple[float, float, float]:
    """
    Convert ternary string to FibPi3D coordinates using golden ratio spiral.

    Returns (x, y, z) coordinates on a unit sphere-like structure.
    """
    tern_val = ternary_to_int(tern)

    # Golden angle spiral mapping
    angle = (tern_val * PHI * GOLDEN_ANGLE) % (2 * math.pi)
    z_norm = (tern_val % 1000000) / 1000000.0

    r = math.sqrt(1 - z_norm * z_norm) if z_norm < 1 else 0
    x = r * math.cos(angle)
    y = r * math.sin(angle)
    z = z_norm

    return (x, y, z)


def shake256_regrow(seed: bytes, size: int) -> bytes:
    """
    Use SHAKE256 with domain tag to regenerate bytes from seed.

    This is the core regeneration primitive from the historical specs.
    """
    ctx = hashlib.shake_256(DOMAIN_TAG + seed)
    return ctx.digest(size)


def compute_entropy(data: bytes) -> float:
    """Compute Shannon entropy of data in bits per byte."""
    if not data:
        return 0.0

    freq = {}
    for b in data:
        freq[b] = freq.get(b, 0) + 1

    length = len(data)
    entropy = 0.0
    for count in freq.values():
        p = count / length
        if p > 0:
            entropy -= p * math.log2(p)

    return entropy


def chunk_data(data: bytes, chunk_size: int = 4096) -> list:
    """Split data into chunks for processing."""
    return [data[i:i+chunk_size] for i in range(0, len(data), chunk_size)]


# =============================================================================
# ENCODER
# =============================================================================

def encode_to_seed_vault(input_bytes: bytes) -> Tuple[bytes, bytes, Dict[str, Any]]:
    """
    Encode input bytes into (seed, vault, meta).

    The encoding strategy based on historical specs:
    1. Compute SHA256 hash of input (anchor)
    2. Convert hash to base64 → ternary → FibPi3D coords
    3. Compress input with zlib
    4. Create seed = SHAKE256 initialization from hash
    5. Create vault = compressed data + routing info
    6. Meta = sizes, entropy, coordinates

    For true "regeneration", we store the compressed data in the vault.
    The seed carries the cryptographic anchor and FibPi3D position.

    NOTE: This is a HYBRID codec - it's not magic compression.
    The vault carries the actual data (compressed), while the seed
    carries the verification/routing information.
    """
    # Step 1: Compute content hash (the anchor)
    content_sha = sha256_hash(input_bytes)
    content_sha_hex = content_sha.hex()

    # Step 2: Derive FibPi3D coordinates from hash
    content_b64 = data_to_base64(content_sha)
    content_ternary = base64_to_ternary(content_b64)
    fibpi3d_coords = ternary_to_fibpi3d(content_ternary)

    # Step 3: Compress the input data
    compressed = zlib.compress(input_bytes, level=9)

    # Step 4: Create seed
    # Seed = domain-tagged hash that can regenerate verification bytes
    seed_data = {
        "version": "1.0",
        "sha256": content_sha_hex,
        "size": len(input_bytes),
        "ternary_prefix": content_ternary[:32],  # First 32 ternary digits
        "fibpi3d": {
            "x": round(fibpi3d_coords[0], 8),
            "y": round(fibpi3d_coords[1], 8),
            "z": round(fibpi3d_coords[2], 8)
        }
    }
    seed_bytes = json.dumps(seed_data, separators=(',', ':')).encode('utf-8')

    # Step 5: Create vault
    # Vault = compressed data + verification info
    # In the historical design, this was the "symbol_map" + actual payload
    vault_header = {
        "type": "seed_vault_v1",
        "compressed_size": len(compressed),
        "original_size": len(input_bytes),
        "compression": "zlib",
        "entropy": round(compute_entropy(input_bytes), 4)
    }
    vault_header_bytes = json.dumps(vault_header, separators=(',', ':')).encode('utf-8')

    # Vault format: [4-byte header length][header JSON][compressed data]
    vault_bytes = (
        len(vault_header_bytes).to_bytes(4, 'little') +
        vault_header_bytes +
        compressed
    )

    # Step 6: Meta (returned separately for inspection)
    meta = {
        "original_size": len(input_bytes),
        "seed_size": len(seed_bytes),
        "vault_size": len(vault_bytes),
        "compressed_size": len(compressed),
        "sha256": content_sha_hex,
        "entropy_bits_per_byte": round(compute_entropy(input_bytes), 4),
        "fibpi3d": {
            "x": fibpi3d_coords[0],
            "y": fibpi3d_coords[1],
            "z": fibpi3d_coords[2]
        },
        "compression_ratio": len(input_bytes) / len(vault_bytes) if len(vault_bytes) > 0 else 0
    }

    return seed_bytes, vault_bytes, meta


# =============================================================================
# DECODER
# =============================================================================

def decode_from_seed_vault(seed_bytes: bytes, vault_bytes: bytes, meta: Dict[str, Any] = None) -> bytes:
    """
    Decode (seed, vault, meta) back to original bytes.

    The decoding process:
    1. Parse seed JSON to get expected SHA256 and size
    2. Parse vault header + extract compressed data
    3. Decompress data
    4. Verify SHA256 matches seed
    5. Return original data

    Raises ValueError if verification fails.
    """
    # Step 1: Parse seed
    seed_data = json.loads(seed_bytes.decode('utf-8'))
    expected_sha = seed_data["sha256"]
    expected_size = seed_data["size"]

    # Step 2: Parse vault
    header_len = int.from_bytes(vault_bytes[:4], 'little')
    vault_header = json.loads(vault_bytes[4:4+header_len].decode('utf-8'))
    compressed = vault_bytes[4+header_len:]

    # Step 3: Decompress
    if vault_header["compression"] == "zlib":
        decompressed = zlib.decompress(compressed)
    else:
        raise ValueError(f"Unknown compression: {vault_header['compression']}")

    # Step 4: Verify
    actual_sha = sha256_hex(decompressed)
    if actual_sha != expected_sha:
        raise ValueError(
            f"SHA256 mismatch! Expected {expected_sha[:16]}..., got {actual_sha[:16]}..."
        )

    if len(decompressed) != expected_size:
        raise ValueError(
            f"Size mismatch! Expected {expected_size}, got {len(decompressed)}"
        )

    return decompressed


# =============================================================================
# VERIFICATION HELPERS
# =============================================================================

def verify_round_trip(data: bytes) -> Tuple[bool, Dict[str, Any]]:
    """
    Test that encode → decode returns exactly the original data.

    Returns (success, info_dict).
    """
    try:
        seed, vault, meta = encode_to_seed_vault(data)
        decoded = decode_from_seed_vault(seed, vault, meta)

        success = (decoded == data)

        return success, {
            "success": success,
            "original_size": len(data),
            "seed_size": len(seed),
            "vault_size": len(vault),
            "total_encoded": len(seed) + len(vault),
            "compression_ratio": len(data) / (len(seed) + len(vault)),
            "sha256_match": sha256_hex(decoded) == sha256_hex(data),
            "size_match": len(decoded) == len(data),
            "byte_match": decoded == data
        }
    except Exception as e:
        return False, {
            "success": False,
            "error": str(e)
        }


def test_determinism(data: bytes, iterations: int = 3) -> bool:
    """
    Verify that encoding is deterministic across multiple runs.
    """
    results = []
    for _ in range(iterations):
        seed, vault, meta = encode_to_seed_vault(data)
        results.append((seed, vault))

    # All results should be identical
    first_seed, first_vault = results[0]
    for seed, vault in results[1:]:
        if seed != first_seed or vault != first_vault:
            return False
    return True


def test_mutation_sensitivity(data: bytes) -> Dict[str, bool]:
    """
    Verify that small mutations in seed/vault produce different decodes.
    """
    seed, vault, meta = encode_to_seed_vault(data)

    results = {}

    # Test seed mutation
    if len(seed) > 10:
        mutated_seed = bytearray(seed)
        mutated_seed[10] ^= 0x01  # Flip one bit
        try:
            decoded = decode_from_seed_vault(bytes(mutated_seed), vault, meta)
            results["seed_mutation_detected"] = (decoded != data)
        except Exception:
            results["seed_mutation_detected"] = True  # Verification failed = good

    # Test vault mutation - target the compressed payload, not header
    # Vault format: [4-byte header length][header JSON][compressed data]
    header_len = int.from_bytes(vault[:4], 'little')
    payload_start = 4 + header_len + 10  # Skip header + some bytes into payload

    if len(vault) > payload_start + 10:
        mutated_vault = bytearray(vault)
        # Corrupt multiple bytes in compressed payload to ensure zlib detects it
        for i in range(5):
            if payload_start + i < len(mutated_vault):
                mutated_vault[payload_start + i] ^= 0xFF
        try:
            decoded = decode_from_seed_vault(seed, bytes(mutated_vault), meta)
            results["vault_mutation_detected"] = (decoded != data)
        except Exception:
            results["vault_mutation_detected"] = True  # Verification failed = good
    else:
        results["vault_mutation_detected"] = True  # Too small to test meaningfully

    return results


# =============================================================================
# CLI
# =============================================================================

if __name__ == "__main__":
    import sys

    print("=" * 70)
    print("SEED+VAULT CODEC TEST")
    print("=" * 70)
    print()

    # Test with a simple string
    test_data = b"Hello, this is a test of the Seed+Vault codec from historical specs!"

    print(f"Input: {test_data[:50]}...")
    print(f"Input size: {len(test_data)} bytes")
    print()

    # Encode
    seed, vault, meta = encode_to_seed_vault(test_data)
    print("ENCODE RESULTS:")
    print(f"  Seed size: {len(seed)} bytes")
    print(f"  Vault size: {len(vault)} bytes")
    print(f"  Total: {len(seed) + len(vault)} bytes")
    print(f"  Compression ratio: {meta['compression_ratio']:.2f}x")
    print(f"  SHA256: {meta['sha256'][:32]}...")
    print(f"  FibPi3D: ({meta['fibpi3d']['x']:.4f}, {meta['fibpi3d']['y']:.4f}, {meta['fibpi3d']['z']:.4f})")
    print()

    # Decode
    decoded = decode_from_seed_vault(seed, vault, meta)
    print("DECODE RESULTS:")
    print(f"  Decoded size: {len(decoded)} bytes")
    print(f"  Match: {decoded == test_data}")
    print()

    # Round-trip test
    success, info = verify_round_trip(test_data)
    print(f"ROUND-TRIP: {'PASS' if success else 'FAIL'}")
    print()

    # Determinism test
    det = test_determinism(test_data)
    print(f"DETERMINISM: {'PASS' if det else 'FAIL'}")
    print()

    # Mutation test
    mut = test_mutation_sensitivity(test_data)
    print(f"MUTATION SENSITIVITY:")
    for k, v in mut.items():
        print(f"  {k}: {'PASS' if v else 'FAIL'}")
    print()

    print("=" * 70)
    print("CODEC MODULE READY")
    print("=" * 70)
