#!/usr/bin/env python3
"""
HELIX Codec Registry

Central registry for HELIX compression codecs (algo_id → implementation)
"""
from typing import Dict, Callable, Any, Tuple, List
import hashlib
import sys

class CodecRegistry:
    """Registry for HELIX compression codecs"""

    def __init__(self):
        self._codecs: Dict[str, Dict[str, Any]] = {}

    def register(
        self,
        algo_id: str,
        name: str,
        pack_fn: Callable,
        unpack_fn: Callable,
        description: str = "",
        version: str = "1.0",
        **kwargs
    ):
        """Register a codec with optional block-level random access"""
        self._codecs[algo_id] = {
            "algo_id": algo_id,
            "name": name,
            "pack": pack_fn,
            "unpack": unpack_fn,
            "description": description,
            "version": version,
            # Optional CDC block ops
            "blocks": kwargs.get("blocks"),
            "read_block": kwargs.get("read_block"),
            "write_block": kwargs.get("write_block"),
            "random_access": kwargs.get("blocks") is not None
        }

    def get_codec(self, algo_id: str) -> Dict[str, Any]:
        """Get codec by algo_id"""
        if algo_id not in self._codecs:
            raise ValueError(f"Unknown algo_id: {algo_id}")
        return self._codecs[algo_id]

    def pack(self, algo_id: str, data: bytes, label: str, seed: str, **kwargs) -> Tuple[Any, ...]:
        """Pack data using specified codec"""
        codec = self.get_codec(algo_id)
        return codec["pack"](data, label, seed, **kwargs)

    def unpack(self, algo_id: str, *args, **kwargs) -> bytes:
        """Unpack data using specified codec"""
        codec = self.get_codec(algo_id)
        return codec["unpack"](*args, **kwargs)

    def blocks(self, algo_id: str, capsule: Any, **kwargs) -> int:
        """Get number of blocks in capsule (random access codec only)"""
        codec = self.get_codec(algo_id)
        if not codec.get("blocks"):
            raise NotImplementedError(f"Codec {codec['name']} does not support block-level random access")
        return codec["blocks"](capsule, **kwargs)

    def read_block(self, algo_id: str, capsule: Any, block_idx: int, **kwargs) -> bytes:
        """Read specific block from capsule (random access codec only)"""
        codec = self.get_codec(algo_id)
        if not codec.get("read_block"):
            raise NotImplementedError(f"Codec {codec['name']} does not support block-level random access")
        return codec["read_block"](capsule, block_idx, **kwargs)

    def write_block(self, algo_id: str, capsule: Any, block_idx: int, payload: bytes, **kwargs) -> Any:
        """Write specific block to capsule (random access codec only)"""
        codec = self.get_codec(algo_id)
        if not codec.get("write_block"):
            raise NotImplementedError(f"Codec {codec['name']} does not support block-level random access")
        return codec["write_block"](capsule, block_idx, payload, **kwargs)

    def list_codecs(self) -> List[Dict[str, Any]]:
        """List all registered codecs"""
        return [
            {
                "algo_id": c["algo_id"],
                "name": c["name"],
                "description": c["description"],
                "version": c["version"]
            }
            for c in self._codecs.values()
        ]

# Global registry instance
registry = CodecRegistry()

# ============================================================================
# Codec Implementations
# ============================================================================

def pack_helix64_min(data: bytes, label: str, seed: str, **kwargs) -> Tuple[List[int], str]:
    """
    Pack with HELIX64-MIN (algo_id 0x32)

    Modes:
      - "native" (default): NF1 → RLE+VBE → HELIX64 (pure Helix)
      - "compat": zlib → HELIX64 (legacy)

    kwargs:
      - mode: "native" or "compat" (default: "native")
      - deterministic: True/False (default: True)
    """
    from helix64_min import helix64_pack
    mode = kwargs.pop("mode", "native")  # Pop to avoid passing twice
    return helix64_pack(data, label, seed, mode=mode, **kwargs)

def unpack_helix64_min(payload_b64: str, frame_sizes: List[int], label: str, seed: str, **kwargs) -> bytes:
    """
    Unpack with HELIX64-MIN (algo_id 0x32)

    Modes:
      - "native" (default): HELIX64 → RLE+VBE decode → NF1 unfold
      - "compat": HELIX64 → zlib decompress (legacy)

    kwargs:
      - mode: "native" or "compat" (default: "native")
    """
    from helix64_min import helix64_unpack
    mode = kwargs.pop("mode", "native")  # Pop to avoid passing twice
    return helix64_unpack(payload_b64, frame_sizes, label, seed, mode=mode, **kwargs)

def pack_helix_seed_xor(data: bytes, label: str, seed: str) -> Tuple[bytes, bytes]:
    """
    Pack with HELIX seed-xor (algo_id 0x31)

    Legacy: seed_to_bytes → XOR diff → zlib
    """
    import sys
    import os
    sys.path.insert(0, os.path.expanduser("~/HELIX_RSA_LABS"))

    from helix_regrow import seed_to_bytes
    import zlib

    # Generate pattern from seed
    pattern = seed_to_bytes(seed, len(data))

    # XOR diff
    diff = bytes(a ^ b for a, b in zip(data, pattern))

    # Compress diff
    compressed_diff = zlib.compress(diff, level=9)

    return (diff, compressed_diff)

def unpack_helix_seed_xor(compressed_diff: bytes, length: int, label: str, seed: str) -> bytes:
    """
    Unpack with HELIX seed-xor (algo_id 0x31)

    Reverse: zlib decompress → XOR with seed_to_bytes
    """
    import sys
    import os
    sys.path.insert(0, os.path.expanduser("~/HELIX_RSA_LABS"))

    from helix_regrow import seed_to_bytes
    import zlib

    # Decompress diff
    diff = zlib.decompress(compressed_diff)

    # Regenerate pattern
    pattern = seed_to_bytes(seed, length)

    # XOR to recover original
    data = bytes(a ^ b for a, b in zip(pattern, diff))

    return data

# ============================================================================
# Block-level API implementations (must be defined before registrations)
# ============================================================================

# 0x32 Block API (mutable frames with in-place edits)
def _blocks_32(capsule: Dict[str, Any], block_size: int = 32768, **kwargs) -> int:
    """Get number of frames in 0x32 capsule"""
    frames = capsule.get("frames", [])
    return len(frames)

def _read_block_32(capsule: Dict[str, Any], block_idx: int, **kwargs) -> bytes:
    """
    Read frame from 0x32 capsule (frame-level random access)

    Note: 0x32 frames are variable-size, so block_idx = frame_idx
    """
    from helix64_min import helix64_unpack, FRAME_SIZE

    frames = capsule.get("frames", [])
    if block_idx >= len(frames):
        return b""

    label = capsule.get("label", "response")
    seed = capsule.get("seed_hex", "")
    payload_b64 = capsule.get("payload_b64", "")
    mode = capsule.get("mode", "native")  # Get mode from capsule, not kwargs

    # Decode just the requested frame
    offset = 0
    for i in range(block_idx):
        frame_size = frames[i]
        enc_len = ((frame_size + 2) // 3) * 4  # Standard base64 size formula
        offset += enc_len

    frame_size = frames[block_idx]
    enc_len = ((frame_size + 2) // 3) * 4  # Standard base64 size formula
    frame_b64 = payload_b64[offset:offset + enc_len]

    # Debug: Check if we're extracting the right amount
    if len(frame_b64) != enc_len:
        print(f"Warning: Expected {enc_len} bytes, got {len(frame_b64)} bytes for frame {block_idx}")
        print(f"  Frame size: {frame_size}, Total payload: {len(payload_b64)}, Offset: {offset}")

    # Unpack single frame
    return helix64_unpack(frame_b64, [frame_size], label, seed, mode=mode)

def _write_block_32(capsule: Dict[str, Any], block_idx: int, payload: bytes, **kwargs) -> Dict[str, Any]:
    """
    Write frame to 0x32 capsule (mutable in-place edit)

    Returns new capsule with updated frame.
    """
    from helix64_min import helix64_pack, FRAME_SIZE
    import copy

    capsule_copy = copy.deepcopy(capsule)
    label = capsule_copy.get("label", "response")
    seed = capsule_copy.get("seed_hex", "")
    mode = capsule_copy.get("mode", "native")  # Get mode from capsule, not kwargs

    # Re-encode just this frame
    frame_sizes, new_frame_b64 = helix64_pack(payload, label, seed, mode=mode)

    # Update frame size and rebuild payload
    frames = capsule_copy.get("frames", [])
    frames[block_idx] = frame_sizes[0]

    # Rebuild full payload_b64 with new frame
    # (For simplicity, re-decode all frames and re-encode — optimize later)
    all_frames = []
    for i in range(len(frames)):
        if i == block_idx:
            all_frames.append(payload)
        else:
            all_frames.append(_read_block_32(capsule, i, mode=mode))

    # Re-pack all frames together
    full_data = b"".join(all_frames)
    new_frame_sizes, new_payload_b64 = helix64_pack(full_data, label, seed, mode=mode)

    capsule_copy["frames"] = new_frame_sizes
    capsule_copy["payload_b64"] = new_payload_b64
    capsule_copy["payload_bytes"] = len(new_payload_b64)

    return capsule_copy

# 0x33 Block API (imported from helix_dna_regen.py)
from helix_dna_regen import _blocks_33, _read_block_33, _write_block_33

# ============================================================================
# Register Codecs
# ============================================================================

registry.register(
    algo_id="0x31",
    name="helix-seed-xor",
    pack_fn=pack_helix_seed_xor,
    unpack_fn=unpack_helix_seed_xor,
    description="Legacy HELIX seed pattern + XOR diff + zlib",
    version="1.0"
)

registry.register(
    algo_id="0x32",
    name="helix64-min",
    pack_fn=pack_helix64_min,
    unpack_fn=unpack_helix64_min,
    description="Transport wrapper: zlib → HELIX64 (seed-derived per-frame alphabets)",
    version="1.0",
    blocks=_blocks_32,
    read_block=_read_block_32,
    write_block=_write_block_32
)

# ============================================================================
# 0x33: Pure Seed Regeneration (DNA Compression)
# ============================================================================

def pack_helix_dna_regen(data: bytes, label: str, seed_hex: str, **kw) -> Dict[str, Any]:
    """
    Pack with HELIX DNA regeneration (algo_id 0x33)

    0x33 'pack' is declarative: we don't compress data—we declare how to regrow it.

    Required kw:
      - total_len: int (bytes of the target artifact)
      - regrow_fn: str (name in helix_regrow.REGROW)
      - params: dict (optional) — anything your regrow() needs
    """
    import time
    from base64 import b64encode

    total_len = kw["total_len"]
    regrow_fn = kw.get("reggrow_fn", kw.get("regrow_fn"))  # tolerate typo
    params = kw.get("params", {})

    # Build HXZ1 header
    hdr = bytearray(96)
    hdr[0:4] = b"HXZ1"
    hdr[4] = 1          # version
    hdr[5] = 0x33       # algo_id
    hdr[6:8] = (total_len & 0xFFFF).to_bytes(2, 'little')
    hdr[8:40] = bytes.fromhex(seed_hex)[:32]
    hdr[40:72] = b"\x00" * 32  # sha3 unknown until regen
    hdr[72:80] = int(time.time()).to_bytes(8, 'little')
    hdr[80:96] = label.encode()[:16].ljust(16, b"\x00")

    capsule = {
        "schema": "helix.v1",
        "label": label,
        "algo_id": "0x33",
        "len": int(total_len),
        "hxz_header_b64": b64encode(hdr).decode("ascii"),
        "seed_hex": seed_hex,
        "attest": hashlib.sha3_256(bytes.fromhex(seed_hex)).hexdigest(),
        "regrow_fn": regrow_fn,
        "params": params,
        "compression_method": "helix-dna-regen"
        # NOTE: no payload_b64 / diff_b64 — data is implied by (seed,label,fn,params,len)
    }
    return capsule

def unpack_helix_dna_regen(capsule: Dict[str, Any]) -> bytes:
    """
    Unpack with HELIX DNA regeneration (algo_id 0x33)

    Fully regrows from seed (or can expose streaming for CDC)
    """
    import sys
    import os
    sys.path.insert(0, os.path.expanduser("~/HELIX_RSA_LABS"))

    from helix_regrow import regrow

    return regrow(
        capsule["seed_hex"],
        label=capsule.get("label", "response"),
        offset=0,
        size=capsule["len"],
        fn=capsule["regrow_fn"],
        params=capsule.get("params", {})
    )

registry.register(
    algo_id="0x33",
    name="helix-dna-regen",
    pack_fn=pack_helix_dna_regen,
    unpack_fn=unpack_helix_dna_regen,
    description="Pure seed regeneration: millions:1 DNA compression with true CDC/random access",
    version="1.0",
    blocks=_blocks_33,
    read_block=_read_block_33,
    write_block=_write_block_33
)

# ============================================================================
# 0x36: Recursive Base64 Branching (Random Data Oriented)
# ============================================================================

def pack_helix_branch64(data: bytes, label: str, seed_hex: str, **kw) -> Dict[str, Any]:
    """
    Pack with HELIX Branch64 (algo_id 0x36)

    Recursive Base64 Branching → HELIX64 for high-entropy/random data
    """
    from helix_branch64 import pack_helix_branch64 as pack_rbb
    return pack_rbb(data, label, seed_hex, **kw)

def unpack_helix_branch64(capsule: Dict[str, Any]) -> bytes:
    """
    Unpack with HELIX Branch64 (algo_id 0x36)
    """
    from helix_branch64 import unpack_helix_branch64 as unpack_rbb
    return unpack_rbb(capsule)

registry.register(
    algo_id="0x36",
    name="helix-branch64",
    pack_fn=pack_helix_branch64,
    unpack_fn=unpack_helix_branch64,
    description="Recursive Base64 Branching → HELIX64 (random-data oriented)",
    version="1.0"
)

# ============================================================================
# Convenience Functions
# ============================================================================

def pack(data: bytes, label: str = "response", algo_id: str = "0x32", **kwargs) -> Dict[str, Any]:
    """
    Pack data with specified codec (convenience wrapper)

    Returns capsule dict ready for transmission
    """
    # Generate seed from content
    seed = hashlib.sha256(data).hexdigest()[:64]

    # SHA3 hash for verification
    sha3_hash = hashlib.sha3_256(data).hexdigest()

    # Pack with codec
    if algo_id == "0x32":
        frame_sizes, payload_b64 = registry.pack(algo_id, data, label, seed, **kwargs)

        return {
            "schema": "helix.v1",
            "algo_id": algo_id,
            "label": label,
            "len": len(data),
            "encoding": "utf-8",
            "seed_hex": seed,
            "attest": hashlib.sha3_256(bytes.fromhex(seed)).hexdigest(),
            "sha3": sha3_hash,
            "frames": frame_sizes,
            "payload_b64": payload_b64,
            "payload_bytes": len(payload_b64),
            "helix_ratio": round(len(data) / len(payload_b64), 2) if payload_b64 else 0.0,
            "compression_method": "helix64-min",
            "mode": kwargs.get("mode", "native")
        }

    elif algo_id == "0x31":
        diff, compressed_diff = registry.pack(algo_id, data, label, seed)

        from base64 import b64encode
        return {
            "schema": "helix.v1",
            "algo_id": algo_id,
            "label": label,
            "len": len(data),
            "encoding": "utf-8",
            "seed_hex": seed,
            "attest": hashlib.sha3_256(bytes.fromhex(seed)).hexdigest(),
            "sha3": sha3_hash,
            "diff_b64": b64encode(compressed_diff).decode('ascii'),
            "diff_bytes": len(compressed_diff),
            "helix_ratio": round(len(data) / len(compressed_diff), 2) if compressed_diff else 0.0,
            "compression_method": "helix-seed-xor"
        }

    elif algo_id == "0x36":
        # Pack with Recursive Base64 Branching
        capsule = registry.pack(algo_id, data, label, seed)
        return capsule  # Already returns full capsule dict

    else:
        raise ValueError(f"Unknown algo_id: {algo_id}")

def unpack(capsule: Dict[str, Any]) -> bytes:
    """
    Unpack capsule (convenience wrapper)

    Args:
        capsule: Can be either a dict (capsule object) or algo_id string

    Returns original data bytes
    """
    # Handle both registry.unpack(capsule_dict) and registry.unpack(algo_id, ...)
    if isinstance(capsule, str):
        # Called as registry.unpack("0x32", ...) - this is handled by CodecRegistry.unpack()
        raise TypeError("Use CodecRegistry.unpack() for raw codec calls, not the convenience unpack()")

    algo_id = capsule.get("algo_id", "0x32")

    if algo_id == "0x32":
        data = registry.unpack(
            algo_id,
            capsule["payload_b64"],
            capsule["frames"],
            capsule.get("label", "response"),
            capsule["seed_hex"],
            mode=capsule.get("mode", "native")
        )
    elif algo_id == "0x31":
        from base64 import b64decode
        compressed_diff = b64decode(capsule.get("diff_b64", capsule.get("scrambled_b64", "")))

        data = registry.unpack(
            algo_id,
            compressed_diff,
            capsule["len"],
            capsule.get("label", "response"),
            capsule["seed_hex"]
        )
    elif algo_id == "0x33":
        # DNA regeneration
        data = unpack_helix_dna_regen(capsule)
    else:
        raise ValueError(f"Unknown algo_id: {algo_id}")

    # Verify SHA3 if present
    if "sha3" in capsule:
        sha3_hash = hashlib.sha3_256(data).hexdigest()
        if sha3_hash != capsule["sha3"]:
            raise ValueError(f"SHA3 mismatch: expected {capsule['sha3']}, got {sha3_hash}")

    return data

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

if __name__ == "__main__":
    print("🧬 HELIX Codec Registry")
    print("=" * 60)

    for codec in registry.list_codecs():
        print(f"\n📦 {codec['algo_id']}: {codec['name']}")
        print(f"   {codec['description']}")
        print(f"   Version: {codec['version']}")

    print("\n" + "=" * 60)

    # Test roundtrip
    test_data = b"HELIX codec registry test: fractal compression via seed-derived tables!"
    print(f"\n🧪 Testing roundtrip with {len(test_data)} bytes...")

    for algo_id in ["0x32", "0x31"]:
        print(f"\n   {algo_id} ({registry.get_codec(algo_id)['name']}):")

        capsule = pack(test_data, algo_id=algo_id)
        print(f"      Compressed: {capsule.get('payload_bytes', capsule.get('diff_bytes', 0))} bytes")
        print(f"      Ratio: {capsule['helix_ratio']}:1")

        recovered = unpack(capsule)
        print(f"      Verified: {recovered == test_data} ✅" if recovered == test_data else "      Verified: ❌ FAIL")

    print("\n✅ Registry test complete")
