#!/usr/bin/env python3
"""
DNA Seed Loader for text-format HXZ files

!!! DISPROVEN - DO NOT USE FOR MODEL REGENERATION !!!

This module was DISPROVEN on 2026-01-02. DNA seeds (#HXZ_START format)
are PSEUDO-RANDOM GENERATORS via SHAKE256 expansion, NOT encoders.
They produce DIFFERENT tensors than the original model weights.

Fidelity check results:
  - 3/3 tensors tested: Output != original model
  - The underscore-named files (blk_0_*) are NOT compressed copies
    of the dot-named files (blk.0.*)

DO NOT:
  - Delete HuggingFace cache based on DNA seed "regeneration"
  - Claim compression ratios based on this format

USE INSTEAD:
  - CDNA v2 with outlier sidecars (PROVEN, max error < 0.001)
  - Chunk Vault CAS for file-level deduplication

Receipt: receipts/fidelity_checks/fidelity_check_multi_tensor.json

---

Parses HelixDNA text format with:
- Header: entropy, sha256
- @SEG anchors with FibPi3D coordinates
- folded_braid() symbol sequence

Uses hierarchical decoding:
  Level 1: a-h → 3 bits each
  Level 2: A-H → 9 bits each (3 level-1 symbols)

Expands symbols to binary using SHAKE256 sponge.

Usage:
    from helix_cdc.regrow.dna_seed_loader import load_dna_seed, parse_dna_seed

    # Parse without loading tensors
    info = parse_dna_seed("path/to/seed.hxz")

    # Load and expand to state dict
    state_dict, metadata = load_dna_seed("path/to/seed.hxz")
"""

from __future__ import annotations

import hashlib
import re
import warnings
import zlib
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

# DISPROVEN WARNING - DNA seeds cannot regenerate original model weights
# Moved to function-level to avoid noise on every import (WO-ECHO-MEMORY-PIPELINE-02)
_WARNED = False


def _warn_once():
    """Emit disproven warning once per session."""
    global _WARNED
    if not _WARNED:
        warnings.warn(
            "dna_seed_loader is DISPROVEN for model regeneration. DNA seeds produce "
            "PSEUDO-RANDOM tensors via SHAKE256, NOT original weights. "
            "Use CDNA v2 + sidecars instead. See receipts/fidelity_checks/",
            UserWarning,
            stacklevel=3,
        )
        _WARNED = True


@dataclass
class DNASeedInfo:
    """Parsed DNA seed information."""
    model_name: str = ""
    timestamp: str = ""
    entropy: float = 0.0
    sha256: str = ""
    segments: List[Dict[str, Any]] = field(default_factory=list)
    braid: List[str] = field(default_factory=list)
    raw_text: str = ""


# Hierarchical encoding tables from CLAUDE.md
LETTER_TO_BITS = {
    'a': '000', 'b': '001', 'c': '010', 'd': '011',
    'e': '100', 'f': '101', 'g': '110', 'h': '111',
}

# Uppercase = level 2 (expands to 3 lowercase)
SYMBOL_TO_LETTERS = {
    'A': 'aaa', 'B': 'bbb', 'C': 'ccc', 'D': 'ddd',
    'E': 'eee', 'F': 'fff', 'G': 'ggg', 'H': 'hhh',
}

# Level 3 glyphs (optional, for future use)
GLYPH_TO_SYMBOLS = {
    '@': 'AAA', '#': 'BBB', '$': 'CCC', '%': 'DDD',
    '^': 'EEE', '&': 'FFF', '*': 'GGG', '!': 'HHH',
}


def parse_dna_seed(hxz_path: str) -> DNASeedInfo:
    """
    Parse DNA seed file and extract structure without expanding.

    Args:
        hxz_path: Path to .hxz file with text DNA format

    Returns:
        DNASeedInfo with parsed fields
    """
    _warn_once()
    path = Path(hxz_path)
    if not path.exists():
        raise FileNotFoundError(f"DNA seed not found: {hxz_path}")

    with open(path, "rb") as f:
        data = f.read()

    # Check for text format
    if not data.startswith(b"#HXZ_START"):
        raise ValueError(f"Not a text-format DNA seed: {hxz_path}")

    # Find header end and decompress blob
    sep_idx = data.find(b"#====")
    if sep_idx == -1:
        raise ValueError("Missing #==== separator in DNA seed")

    # Skip separator + newline
    blob_start = sep_idx + 5
    if blob_start < len(data) and data[blob_start:blob_start+1] == b'\n':
        blob_start += 1

    blob = data[blob_start:]

    # Decompress
    try:
        text = zlib.decompress(blob).decode("utf-8")
    except zlib.error:
        text = blob.decode("utf-8", errors="replace")

    return _parse_dna_text(text)


def _parse_dna_text(text: str) -> DNASeedInfo:
    """Parse the decompressed DNA text content."""
    info = DNASeedInfo(raw_text=text)

    # Parse header comments
    for line in text.split('\n'):
        line = line.strip()
        if line.startswith("# HelixDNA Braid for "):
            info.model_name = line.split("for ")[-1]
        elif line.startswith("# Timestamp:"):
            info.timestamp = line.split(":", 1)[-1].strip()
        elif line.startswith("# Entropy:"):
            try:
                info.entropy = float(line.split(":")[-1].strip().split()[0])
            except ValueError:
                pass

    # Parse @SEG anchors
    seg_pattern = r'@SEG(\d+)\s*\{([^}]+)\}'
    for match in re.finditer(seg_pattern, text):
        seg_num = match.group(1)
        seg_data = match.group(2)

        segment = {"id": f"SEG{seg_num}"}

        # Parse X, Y, Z coordinates
        for coord in ['X', 'Y', 'Z']:
            coord_match = re.search(rf'{coord}=([^,;}}]+)', seg_data)
            if coord_match:
                try:
                    segment[coord] = float(coord_match.group(1))
                except ValueError:
                    segment[coord] = coord_match.group(1)

        # Parse SHA
        sha_match = re.search(r'SHA=([a-f0-9]+)', seg_data)
        if sha_match:
            segment['sha'] = sha_match.group(1)
            info.sha256 = sha_match.group(1)

        info.segments.append(segment)

    # Parse folded_braid()
    braid_match = re.search(r'folded_braid\(\)\s*=\s*\(([\s\S]*?)\)', text)
    if braid_match:
        braid_content = braid_match.group(1)
        # Extract symbols (letters separated by commas/whitespace)
        symbols = re.findall(r'[a-hA-H]', braid_content)
        info.braid = symbols

    return info


def decode_braid_to_bits(braid: List[str]) -> str:
    """
    Decode folded_braid symbols to bit string.

    Level 1: a-h → 3 bits
    Level 2: A-H → 9 bits (3 level-1 chars)
    """
    bits = []
    for symbol in braid:
        if symbol in SYMBOL_TO_LETTERS:
            # Level 2: expand to 3 level-1 symbols
            letters = SYMBOL_TO_LETTERS[symbol]
            for letter in letters:
                bits.append(LETTER_TO_BITS[letter])
        elif symbol in LETTER_TO_BITS:
            # Level 1: direct mapping
            bits.append(LETTER_TO_BITS[symbol])
        else:
            # Unknown symbol - skip or use default
            pass

    return ''.join(bits)


def expand_braid_to_bytes(
    braid: List[str],
    target_size: int,
    seed_context: str = "",
) -> bytes:
    """
    Expand folded_braid to target bytes using SHAKE256.

    The braid symbols provide the initial structure, then SHAKE256
    expands deterministically to the target size.

    Args:
        braid: List of symbols from folded_braid()
        target_size: Target byte count
        seed_context: Additional context (model name, sha256, etc.)

    Returns:
        Expanded bytes
    """
    # Decode to bits
    bit_string = decode_braid_to_bits(braid)

    # Convert bits to bytes (pad if needed)
    bit_bytes = len(bit_string) // 8
    if bit_bytes > 0:
        bit_data = int(bit_string[:bit_bytes * 8], 2).to_bytes(bit_bytes, 'big')
    else:
        bit_data = b''

    # Create seed from braid data + context
    seed = bit_data + seed_context.encode('utf-8')

    # Use SHAKE256 to expand to target size
    shake = hashlib.shake_256(seed)
    return shake.digest(target_size)


def expand_braid_to_floats(
    braid: List[str],
    num_elements: int,
    seed_context: str = "",
    dtype: str = "float32",
    scale: float = 0.1,
) -> np.ndarray:
    """
    Expand folded_braid to valid float array using SHAKE256.

    Uses uniform byte distribution to create Gaussian-like float values
    via Box-Muller-like transform.

    Args:
        braid: List of symbols from folded_braid()
        num_elements: Number of float elements to generate
        seed_context: Additional context (model name, sha256, etc.)
        dtype: "float32" or "float16"
        scale: Scale factor for output values (default 0.1)

    Returns:
        numpy array with valid float values
    """
    # Decode braid to bits
    bit_string = decode_braid_to_bits(braid)

    # Convert bits to bytes
    bit_bytes = len(bit_string) // 8
    if bit_bytes > 0:
        bit_data = int(bit_string[:bit_bytes * 8], 2).to_bytes(bit_bytes, 'big')
    else:
        bit_data = b''

    # Create seed from braid data + context
    seed = bit_data + seed_context.encode('utf-8')

    # Generate enough random bytes: 2 uint32 per float for Box-Muller
    shake = hashlib.shake_256(seed)
    raw_bytes = shake.digest(num_elements * 4)  # 4 bytes per uint32

    # Convert to uint32 array
    uint_arr = np.frombuffer(raw_bytes, dtype=np.uint32)

    # Convert to uniform [0, 1) using uint32 max
    uniform = uint_arr.astype(np.float64) / (2**32)

    # Clamp to avoid log(0)
    uniform = np.clip(uniform, 1e-10, 1.0 - 1e-10)

    # Use inverse normal CDF approximation for Gaussian-like distribution
    # Simpler: just scale uniform to [-scale, scale]
    values = (uniform * 2 - 1) * scale

    # Convert to target dtype
    np_dtype = np.float32 if dtype == "float32" else np.float16
    return values.astype(np_dtype)


def load_dna_seed(
    hxz_path: str,
    target_shape: Optional[Tuple[int, ...]] = None,
    dtype: str = "float32",
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
    """
    Load DNA seed and expand to state dict.

    Args:
        hxz_path: Path to .hxz file
        target_shape: Optional target shape (defaults to inferred)
        dtype: Target dtype ("float32" or "float16")

    Returns:
        Tuple of (state_dict, metadata)
    """
    _warn_once()
    info = parse_dna_seed(hxz_path)

    if not info.braid:
        raise ValueError(f"No folded_braid() found in {hxz_path}")

    # Create seed context from available info
    seed_context = f"{info.model_name}|{info.sha256}|{info.entropy}"

    # Calculate target size
    np_dtype = np.float32 if dtype == "float32" else np.float16
    element_size = np.dtype(np_dtype).itemsize

    if target_shape:
        num_elements = int(np.prod(target_shape))
    else:
        # Infer from braid length
        # Each symbol contributes structure, expand proportionally
        # For Mistral-7B shard: 104 symbols → ~2.36MB
        # That's ~22KB per symbol as a rough heuristic
        num_elements = len(info.braid) * 5632  # ~22KB / 4 bytes

    # Expand braid to valid float values (not raw bytes)
    arr = expand_braid_to_floats(info.braid, num_elements, seed_context, dtype=dtype)

    if target_shape:
        arr = arr.reshape(target_shape)

    # Build state dict with meaningful name
    tensor_name = info.model_name or "_dna_expanded"
    tensor_name = tensor_name.replace(" ", "_").replace("/", "_")

    state_dict = {tensor_name: arr}

    metadata = {
        "format": "dna_text_v1",
        "model_name": info.model_name,
        "timestamp": info.timestamp,
        "entropy": info.entropy,
        "sha256": info.sha256,
        "segments": info.segments,
        "braid_length": len(info.braid),
        "num_elements": arr.size,
        "expanded_bytes": arr.nbytes,
        "expanded_shape": list(arr.shape),
        "source_path": str(hxz_path),
        "expansion_method": "shake256_to_floats",
    }

    return state_dict, metadata


# =============================================================================
# UNIFIED FORMAT DETECTION + CANONICAL LOADER (WO-DNA-TEXTV1-BRIDGE-01)
# =============================================================================

from typing import Literal
import json as _json

DNA_FORMAT = Literal["text_v1", "hxz1", "cdna"]


def detect_dna_seed_format(path: Path) -> DNA_FORMAT:
    """
    Detect whether a seed is DNA text_v1, HXZ1 seedpack, or CDNA.

    Heuristics:
    - If file starts with b"#HXZ_START" -> text_v1
    - If file starts with b"HXZ1"       -> hxz1 seedpack
    - If file starts with b"HXZC" or b"CDNA" -> cdna format
    - Otherwise: raise ValueError
    """
    p = Path(path) if not isinstance(path, Path) else path
    with open(p, "rb") as f:
        prefix = f.read(16)

    if prefix.startswith(b"#HXZ_START"):
        return "text_v1"
    if prefix.startswith(b"HXZ1"):
        return "hxz1"
    if prefix.startswith(b"HXZC") or prefix.startswith(b"CDNA"):
        return "cdna"
    raise ValueError(f"Unknown DNA seed format for {path}: prefix={prefix[:8]!r}")


@dataclass
class DNASeedCanonical:
    """Unified canonical representation of any DNA seed format."""
    path: str
    format: DNA_FORMAT
    meta: Dict[str, Any]
    payload_bytes: bytes
    file_size: int


def load_hxz1_seed_header(path: Path) -> Tuple[Dict[str, Any], bytes]:
    """
    Load HXZ1 seedpack header and payload.

    HXZ1 format:
    - Magic: b"HXZ1" (4 bytes)
    - Header length: uint32 big-endian (4 bytes)
    - Header JSON: UTF-8
    - Payload: zlib-compressed tensor data
    """
    with open(path, "rb") as f:
        magic = f.read(4)
        if magic != b"HXZ1":
            raise ValueError(f"Expected HXZ1 magic, got {magic!r}")

        header_len_bytes = f.read(4)
        header_len = int.from_bytes(header_len_bytes, "big")
        header_json = f.read(header_len)
        header = _json.loads(header_json.decode("utf-8"))
        payload = f.read()

    return header, payload


def load_dna_seed_any(path: str) -> DNASeedCanonical:
    """
    Unified DNA seed loader: auto-detects format and returns canonical representation.

    Supports:
    - text_v1: DNA braid format with #HXZ_START header
    - hxz1: Binary seedpack with HXZ1 magic
    - cdna: CDNA compressed format with CDNA magic

    Args:
        path: Path to .hxz file

    Returns:
        DNASeedCanonical with format, meta, and payload
    """
    _warn_once()
    p = Path(path) if not isinstance(path, Path) else path
    file_size = p.stat().st_size
    fmt = detect_dna_seed_format(p)

    if fmt == "text_v1":
        # Parse text_v1 format
        info = parse_dna_seed(str(p))

        # Get raw payload (the zlib blob after #====)
        with open(p, "rb") as f:
            data = f.read()
        sep_idx = data.find(b"#====")
        if sep_idx != -1:
            blob_start = sep_idx + 5
            if blob_start < len(data) and data[blob_start:blob_start+1] == b'\n':
                blob_start += 1
            payload = data[blob_start:]
        else:
            payload = b''

        meta = {
            "schema": "dna.text_v1",
            "model_name": info.model_name,
            "timestamp": info.timestamp,
            "entropy": info.entropy,
            "sha256_header": info.sha256,
            "segments": info.segments,
            "braid_length": len(info.braid),
            "braid_preview": info.braid[:10] if info.braid else [],
        }

        return DNASeedCanonical(
            path=str(p),
            format=fmt,
            meta=meta,
            payload_bytes=payload,
            file_size=file_size,
        )

    elif fmt == "hxz1":
        # Parse HXZ1 seedpack format
        header, payload = load_hxz1_seed_header(p)

        meta = {
            "schema": header.get("schema", "seedpack.v2.0"),
            **{k: v for k, v in header.items() if k != "schema"},
        }

        return DNASeedCanonical(
            path=str(p),
            format=fmt,
            meta=meta,
            payload_bytes=payload,
            file_size=file_size,
        )

    elif fmt == "cdna":
        # Parse CDNA format (HXZC magic) - use existing cdna_stream if available
        try:
            from helix_cdc.regrow.cdna_stream import inspect_cdna_hxz_header
            info = inspect_cdna_hxz_header(str(p))
            if not info.get("valid"):
                raise ValueError(info.get("error", "Invalid CDNA file"))

            meta = {
                "schema": "cdna.hxzc.v1",
                "tensor_name": info.get("tensor_name"),
                "shape": info.get("shape"),
                "rows": info.get("rows"),
                "cols": info.get("cols"),
                "codebook_size": info.get("codebook_size"),
                "indices_dtype": info.get("indices_dtype"),
                "codebook_dtype": info.get("codebook_dtype"),
                "encoding": info.get("encoding"),
            }

            # Read payload (everything after header)
            with open(p, "rb") as f:
                data = f.read()
            # HXZC: 4-byte magic + 2-byte version + 2-byte header_len + header + blob
            header_len = int.from_bytes(data[6:8], "big")
            payload = data[8 + header_len:]

            return DNASeedCanonical(
                path=str(p),
                format=fmt,
                meta=meta,
                payload_bytes=payload,
                file_size=file_size,
            )
        except ImportError:
            # Fallback: minimal HXZC parsing
            with open(p, "rb") as f:
                data = f.read()

            if len(data) < 8:
                raise ValueError(f"CDNA file too small: {len(data)} bytes")

            # HXZC: 4-byte magic + 2-byte version + 2-byte header_len
            header_len = int.from_bytes(data[6:8], "big")
            if len(data) < 8 + header_len:
                raise ValueError(f"CDNA header truncated")

            header_json = data[8:8+header_len]
            header = _json.loads(header_json.decode("utf-8"))
            payload = data[8+header_len:]

            meta = {
                "schema": header.get("schema", "cdna.hxzc.v1"),
                **{k: v for k, v in header.items() if k != "schema"},
            }

            return DNASeedCanonical(
                path=str(p),
                format=fmt,
                meta=meta,
                payload_bytes=payload,
                file_size=file_size,
            )

    else:
        raise ValueError(f"Unsupported format: {fmt}")


# CLI for testing
if __name__ == "__main__":
    import sys
    import json

    if len(sys.argv) < 2:
        print("Usage: python -m helix_cdc.regrow.dna_seed_loader <seed.hxz>")
        sys.exit(1)

    hxz_path = sys.argv[1]

    print(f"=== DNA Seed Loader ===")
    print(f"File: {hxz_path}")
    print()

    # Parse
    info = parse_dna_seed(hxz_path)
    print(f"Model: {info.model_name}")
    print(f"Timestamp: {info.timestamp}")
    print(f"Entropy: {info.entropy} bits/byte")
    print(f"SHA256: {info.sha256[:32]}..." if info.sha256 else "SHA256: (none)")
    print(f"Segments: {len(info.segments)}")
    for seg in info.segments:
        print(f"  - {seg['id']}: X={seg.get('X')}, Y={seg.get('Y')}, Z={seg.get('Z')}")
    print(f"Braid length: {len(info.braid)} symbols")
    print(f"Braid preview: {','.join(info.braid[:10])}...")
    print()

    # Decode braid to bits
    bits = decode_braid_to_bits(info.braid)
    print(f"Decoded bits: {len(bits)} bits ({len(bits)//8} bytes)")
    print(f"Bit preview: {bits[:64]}...")
    print()

    # Load and expand
    print("Expanding to state dict...")
    state_dict, metadata = load_dna_seed(hxz_path)
    print(f"Tensors loaded: {len(state_dict)}")
    for name, arr in state_dict.items():
        print(f"  - {name}: {arr.shape} ({arr.dtype})")
    print()
    print(f"Metadata: {json.dumps(metadata, indent=2, default=str)}")
