#!/usr/bin/env python3
"""
hxz1_stream.py - Streaming tensor access from HXZ1 seeds

Enables compute-from-seed: run forward passes where weight tensors
are never fully materialized in RAM. Instead, rows/blocks are streamed
from the compressed seed on demand.

Key difference from full decode:
  - Full decode: seed → decompress → full tensor [M, N] → compute
  - Streaming:   seed → decompress → yield row blocks [B, N] → compute per block

This is the foundation for "compressed domain compute" claims.

Usage:
    from helix_cdc.regrow.hxz1_stream import stream_hxz_tensor_rows

    for rows, start_idx in stream_hxz_tensor_rows("shard.hxz", block_rows=32):
        y_chunk = rows @ x
        ...
"""

from __future__ import annotations

import json
import zlib
from pathlib import Path
from typing import Any, Dict, Generator, Optional, Tuple

import numpy as np

from helix_cdc.regrow.hxz_model_loader import inspect_hxz_header, HXZHeader


# Dtype mapping (same as hxz_model_loader)
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,
}


class HXZ1StreamReader:
    """
    Streaming reader for HXZ1 seeds.

    Decompresses the blob once, then provides row-block views
    without creating the full [M, N] tensor as a first-class object.
    """

    def __init__(self, hxz_path: str, target_dtype: str = "float32"):
        """
        Initialize stream reader.

        Args:
            hxz_path: Path to .hxz file
            target_dtype: Target dtype for output blocks ("float32" or "float16")
        """
        self.hxz_path = Path(hxz_path)
        self.target_dtype = np.float32 if target_dtype == "float32" else np.float16

        # Parse header
        self.header = inspect_hxz_header(str(self.hxz_path))

        if self.header.format != "binary_v1":
            raise ValueError(f"HXZ1 stream only supports binary_v1, got {self.header.format}")

        if not self.header.registry:
            raise ValueError("HXZ1 stream requires registry in header")

        # Get first tensor entry
        self.entry = self.header.registry[0]
        self.tensor_name = self.entry["name"]
        self.shape = tuple(self.entry.get("shape", []))
        self.source_dtype = DTYPE_MAP.get(self.entry.get("dtype", "float32"), np.float32)

        if len(self.shape) != 2:
            raise ValueError(f"HXZ1 stream requires 2D tensor, got shape {self.shape}")

        self.rows, self.cols = self.shape

        # Decompress blob (done once)
        self._blob: Optional[bytes] = None

    def _ensure_blob(self) -> bytes:
        """Load and decompress blob if not already done."""
        if self._blob is not None:
            return self._blob

        with open(self.hxz_path, "rb") as f:
            f.seek(self.header.blob_offset)
            compressed = f.read()

        if self.header.encoding == "zlib":
            self._blob = zlib.decompress(compressed)
        else:
            self._blob = compressed

        return self._blob

    def stream_rows(
        self,
        block_rows: int = 32,
    ) -> Generator[Tuple[np.ndarray, int], None, None]:
        """
        Stream tensor rows in blocks.

        Args:
            block_rows: Number of rows per block

        Yields:
            (row_block, start_row_idx) where row_block is [B, cols]
        """
        blob = self._ensure_blob()

        # Calculate bytes per row
        bytes_per_elem = self.source_dtype().itemsize
        bytes_per_row = self.cols * bytes_per_elem

        offset = self.entry.get("offset", 0)

        for start_row in range(0, self.rows, block_rows):
            end_row = min(start_row + block_rows, self.rows)
            num_rows = end_row - start_row

            # Calculate byte range
            byte_start = offset + start_row * bytes_per_row
            byte_end = offset + end_row * bytes_per_row

            # Extract bytes for this block
            block_bytes = blob[byte_start:byte_end]

            # Convert to array
            block = np.frombuffer(block_bytes, dtype=self.source_dtype)
            block = block.reshape(num_rows, self.cols)

            # Convert to target dtype if needed
            if block.dtype != self.target_dtype:
                block = block.astype(self.target_dtype)

            yield block, start_row

    def get_full_tensor(self) -> np.ndarray:
        """
        Get full tensor (for comparison baseline).

        This defeats the purpose of streaming but is useful for verification.
        """
        blob = self._ensure_blob()

        offset = self.entry.get("offset", 0)
        length = self.entry.get("length", len(blob) - offset)

        tensor_bytes = blob[offset:offset + length]
        arr = np.frombuffer(tensor_bytes, dtype=self.source_dtype)
        arr = arr.reshape(self.shape)

        if arr.dtype != self.target_dtype:
            arr = arr.astype(self.target_dtype)

        return arr

    @property
    def memory_stats(self) -> Dict[str, Any]:
        """Get memory statistics for the tensor."""
        full_tensor_bytes = self.rows * self.cols * self.target_dtype().itemsize
        return {
            "shape": list(self.shape),
            "dtype": str(self.source_dtype),
            "target_dtype": str(self.target_dtype),
            "full_tensor_bytes": full_tensor_bytes,
            "rows": self.rows,
            "cols": self.cols,
        }


def stream_hxz_tensor_rows(
    hxz_path: str,
    block_rows: int = 32,
    target_dtype: str = "float32",
) -> Generator[Tuple[np.ndarray, int], None, None]:
    """
    Stream tensor rows from HXZ1 seed.

    This is the main API for compute-from-seed.

    Args:
        hxz_path: Path to .hxz file
        block_rows: Number of rows per block
        target_dtype: Target dtype ("float32" or "float16")

    Yields:
        (row_block, start_row_idx) where row_block is [B, cols]

    Example:
        y_chunks = []
        for rows, _ in stream_hxz_tensor_rows("shard.hxz", block_rows=32):
            y_chunks.append(rows @ x)
        y = np.concatenate(y_chunks, axis=0)
    """
    reader = HXZ1StreamReader(hxz_path, target_dtype)
    yield from reader.stream_rows(block_rows)


def stream_matmul(
    hxz_path: str,
    x: np.ndarray,
    block_rows: int = 32,
) -> Tuple[np.ndarray, Dict[str, Any]]:
    """
    Compute y = W @ x where W is streamed from HXZ1 seed.

    The full weight matrix W is never materialized.

    Args:
        hxz_path: Path to .hxz file containing W
        x: Input vector [N] or [N, K]
        block_rows: Number of rows per block

    Returns:
        (y, stats) where y is [M] or [M, K] and stats has compute info
    """
    reader = HXZ1StreamReader(hxz_path, target_dtype="float32")

    # Ensure x is float32 and correct shape
    x = np.asarray(x, dtype=np.float32)
    if x.ndim == 1:
        x = x.reshape(-1, 1)
        squeeze_output = True
    else:
        squeeze_output = False

    # Check dimensions
    if x.shape[0] != reader.cols:
        raise ValueError(
            f"Input dimension mismatch: W is [{reader.rows}, {reader.cols}], "
            f"x is [{x.shape[0]}, ...]"
        )

    # Compute in blocks
    y_chunks = []
    blocks_processed = 0
    max_block_bytes = 0

    for rows, start_idx in reader.stream_rows(block_rows):
        # rows: [B, cols], x: [cols, K]
        y_chunk = rows @ x  # [B, K]
        y_chunks.append(y_chunk)

        blocks_processed += 1
        block_bytes = rows.nbytes
        max_block_bytes = max(max_block_bytes, block_bytes)

    # Concatenate results
    y = np.concatenate(y_chunks, axis=0)

    if squeeze_output:
        y = y.squeeze(-1)

    stats = {
        "tensor_name": reader.tensor_name,
        "tensor_shape": list(reader.shape),
        "block_rows": block_rows,
        "blocks_processed": blocks_processed,
        "max_block_bytes": max_block_bytes,
        "full_tensor_bytes": reader.memory_stats["full_tensor_bytes"],
        "memory_ratio": reader.memory_stats["full_tensor_bytes"] / max_block_bytes,
    }

    return y, stats


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

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

    hxz_path = sys.argv[1]
    block_rows = int(sys.argv[2]) if len(sys.argv) > 2 else 32

    print(f"=== HXZ1 Stream Test ===")
    print(f"File: {hxz_path}")
    print(f"Block rows: {block_rows}")
    print()

    reader = HXZ1StreamReader(hxz_path)
    print(f"Tensor: {reader.tensor_name}")
    print(f"Shape: {reader.shape}")
    print(f"Source dtype: {reader.source_dtype}")
    print()

    # Stream and count
    block_count = 0
    for rows, start_idx in reader.stream_rows(block_rows):
        block_count += 1
        if block_count <= 3:
            print(f"Block {block_count}: rows {start_idx}-{start_idx + len(rows) - 1}, shape {rows.shape}")

    print(f"...")
    print(f"Total blocks: {block_count}")
