#!/usr/bin/env python3
"""
Crystal Pack Format (.cryst)
============================

Binary file format for storing crystal-encoded neural network weights.

File Structure:
---------------
[HEADER]           Fixed 64 bytes
  - Magic:         4 bytes "CRYS"
  - Version:       4 bytes (uint32, little-endian)
  - Model hash:    32 bytes (SHA256 of original model)
  - Num tensors:   4 bytes (uint32)
  - Num bases:     4 bytes (uint32)
  - Adapter mode:  4 bytes (0=none, 1=int8, 2=int4, 3=sparse)
  - Seed hash:     8 bytes (first 8 bytes of basis seed hash)
  - Reserved:      4 bytes

[BASIS PARAMS]     Variable (JSON + padding to 8-byte alignment)
  - Length:        4 bytes (uint32)
  - JSON data:     Variable (basis generation parameters)
  - Padding:       0-7 bytes

[TENSOR INDEX]     8 + 24 * num_tensors bytes
  - Total size:    8 bytes (uint64, total data section size)
  - Per tensor:
    - Name hash:   8 bytes (first 8 bytes of tensor name SHA256)
    - Offset:      8 bytes (uint64, offset into data section)
    - Size:        8 bytes (uint64, size of tensor block)

[DATA SECTION]     Variable
  - Per tensor block:
    [TENSOR BLOCK]
    - Name length: 2 bytes (uint16)
    - Name:        Variable (UTF-8)
    - Shape dims:  2 bytes (uint16, number of dimensions)
    - Shape:       4 * dims bytes (uint32 per dimension)
    - Coeffs size: 4 bytes (uint32, num coefficients)
    - Coeffs:      4 * coeffs_size bytes (float32)
    - Adapter len: 4 bytes (uint32)
    - Adapter:     Variable (adapter data)
    - Checksum:    4 bytes (CRC32 of block)

Usage:
------
    from helix_cdc.crystal_pack import CrystalPack

    # Write
    pack = CrystalPack.create(model_hash="...", num_bases=64, adapter_mode="int8")
    pack.add_tensor("layer.0.weight", encoding)
    pack.write("model.cryst")

    # Read
    pack = CrystalPack.load("model.cryst")
    encoding = pack.get_tensor("layer.0.weight")
    decoded = decode_crystal_to_tensor(encoding)
"""

import struct
import json
import hashlib
import zlib
import numpy as np
from pathlib import Path
from typing import Dict, List, Tuple, Optional, Any
from dataclasses import dataclass

# Constants
MAGIC = b'CRYS'
VERSION = 1
HEADER_SIZE = 64

ADAPTER_MODES = {
    "none": 0,
    "int8": 1,
    "int4": 2,
    "sparse": 3,
}
ADAPTER_MODES_REV = {v: k for k, v in ADAPTER_MODES.items()}


@dataclass
class TensorEntry:
    """Metadata for a tensor in the pack."""
    name: str
    name_hash: bytes
    offset: int
    size: int
    shape: Tuple[int, ...]
    coeffs: np.ndarray
    adapter_data: Dict[str, Any]
    checksum: int


class CrystalPack:
    """Crystal-encoded model pack file."""

    def __init__(self):
        self.model_hash: bytes = b'\x00' * 32
        self.num_bases: int = 64
        self.adapter_mode: str = "int8"
        self.seed: str = "crystal"
        self.basis_params: Dict[str, Any] = {}
        self.tensors: Dict[str, TensorEntry] = {}
        self._data_buffer: bytes = b''

    @classmethod
    def create(
        cls,
        model_hash: str,
        num_bases: int = 64,
        adapter_mode: str = "int8",
        seed: str = "crystal"
    ) -> 'CrystalPack':
        """Create a new crystal pack."""
        pack = cls()
        pack.model_hash = bytes.fromhex(model_hash) if len(model_hash) == 64 else hashlib.sha256(model_hash.encode()).digest()
        pack.num_bases = num_bases
        pack.adapter_mode = adapter_mode
        pack.seed = seed
        pack.basis_params = {
            "num_bases": num_bases,
            "seed": seed,
            "version": VERSION,
        }
        return pack

    def add_tensor(self, name: str, encoding: 'CrystalEncoding') -> None:
        """Add a tensor encoding to the pack."""
        # Serialize tensor block
        block = self._serialize_tensor_block(name, encoding)

        name_hash = hashlib.sha256(name.encode()).digest()[:8]
        offset = len(self._data_buffer)

        entry = TensorEntry(
            name=name,
            name_hash=name_hash,
            offset=offset,
            size=len(block),
            shape=encoding.shape,
            coeffs=encoding.crystal_coeffs,
            adapter_data=encoding.adapter_data,
            checksum=zlib.crc32(block) & 0xffffffff,
        )

        self.tensors[name] = entry
        self._data_buffer += block

    def _serialize_tensor_block(self, name: str, encoding: 'CrystalEncoding') -> bytes:
        """Serialize a single tensor block."""
        parts = []

        # Name
        name_bytes = name.encode('utf-8')
        parts.append(struct.pack('<H', len(name_bytes)))
        parts.append(name_bytes)

        # Shape
        parts.append(struct.pack('<H', len(encoding.shape)))
        for dim in encoding.shape:
            parts.append(struct.pack('<I', dim))

        # Coefficients
        coeffs = encoding.crystal_coeffs.astype(np.float32)
        parts.append(struct.pack('<I', len(coeffs)))
        parts.append(coeffs.tobytes())

        # Adapter data (JSON + binary if needed)
        adapter_json = json.dumps(self._prepare_adapter_for_json(encoding.adapter_data))
        adapter_bytes = adapter_json.encode('utf-8')
        parts.append(struct.pack('<I', len(adapter_bytes)))
        parts.append(adapter_bytes)

        # Combine and add checksum
        block = b''.join(parts)
        checksum = zlib.crc32(block) & 0xffffffff
        block += struct.pack('<I', checksum)

        return block

    def _prepare_adapter_for_json(self, adapter_data: Dict[str, Any]) -> Dict[str, Any]:
        """Convert numpy arrays to lists for JSON serialization."""
        result = {}
        for key, value in adapter_data.items():
            if isinstance(value, np.ndarray):
                result[key] = value.tolist()
            else:
                result[key] = value
        return result

    def write(self, path: str) -> int:
        """Write pack to file. Returns file size."""
        path = Path(path)

        with open(path, 'wb') as f:
            # Write header
            header = self._build_header()
            f.write(header)

            # Write basis params
            basis_json = json.dumps(self.basis_params).encode('utf-8')
            padding = (8 - (len(basis_json) % 8)) % 8
            f.write(struct.pack('<I', len(basis_json)))
            f.write(basis_json)
            f.write(b'\x00' * padding)

            # Write tensor index
            f.write(struct.pack('<Q', len(self._data_buffer)))
            for entry in self.tensors.values():
                f.write(entry.name_hash)
                f.write(struct.pack('<Q', entry.offset))
                f.write(struct.pack('<Q', entry.size))

            # Write data section
            f.write(self._data_buffer)

        return path.stat().st_size

    def _build_header(self) -> bytes:
        """Build 64-byte header."""
        header = bytearray(HEADER_SIZE)

        # Magic
        header[0:4] = MAGIC

        # Version
        struct.pack_into('<I', header, 4, VERSION)

        # Model hash
        header[8:40] = self.model_hash

        # Num tensors
        struct.pack_into('<I', header, 40, len(self.tensors))

        # Num bases
        struct.pack_into('<I', header, 44, self.num_bases)

        # Adapter mode
        struct.pack_into('<I', header, 48, ADAPTER_MODES.get(self.adapter_mode, 0))

        # Seed hash (first 8 bytes of SHA256)
        seed_hash = hashlib.sha256(self.seed.encode()).digest()[:8]
        header[52:60] = seed_hash

        # Reserved (already zero)

        return bytes(header)

    @classmethod
    def load(cls, path: str) -> 'CrystalPack':
        """Load pack from file."""
        path = Path(path)
        pack = cls()

        with open(path, 'rb') as f:
            # Read header
            header = f.read(HEADER_SIZE)
            pack._parse_header(header)

            # Read basis params
            basis_len = struct.unpack('<I', f.read(4))[0]
            basis_json = f.read(basis_len)
            padding = (8 - (basis_len % 8)) % 8
            f.read(padding)
            pack.basis_params = json.loads(basis_json.decode('utf-8'))

            # Read tensor index
            total_data_size = struct.unpack('<Q', f.read(8))[0]
            num_tensors = len(pack.tensors)  # Set by header

            # Index position
            index_start = f.tell()

            # Read index entries
            tensor_index = []
            for _ in range(pack._num_tensors_from_header):
                name_hash = f.read(8)
                offset = struct.unpack('<Q', f.read(8))[0]
                size = struct.unpack('<Q', f.read(8))[0]
                tensor_index.append((name_hash, offset, size))

            # Data section start
            data_start = f.tell()

            # Read tensor blocks
            for name_hash, offset, size in tensor_index:
                f.seek(data_start + offset)
                block = f.read(size)
                entry = pack._parse_tensor_block(block, name_hash, offset)
                pack.tensors[entry.name] = entry

        return pack

    def _parse_header(self, header: bytes) -> None:
        """Parse 64-byte header."""
        if header[0:4] != MAGIC:
            raise ValueError(f"Invalid magic: {header[0:4]}")

        version = struct.unpack_from('<I', header, 4)[0]
        if version > VERSION:
            raise ValueError(f"Unsupported version: {version}")

        self.model_hash = header[8:40]
        self._num_tensors_from_header = struct.unpack_from('<I', header, 40)[0]
        self.num_bases = struct.unpack_from('<I', header, 44)[0]
        self.adapter_mode = ADAPTER_MODES_REV.get(struct.unpack_from('<I', header, 48)[0], "none")

    def _parse_tensor_block(self, block: bytes, name_hash: bytes, offset: int) -> TensorEntry:
        """Parse a tensor block."""
        pos = 0

        # Name
        name_len = struct.unpack_from('<H', block, pos)[0]
        pos += 2
        name = block[pos:pos + name_len].decode('utf-8')
        pos += name_len

        # Shape
        num_dims = struct.unpack_from('<H', block, pos)[0]
        pos += 2
        shape = []
        for _ in range(num_dims):
            shape.append(struct.unpack_from('<I', block, pos)[0])
            pos += 4
        shape = tuple(shape)

        # Coefficients
        num_coeffs = struct.unpack_from('<I', block, pos)[0]
        pos += 4
        coeffs = np.frombuffer(block[pos:pos + num_coeffs * 4], dtype=np.float32).copy()
        pos += num_coeffs * 4

        # Adapter data
        adapter_len = struct.unpack_from('<I', block, pos)[0]
        pos += 4
        adapter_json = block[pos:pos + adapter_len].decode('utf-8')
        adapter_data = json.loads(adapter_json)
        pos += adapter_len

        # Convert lists back to arrays
        adapter_data = self._restore_adapter_arrays(adapter_data)

        # Checksum
        stored_checksum = struct.unpack_from('<I', block, pos)[0]
        computed_checksum = zlib.crc32(block[:pos]) & 0xffffffff

        if stored_checksum != computed_checksum:
            raise ValueError(f"Checksum mismatch for tensor {name}")

        return TensorEntry(
            name=name,
            name_hash=name_hash,
            offset=offset,
            size=len(block),
            shape=shape,
            coeffs=coeffs,
            adapter_data=adapter_data,
            checksum=stored_checksum,
        )

    def _restore_adapter_arrays(self, adapter_data: Dict[str, Any]) -> Dict[str, Any]:
        """Convert lists back to numpy arrays."""
        result = {}
        for key, value in adapter_data.items():
            if isinstance(value, list):
                result[key] = np.array(value)
            else:
                result[key] = value
        return result

    def get_tensor(self, name: str) -> Optional[TensorEntry]:
        """Get a tensor entry by name."""
        return self.tensors.get(name)

    def get_encoding(self, name: str) -> Optional['CrystalEncoding']:
        """Reconstruct a CrystalEncoding from a tensor entry."""
        entry = self.tensors.get(name)
        if entry is None:
            return None

        # Import here to avoid circular import
        from helix_cdc.crystal_codec import CrystalEncoding

        return CrystalEncoding(
            shape=entry.shape,
            dtype="float32",
            num_bases=self.num_bases,
            seed=self.basis_params.get("seed", "crystal"),
            crystal_coeffs=entry.coeffs,
            adapter_mode=self.adapter_mode,
            adapter_data=entry.adapter_data,
            checksum="",  # Will need to recompute if needed
        )

    def list_tensors(self) -> List[str]:
        """List all tensor names."""
        return list(self.tensors.keys())

    def summary(self) -> Dict[str, Any]:
        """Get pack summary."""
        total_coeffs = sum(len(e.coeffs) for e in self.tensors.values())
        total_data = sum(e.size for e in self.tensors.values())

        return {
            "model_hash": self.model_hash.hex(),
            "num_tensors": len(self.tensors),
            "num_bases": self.num_bases,
            "adapter_mode": self.adapter_mode,
            "total_coefficients": total_coeffs,
            "total_data_bytes": total_data,
            "basis_seed": self.basis_params.get("seed", "crystal"),
        }


def pack_model_gguf(
    gguf_path: str,
    output_path: str,
    num_bases: int = 64,
    adapter_mode: str = "int8",
    max_tensors: Optional[int] = None,
) -> Dict[str, Any]:
    """
    Pack a GGUF model to crystal format.

    Args:
        gguf_path: Path to input GGUF file
        output_path: Path to output .cryst file
        num_bases: Number of FibPi3D bases
        adapter_mode: Adapter quantization mode
        max_tensors: Limit number of tensors (for testing)

    Returns:
        Dictionary with packing statistics
    """
    from helix_cdc.gguf.scanner import GGUFScanner
    from helix_cdc.crystal_codec import encode_tensor_to_crystal, init_fibpi3d_basis

    # Scan GGUF
    scanner = GGUFScanner(gguf_path).scan()

    # Create pack
    model_hash = hashlib.sha256(Path(gguf_path).read_bytes()[:1024]).hexdigest()
    pack = CrystalPack.create(model_hash, num_bases, adapter_mode)

    # Encode tensors
    stats = {
        "tensors_processed": 0,
        "tensors_skipped": 0,
        "total_elements": 0,
        "cosine_sims": [],
    }

    tensors = scanner.tensors
    if max_tensors:
        tensors = tensors[:max_tensors]

    for tensor_info in tensors:
        tensor = scanner.extract_tensor(tensor_info)
        if tensor is None:
            stats["tensors_skipped"] += 1
            continue

        # Pre-warm basis
        flat = tensor.flatten().astype(np.float32)
        init_fibpi3d_basis(num_bases, [len(flat)])

        # Encode
        encoded = encode_tensor_to_crystal(flat, num_bases=num_bases, adapter_mode=adapter_mode)

        # Add to pack (reshape encoding to match original tensor shape)
        # Note: We store flat but record original shape
        encoded.shape = tensor_info['dims'] if 'dims' in tensor_info else tensor.shape
        pack.add_tensor(tensor_info['name'], encoded)

        stats["tensors_processed"] += 1
        stats["total_elements"] += len(flat)
        stats["cosine_sims"].append(encoded.cosine_similarity)

    # Write pack
    file_size = pack.write(output_path)

    stats["output_file"] = output_path
    stats["file_size_bytes"] = file_size
    stats["avg_cosine_sim"] = float(np.mean(stats["cosine_sims"])) if stats["cosine_sims"] else 0
    stats["min_cosine_sim"] = float(np.min(stats["cosine_sims"])) if stats["cosine_sims"] else 0

    return stats
