#!/usr/bin/env python3
"""
HXZ Model Loader for helix64 v0.3.0

Loads Helix capsule (.hxz) seeds and regenerates them to PyTorch state dicts.

Supports two HXZ formats:
  1. Text format (#HXZ_START header) - DNA seed with entropy + sha256
  2. Binary format (HXZ1 magic) - seedpack.v1.x with tensor registry

Usage:
    from helix_cdc.regrow.hxz_model_loader import load_hxz_to_state_dict

    state_dict, metadata = load_hxz_to_state_dict("model.hxz")
    # state_dict is Dict[str, torch.Tensor]
    # metadata is dict with loader info

This module owns "how do I turn a capsule into model weights?"
ModelRunner just calls it.
"""

from __future__ import annotations

import hashlib
import json
import struct
import tempfile
import zlib
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Optional, Tuple

import numpy as np


@dataclass
class HXZHeader:
    """Parsed HXZ header information."""
    format: str  # "text_v1" or "binary_v1"
    schema: str
    entropy: Optional[float] = None
    sha256: Optional[str] = None
    created_utc: Optional[str] = None
    encoding: str = "raw"
    registry: list = None  # Tensor registry for binary format
    header_bytes: int = 0
    blob_offset: int = 0

    def __post_init__(self):
        if self.registry is None:
            self.registry = []


def inspect_hxz_header(hxz_path: str) -> HXZHeader:
    """
    Inspect HXZ file and parse header.

    Args:
        hxz_path: Path to .hxz file

    Returns:
        HXZHeader with parsed information
    """
    path = Path(hxz_path)
    if not path.exists():
        raise FileNotFoundError(f"HXZ not found: {hxz_path}")

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

        # Binary format: HXZ1
        if magic == b"HXZ1":
            return _parse_binary_hxz(f, path)

        # Text format: #HXZ_START
        f.seek(0)
        line = f.read(10)
        if line.startswith(b"#HXZ_START"):
            f.seek(0)
            return _parse_text_hxz(f, path)

        # JSON stub format
        f.seek(0)
        try:
            content = f.read().decode("utf-8")
            data = json.loads(content)
            if data.get("stub"):
                return HXZHeader(
                    format="stub",
                    schema="stub.v1",
                    sha256=data.get("source"),
                )
            # Regular JSON seed
            return HXZHeader(
                format="json_v1",
                schema=data.get("schema", "unknown"),
                sha256=data.get("sha256"),
            )
        except (json.JSONDecodeError, UnicodeDecodeError):
            pass

        raise ValueError(f"Unknown HXZ format: magic={magic!r}")


def _parse_binary_hxz(f, path: Path) -> HXZHeader:
    """Parse binary HXZ format (HXZ1 magic)."""
    # Format: HXZ1 (4) + version (2) + json_len (2 BE) + JSON
    f.seek(4)  # After HXZ1 magic

    # Read version (2 bytes) and JSON length (2 bytes big-endian)
    version_bytes = f.read(2)
    len_bytes = f.read(2)
    if len(len_bytes) < 2:
        raise ValueError("Truncated HXZ header")

    # Parse header length (2 bytes big-endian)
    header_len = (len_bytes[0] << 8) | len_bytes[1]

    # Read JSON header
    json_bytes = f.read(header_len)
    try:
        header_data = json.loads(json_bytes.decode("utf-8"))
    except (json.JSONDecodeError, UnicodeDecodeError) as e:
        raise ValueError(f"Invalid HXZ JSON header: {e}")

    blob_offset = 4 + 2 + 2 + header_len  # magic + version + len + json

    return HXZHeader(
        format="binary_v1",
        schema=header_data.get("schema", "seedpack.v1"),
        created_utc=header_data.get("created_utc"),
        encoding=header_data.get("encoding", "raw"),
        registry=header_data.get("registry", []),
        header_bytes=header_len,
        blob_offset=blob_offset,
    )


def _parse_text_hxz(f, path: Path) -> HXZHeader:
    """Parse text HXZ format (#HXZ_START header)."""
    lines = []
    entropy = None
    sha256 = None

    for line in f:
        line_str = line.decode("utf-8", errors="replace")
        if line_str.startswith("#===="):
            break
        lines.append(line_str)

        if line_str.startswith("# entropy:"):
            try:
                entropy = float(line_str.split(":")[-1].strip())
            except ValueError:
                pass
        elif line_str.startswith("# sha256:"):
            sha256 = line_str.split(":")[-1].strip()

    blob_offset = f.tell()

    return HXZHeader(
        format="text_v1",
        schema="hxz_text.v1",
        entropy=entropy,
        sha256=sha256,
        blob_offset=blob_offset,
    )


def load_hxz_to_state_dict(
    hxz_path: str,
    dtype: str = "float32",
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
    """
    Load HXZ capsule and regenerate to PyTorch state dict.

    Args:
        hxz_path: Path to .hxz file
        dtype: Target dtype ("float32" or "float16")

    Returns:
        Tuple of (state_dict, metadata)
        - state_dict: Dict mapping tensor names to numpy arrays (torch-compatible)
        - metadata: Dict with loader info (format, registry, etc.)
    """
    header = inspect_hxz_header(hxz_path)

    if header.format == "stub":
        raise ValueError(
            f"HXZ is a stub pointing to: {header.sha256}. "
            "Real weights not embedded."
        )

    if header.format == "binary_v1":
        return _load_binary_hxz(hxz_path, header, dtype)

    if header.format == "text_v1":
        return _load_text_hxz(hxz_path, header, dtype)

    if header.format == "json_v1":
        raise ValueError(
            f"HXZ is a JSON descriptor (schema={header.schema}). "
            "Real weights not embedded."
        )

    raise ValueError(f"Cannot load HXZ format: {header.format}")


def _load_binary_hxz(
    hxz_path: str,
    header: HXZHeader,
    dtype: str,
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
    """Load binary HXZ format (seedpack.v1.x)."""
    path = Path(hxz_path)
    np_dtype = np.float32 if dtype == "float32" else np.float16

    with open(path, "rb") as f:
        f.seek(header.blob_offset)
        blob = f.read()

    # Decompress if needed
    if header.encoding == "zlib":
        try:
            blob = zlib.decompress(blob)
        except zlib.error as e:
            raise ValueError(f"Failed to decompress HXZ blob: {e}")

    # Reconstruct tensors from registry
    state_dict = {}
    for entry in header.registry:
        name = entry["name"]
        entry_dtype = entry.get("dtype", "float32")
        shape = tuple(entry.get("shape", []))
        offset = entry.get("offset", 0)
        length = entry.get("length", 0)

        # Extract bytes
        tensor_bytes = blob[offset:offset + length]
        if len(tensor_bytes) < length:
            print(f"[HXZ] Warning: Truncated tensor {name}: got {len(tensor_bytes)}, expected {length}")
            # Pad with zeros
            tensor_bytes = tensor_bytes + b'\x00' * (length - len(tensor_bytes))

        # Convert to numpy - handle various dtypes
        dtype_map = {
            "float32": np.float32,
            "float16": np.float16,
            "float64": np.float64,
            "int8": np.int8,
            "int16": np.int16,
            "int32": np.int32,
            "int64": np.int64,
            "uint8": np.uint8,
            "uint16": np.uint16,
            "uint32": np.uint32,
            "uint64": np.uint64,
        }
        source_dtype = dtype_map.get(entry_dtype, np.float32)
        arr = np.frombuffer(tensor_bytes, dtype=source_dtype)

        if shape:
            try:
                arr = arr.reshape(shape)
            except ValueError as e:
                print(f"[HXZ] Warning: Cannot reshape {name} to {shape}: {e}")

        # Convert to target dtype
        if arr.dtype != np_dtype:
            arr = arr.astype(np_dtype)

        state_dict[name] = arr

    metadata = {
        "format": header.format,
        "schema": header.schema,
        "encoding": header.encoding,
        "tensors_loaded": len(state_dict),
        "registry_entries": len(header.registry),
        "created_utc": header.created_utc,
        "blob_size": len(blob),
        "source_path": str(path),
    }

    return state_dict, metadata


def _load_text_hxz(
    hxz_path: str,
    header: HXZHeader,
    dtype: str,
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
    """
    Load text HXZ format (DNA seed with compressed blob).

    CORRECT ARCHITECTURE:
    - One .hxz seed = ONE shard (~2.3 MB tensor)
    - Full model = ~129 seeds (62 KB total from shard manifest)
    - NOT: one seed → entire 7B model

    Returns the single shard tensor from the DNA seed.
    """
    path = Path(hxz_path)
    np_dtype = np.float32 if dtype == "float32" else np.float16

    # Try DNA seed loader first (handles folded_braid format)
    try:
        from helix_cdc.regrow.dna_seed_loader import load_dna_seed, parse_dna_seed

        # Load DNA seed - returns ONE shard tensor (~2.3 MB)
        state_dict, dna_metadata = load_dna_seed(hxz_path, dtype=dtype)

        # Merge header info
        dna_metadata.update({
            "hxz_header_entropy": header.entropy,
            "hxz_header_sha256": header.sha256,
            "architecture_note": "One seed = one shard (~2.3MB). Full model needs ~129 seeds.",
        })
        return state_dict, dna_metadata
    except Exception as e:
        # DNA loader failed, try other formats
        pass

    with open(path, "rb") as f:
        f.seek(header.blob_offset)
        blob = f.read()

    # Try to decompress (usually zlib)
    try:
        decompressed = zlib.decompress(blob)
    except zlib.error:
        decompressed = blob  # Raw if not zlib

    # Try to parse as JSON (BRAID format with ::ANALYSIS:: and ::FOLDING_MAP::)
    try:
        braid_data = json.loads(decompressed.decode("utf-8"))
        # It's a braid - we need the braid_bridge to regrow
        return _regrow_from_braid(braid_data, header, dtype)
    except (json.JSONDecodeError, UnicodeDecodeError):
        pass

    # Otherwise treat as raw tensor data
    # Without registry, we can't know the shapes
    # Return as single flat tensor
    arr = np.frombuffer(decompressed, dtype=np_dtype)

    state_dict = {"_raw_blob": arr}
    metadata = {
        "format": header.format,
        "schema": header.schema,
        "entropy": header.entropy,
        "sha256": header.sha256,
        "blob_size": len(decompressed),
        "source_path": str(path),
        "note": "Text HXZ without registry - returned as raw blob",
    }

    return state_dict, metadata


def _regrow_from_braid(
    braid_data: Dict[str, Any],
    header: HXZHeader,
    dtype: str,
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
    """
    Regrow tensors from BRAID folding map data.

    Uses helix_cdc.regrow.braid_bridge for actual regeneration.
    """
    np_dtype = np.float32 if dtype == "float32" else np.float16

    # Check if it has BRAID structure
    if "::ANALYSIS::" in braid_data and "::FOLDING_MAP::" in braid_data:
        try:
            from helix_cdc.regrow.braid_bridge import braid_regen_callback
            import base64

            analysis = braid_data["::ANALYSIS::"]
            folding_map = braid_data["::FOLDING_MAP::"]

            # Get indices
            indices = np.frombuffer(
                base64.b64decode(folding_map["indices_b64"]),
                dtype=np.uint8
            )

            # Get codebook
            codebook = np.array(analysis["codebook"], dtype=np.float32)

            # Regrow: indices -> codebook lookup
            values = codebook[indices]

            # Apply residuals if present
            sig_residuals = folding_map.get("significant_residuals", {})
            for pos_str, residual in sig_residuals.items():
                pos = int(pos_str)
                if 0 <= pos < len(values):
                    values[pos] += residual

            if values.dtype != np_dtype:
                values = values.astype(np_dtype)

            # Get original shape if stored
            orig_shape = braid_data.get("original_shape")
            if orig_shape:
                values = values.reshape(orig_shape)

            tensor_name = braid_data.get("tensor_name", "_regrown")
            state_dict = {tensor_name: values}

            metadata = {
                "format": "braid_v1",
                "schema": header.schema,
                "entropy": header.entropy,
                "codebook_size": len(codebook),
                "indices_count": len(indices),
                "residuals_count": len(sig_residuals),
                "regrown_shape": list(values.shape),
            }

            return state_dict, metadata

        except ImportError as e:
            pass  # Fall through to simple approach

    # Simple case: braid_data might have tensors directly
    state_dict = {}
    for key, value in braid_data.items():
        if isinstance(value, list):
            arr = np.array(value, dtype=np_dtype)
            state_dict[key] = arr

    if not state_dict:
        # Return raw dict as metadata
        return {}, {
            "format": "braid_json",
            "raw_keys": list(braid_data.keys()),
            "note": "BRAID data present but regrow pipeline not fully wired",
        }

    return state_dict, {
        "format": "braid_simple",
        "tensors_loaded": len(state_dict),
    }


def hxz_to_pytorch_state_dict(hxz_path: str) -> Tuple[Dict[str, Any], Dict[str, Any]]:
    """
    High-level function to convert HXZ to PyTorch-compatible state dict.

    This is the main entry point for ModelRunner.

    Args:
        hxz_path: Path to .hxz file

    Returns:
        Tuple of (state_dict, mapping_info)
        - state_dict: Dict[str, torch.Tensor] (actually numpy, but torch-compatible)
        - mapping_info: Dict with {"mapped": dict, "unmapped": list}
    """
    state_dict, metadata = load_hxz_to_state_dict(hxz_path)

    # Build mapping info similar to GGUF loader
    mapping_info = {
        "mapped": {name: name for name in state_dict.keys()},
        "unmapped": [],
        "hxz_metadata": metadata,
    }

    return state_dict, mapping_info


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

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

    hxz_path = sys.argv[1]

    print(f"=== HXZ Model Loader ===")
    print(f"Loading: {hxz_path}")
    print()

    # Inspect header
    header = inspect_hxz_header(hxz_path)
    print(f"Format: {header.format}")
    print(f"Schema: {header.schema}")
    if header.entropy:
        print(f"Entropy: {header.entropy}")
    if header.sha256:
        print(f"SHA256: {header.sha256[:32]}...")
    if header.registry:
        print(f"Registry entries: {len(header.registry)}")
        for entry in header.registry[:5]:
            print(f"  - {entry['name']}: {entry.get('shape', '?')} ({entry.get('dtype', '?')})")
        if len(header.registry) > 5:
            print(f"  ... and {len(header.registry) - 5} more")
    print()

    # Try to load
    try:
        state_dict, metadata = load_hxz_to_state_dict(hxz_path)
        print(f"Loaded {len(state_dict)} tensors:")
        for name, arr in list(state_dict.items())[:5]:
            print(f"  - {name}: {arr.shape} ({arr.dtype})")
        if len(state_dict) > 5:
            print(f"  ... and {len(state_dict) - 5} more")
        print()
        print(f"Metadata: {json.dumps(metadata, indent=2, default=str)}")
    except Exception as e:
        print(f"Failed to load: {e}")
        sys.exit(1)
