#!/usr/bin/env python3
"""
FlowTorch Bridge to Proven Q4_K_M Centroid Regen
CC2-FIX: Single source of truth for centroid regen

This module bridges FlowTorch's tensor_from_seed() to the proven
regrow_q4km.py implementation. There must be NO second RNG path.

FlowTorch -> flowtorch_bridge -> regrow_q4km (proven, receipted)
"""
import hashlib
import struct
from dataclasses import dataclass
from typing import Tuple, Optional, Dict, Any
from pathlib import Path
from io import BytesIO
import numpy as np

# Import the PROVEN regen implementation
from helix_cdc.regrow_q4km import (
    load_centroid_atlas,
    combine_seeds,
    deterministic_rng_context,
    synthesize_q4km_block,
    grow_q4km_weights,
    grow_q4km_weights_parallel,
)


@dataclass
class TensorSpec:
    """
    Specification for a tensor to regenerate.

    CC2-OFFSET: Understanding offset semantics:
    - The `offset` parameter is for FibPi3D parallel slicing WITHIN a tensor
    - For per-tensor regen, unique `name` creates unique RNG context
    - RNG context = sha256(combined_seed:name:offset)
    - So (same seed, different name, offset=0) → different deterministic RNG
    - offset > 0 is only needed when parallelizing a SINGLE tensor's regen

    The `size_bytes` is computed from shape/dtype if not provided:
    - Q4_K_M: ceil(numel / 32) * 24 bytes
    """
    name: str                    # e.g., "blk.00.attn.q" - MUST BE UNIQUE per tensor!
    shape: Tuple[int, ...]       # e.g., (4096, 4096)
    dtype: str = "float16"       # "float16", "float32"
    offset: int = 0              # FibPi3D fold offset (0 for single-tensor regen)
    size_bytes: Optional[int] = None  # Override byte size (computed from shape if None)


def regen_tensor_q4km_from_seed(
    seed_hex: str,
    spec: TensorSpec,
    atlas_path: str,
    dna_hex: Optional[str] = None,
    entropy_hex: Optional[str] = None,
    use_fibpi: bool = True,
    chunk_size: int = 65520,
) -> Tuple[np.ndarray, Dict[str, Any]]:
    """
    Regenerate tensor using the PROVEN Q4_K_M centroid regen path.

    This calls into regrow_q4km.py / grow_q4km_weights, which has receipts.
    NO synthetic sampling - uses the exact same code path as the proofs.

    CC2-OFFSET: Offset/size semantics:
    - spec.name is used as `label` → unique RNG context per tensor
    - spec.offset is passed through for FibPi3D positioning (default 0)
    - spec.size_bytes overrides computed size if provided
    - Computed size = ceil(numel / 32) * 24 bytes for Q4_K_M

    Args:
        seed_hex: Helix seed SHA256 as hex string
        spec: TensorSpec with name, shape, dtype, offset, size_bytes
        atlas_path: Path to centroid atlas .npy file
        dna_hex: DNA seed hex (derived from seed_hex if None)
        entropy_hex: Entropy seed hex (derived from seed_hex if None)
        use_fibpi: Use FibPi3D parallel regen (default True)
        chunk_size: Chunk size for parallel regen

    Returns:
        (tensor_array, receipt_dict)

    Raises:
        FileNotFoundError: If atlas_path doesn't exist
        RuntimeError: If regen fails
    """
    # Validate atlas exists
    atlas_p = Path(atlas_path).expanduser()
    if not atlas_p.exists():
        raise FileNotFoundError(
            f"Centroid atlas not found: {atlas_path}\n"
            f"FlowTorch CRYSTAL_Q4KM requires a valid atlas_path in regrow config."
        )

    # Load atlas using proven loader
    atlas = load_centroid_atlas(atlas_p)

    # Derive DNA/entropy from seed_hex if not provided
    if dna_hex is None:
        dna_hex = seed_hex[:10]
    if entropy_hex is None:
        entropy_hex = seed_hex[10:18]

    # Calculate numel from shape
    numel = 1
    for d in spec.shape:
        numel *= d

    # CC2-OFFSET: Use provided size_bytes if available, otherwise compute
    if spec.size_bytes is not None:
        size_bytes = spec.size_bytes
        num_blocks = (size_bytes + 23) // 24  # Reverse compute for logging
    else:
        # Q4_K_M stores 32 values in 24 bytes
        # For numel values: ceil(numel / 32) * 24 bytes
        num_blocks = (numel + 31) // 32
        size_bytes = num_blocks * 24

    # Create buffer for streaming output
    buffer = BytesIO()

    # CC2-OFFSET: Pass spec.offset through to proven regen
    # For per-tensor regen, offset=0 is correct because spec.name provides uniqueness
    # offset > 0 would be used for parallel slicing of a SINGLE large tensor
    regen_offset = spec.offset

    # Call PROVEN regen implementation
    if use_fibpi:
        receipt = grow_q4km_weights_parallel(
            dna_hex=dna_hex,
            entropy_seed=entropy_hex,
            label=spec.name,  # Unique per tensor → unique RNG context
            offset=regen_offset,
            size=size_bytes,
            atlas=atlas,
            output_file=buffer,
            chunk_size=chunk_size,
            use_fibpi=True,
        )
    else:
        receipt = grow_q4km_weights(
            dna_hex=dna_hex,
            entropy_seed=entropy_hex,
            label=spec.name,  # Unique per tensor → unique RNG context
            offset=regen_offset,
            size=size_bytes,
            atlas=atlas,
            output_file=buffer,
        )

    # Decode Q4_K_M blocks to float values
    buffer.seek(0)
    raw_bytes = buffer.read()
    values = _decode_q4km_blocks_to_floats(raw_bytes, atlas, numel)

    # Convert to target dtype
    if spec.dtype == "float16":
        values = values.astype(np.float16)
    elif spec.dtype == "float32":
        pass  # Already float32

    # Reshape to spec.shape
    values = values.reshape(spec.shape)

    # Add bridge metadata to receipt
    receipt["bridge"] = {
        "module": "helix_cdc.regrow.flowtorch_bridge",
        "tensor_name": spec.name,
        "tensor_shape": spec.shape,
        "tensor_dtype": spec.dtype,
        "numel": numel,
        "q4km_blocks": num_blocks,
        "output_hash": hashlib.sha256(values.tobytes()).hexdigest()[:16],
    }

    return values, receipt


def _decode_q4km_blocks_to_floats(
    raw_bytes: bytes,
    atlas: np.ndarray,
    numel: int
) -> np.ndarray:
    """
    Decode Q4_K_M block bytes back to float32 values.

    Q4_K_M block format (24 bytes for 32 values):
    - 2 bytes: scale (fp16)
    - 2 bytes: zero (fp16)
    - 4 bytes: sign mask (32 bits)
    - 16 bytes: 4-bit codes (2 per byte)

    This is the INVERSE of synthesize_q4km_block.
    """
    block_size_bytes = 24
    values_per_block = 32
    num_blocks = len(raw_bytes) // block_size_bytes

    all_values = []

    for b in range(num_blocks):
        block = raw_bytes[b * block_size_bytes:(b + 1) * block_size_bytes]

        # Unpack scale, zero, sign_mask
        scale = struct.unpack('<e', block[0:2])[0]
        zero = struct.unpack('<e', block[2:4])[0]
        sign_mask = struct.unpack('<I', block[4:8])[0]

        # Unpack 4-bit codes
        codes = []
        for i in range(16):
            byte = block[8 + i]
            codes.append(byte & 0xF)
            codes.append((byte >> 4) & 0xF)

        # Reconstruct values: magnitude from atlas, sign from mask
        for i in range(values_per_block):
            code = codes[i]
            magnitude = atlas[code] if code < len(atlas) else 0.0
            sign = -1.0 if (sign_mask & (1 << i)) else 1.0
            value = magnitude * sign
            all_values.append(value)

    # Truncate to exact numel
    return np.array(all_values[:numel], dtype=np.float32)


def verify_against_known_digest(
    values: np.ndarray,
    expected_digest: str,
    tensor_name: str,
) -> bool:
    """
    Verify regenerated tensor matches a known digest (optional debug check).

    Args:
        values: Regenerated tensor
        expected_digest: Expected SHA256 hex prefix
        tensor_name: For error messages

    Returns:
        True if matches, False otherwise
    """
    actual = hashlib.sha256(values.tobytes()).hexdigest()
    matches = actual.startswith(expected_digest)

    if not matches:
        print(f"WARNING: {tensor_name} digest mismatch!")
        print(f"  Expected: {expected_digest}...")
        print(f"  Actual:   {actual[:len(expected_digest)]}...")

    return matches
