#!/usr/bin/env python3
"""
cdna_kernel.py - CDNA Kernel Primitive

Low-level kernel for CDNAv2 streaming compute.
Wraps CDNAv2Reader with a clean capability interface.

Three core operations (the kernel contract):
1. address(row_range) -> (block_id, local_start, local_end) tuples
2. fetch_block(block_id) -> raw indices [B, K] uint8
3. verify_block(block_id) -> bool (SHA256 check)

Everything else is built on these primitives.

Work Order: WO-CDNA-KERNEL-V1
"""

from __future__ import annotations

import hashlib
import time
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Iterator, List, Optional, Tuple, Union

import numpy as np

from helix_cdc.regrow.cdna_stream_v2 import (
    CDNAv2Reader,
    cdna2_matmul_block,
    BLOCK_INDEX_ENTRY_SIZE,
)


@dataclass
class KernelReceipt:
    """
    Receipt for kernel operations.

    Design: Every kernel operation emits a receipt that can be audited later
    without trusting agent logs. The receipt includes:
    - File path (canonical)
    - Tensor name + shape
    - Block structure (block_rows, num_blocks)
    - Block IDs touched
    - SHA256 of each block (stored vs verified)
    - Timings
    """
    operation: str
    timestamp: str = field(default_factory=lambda: datetime.utcnow().isoformat() + "Z")
    duration_ms: float = 0.0

    # File/tensor context
    file_path: Optional[str] = None
    tensor_name: Optional[str] = None
    tensor_shape: Optional[Tuple[int, int]] = None
    block_rows: Optional[int] = None
    num_blocks: Optional[int] = None

    # Block access info
    block_id: Optional[int] = None
    blocks_touched: Optional[List[int]] = None
    blocks_verified_count: int = 0

    # Verification
    row_range: Optional[Tuple[int, int]] = None
    verified: Optional[bool] = None
    sha256_stored: Optional[str] = None
    sha256_computed: Optional[str] = None

    # Extra context
    extra: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        d = {
            "schema": "kernel_receipt_v1",
            "operation": self.operation,
            "timestamp": self.timestamp,
            "duration_ms": self.duration_ms,
        }

        # File/tensor context
        if self.file_path is not None:
            d["file_path"] = self.file_path
        if self.tensor_name is not None:
            d["tensor_name"] = self.tensor_name
        if self.tensor_shape is not None:
            d["tensor_shape"] = list(self.tensor_shape)
        if self.block_rows is not None:
            d["block_rows"] = self.block_rows
        if self.num_blocks is not None:
            d["num_blocks"] = self.num_blocks

        # Block access
        if self.block_id is not None:
            d["block_id"] = self.block_id
        if self.blocks_touched is not None:
            d["blocks_touched"] = self.blocks_touched
        if self.blocks_verified_count > 0:
            d["blocks_verified_count"] = self.blocks_verified_count

        # Verification
        if self.row_range is not None:
            d["row_range"] = list(self.row_range)
        if self.verified is not None:
            d["verified"] = self.verified
        if self.sha256_stored is not None:
            d["sha256_stored"] = self.sha256_stored
        if self.sha256_computed is not None:
            d["sha256_computed"] = self.sha256_computed

        # Extra
        if self.extra:
            d.update(self.extra)

        return d


class CDNAKernel:
    """
    Low-level kernel for CDNA streaming compute.

    Wraps CDNAv2Reader with a clean capability interface.
    The kernel is deterministic and verifiable - no policy logic here.

    Core primitives:
    - address(start_row, end_row) -> block tuples
    - fetch_block(block_id) -> indices [B, K]
    - verify_block(block_id) -> bool

    Convenience methods (built on primitives):
    - get_rows(start, end, verify=True) -> indices
    - iter_blocks(verify=True) -> iterator

    Compute methods:
    - matmul_block(block_id, x) -> y_block
    - stream_matmul(x, verify=True) -> y

    Usage:
        k = CDNAKernel("tensor.cdna2.hxz")
        print(k.shape)  # (4096, 1024)

        # Core primitives
        blocks = k.address(0, 100)  # Which blocks cover rows 0-100?
        indices = k.fetch_block(0)  # Get block 0 indices
        ok = k.verify_block(0)      # Verify block 0 SHA256

        # Compute
        y = k.stream_matmul(x, verify=True)  # y = W @ x
    """

    def __init__(self, path: Union[str, Path]):
        """
        Initialize kernel from CDNAv2 file.

        Args:
            path: Path to .cdna2.hxz file
        """
        self.path = Path(path)
        self._reader = CDNAv2Reader(self.path)

        # Cache frequently accessed properties
        self._shape = (self._reader.rows, self._reader.cols)
        self._codebook = self._reader.codebook.copy()
        self._block_rows = self._reader.block_rows
        self._num_blocks = self._reader.num_blocks

        # Block cache (not policy-driven, just for avoiding re-decompression)
        # Agent layer handles cache policy
        self._block_cache: Dict[int, np.ndarray] = {}
        self._cache_enabled = False  # Off by default - agent controls this

    # ========== Core Primitives ==========

    def address(self, start_row: int, end_row: int) -> List[Tuple[int, int, int]]:
        """
        Map row range to (block_id, local_start, local_end) tuples.

        This is the core addressing primitive. Given a row range [start, end),
        returns which blocks contain those rows and the local offsets within
        each block.

        Args:
            start_row: First row (inclusive)
            end_row: Last row (exclusive)

        Returns:
            List of (block_id, local_start_row, local_end_row) tuples
            where local_start/end are offsets within that block.

        Example:
            >>> k = CDNAKernel("tensor.cdna2.hxz")  # block_rows=16
            >>> k.address(10, 40)
            [(0, 10, 16), (1, 0, 16), (2, 0, 8)]
        """
        if start_row < 0 or end_row > self._shape[0] or start_row >= end_row:
            raise ValueError(
                f"Invalid row range [{start_row}, {end_row}) "
                f"for tensor with {self._shape[0]} rows"
            )

        result = []

        # Calculate which blocks we need
        start_block = start_row // self._block_rows
        end_block = (end_row - 1) // self._block_rows + 1

        for block_idx in range(start_block, end_block):
            block_start_row = block_idx * self._block_rows

            # Handle last block which may be shorter
            if block_idx == self._num_blocks - 1:
                block_end_row = self._shape[0]
            else:
                block_end_row = block_start_row + self._block_rows

            # Calculate local offsets within this block
            local_start = max(0, start_row - block_start_row)
            local_end = min(block_end_row - block_start_row, end_row - block_start_row)

            if local_start < local_end:
                result.append((block_idx, local_start, local_end))

        return result

    def fetch_block(self, block_id: int) -> np.ndarray:
        """
        Fetch and decompress a single block.

        Core fetch primitive. Seeks to block position, reads compressed bytes,
        decompresses with brotli, returns as [B, K] uint8 array.

        Args:
            block_id: Block index (0 to num_blocks-1)

        Returns:
            indices [block_rows, K] uint8 (last block may have fewer rows)
        """
        if block_id < 0 or block_id >= self._num_blocks:
            raise ValueError(
                f"Block {block_id} out of range [0, {self._num_blocks})"
            )

        # Check cache if enabled
        if self._cache_enabled and block_id in self._block_cache:
            return self._block_cache[block_id].copy()

        # Fetch via reader's stream_blocks (which handles seek + decompress)
        for indices, start_row in self._reader.stream_blocks(block_id, block_id + 1):
            if self._cache_enabled:
                self._block_cache[block_id] = indices.copy()
            return indices

        # Should never reach here
        raise RuntimeError(f"Failed to fetch block {block_id}")

    def verify_block(self, block_id: int) -> bool:
        """
        Verify block SHA256.

        Core verification primitive. Decompresses block and computes SHA256,
        comparing against the stored hash in the block index.

        Args:
            block_id: Block index (0 to num_blocks-1)

        Returns:
            True if SHA256 matches, False otherwise
        """
        return self._reader.verify_block(block_id)

    def verify_block_with_receipt(self, block_id: int) -> Tuple[bool, str, str]:
        """
        Verify block SHA256 and return stored/computed hashes.

        Used for detailed receipts that prove which blocks were checked.

        Args:
            block_id: Block index (0 to num_blocks-1)

        Returns:
            (match: bool, stored_sha256: str, computed_sha256: str)
        """
        if block_id < 0 or block_id >= self._num_blocks:
            raise ValueError(f"Block {block_id} out of range [0, {self._num_blocks})")

        # Get stored SHA256 from index
        entry = self._reader.block_index[block_id]
        stored_sha = entry.sha256.hex()

        # Read and decompress block to compute SHA256
        import brotli
        with open(self.path, "rb") as f:
            f.seek(entry.offset)
            compressed = f.read(entry.comp_len)
        raw_bytes = brotli.decompress(compressed)
        computed_sha = hashlib.sha256(raw_bytes).hexdigest()

        return (stored_sha == computed_sha, stored_sha, computed_sha)

    # ========== Convenience Methods ==========

    def get_rows(
        self,
        start_row: int,
        end_row: int,
        verify: bool = True,
    ) -> np.ndarray:
        """
        Random access to specific rows with optional verification.

        Built on address() and fetch_block() primitives.

        Args:
            start_row: First row (inclusive)
            end_row: Last row (exclusive)
            verify: If True, verify each block before returning

        Returns:
            indices [end_row - start_row, K] uint8
        """
        blocks_needed = self.address(start_row, end_row)
        result_rows = []

        for block_id, local_start, local_end in blocks_needed:
            if verify and not self.verify_block(block_id):
                raise ValueError(f"Block {block_id} failed SHA256 verification")

            block_indices = self.fetch_block(block_id)
            result_rows.append(block_indices[local_start:local_end])

        return np.vstack(result_rows)

    def iter_blocks(
        self,
        verify: bool = True,
    ) -> Iterator[Tuple[np.ndarray, int, KernelReceipt]]:
        """
        Stream all blocks with optional verification.

        Built on fetch_block() and verify_block() primitives.

        Args:
            verify: If True, verify each block before yielding

        Yields:
            (indices [B, K] uint8, start_row, receipt)
        """
        for block_id in range(self._num_blocks):
            t0 = time.perf_counter()

            verified = None
            if verify:
                verified = self.verify_block(block_id)
                if not verified:
                    raise ValueError(f"Block {block_id} failed SHA256 verification")

            indices = self.fetch_block(block_id)
            start_row = block_id * self._block_rows

            duration_ms = (time.perf_counter() - t0) * 1000

            receipt = KernelReceipt(
                operation="iter_block",
                duration_ms=duration_ms,
                block_id=block_id,
                row_range=(start_row, start_row + indices.shape[0]),
                verified=verified,
            )

            yield indices, start_row, receipt

    # ========== Compute Methods ==========

    def matmul_block(self, block_id: int, x: np.ndarray) -> np.ndarray:
        """
        Compute y = W_block @ x without decoding W to float32.

        Uses CDNA dot product trick: bincount + codebook matmul.

        Args:
            block_id: Block index
            x: Input vector [K]

        Returns:
            y_block [block_rows] (or fewer for last block)
        """
        x = np.asarray(x, dtype=np.float32).ravel()
        if len(x) != self._shape[1]:
            raise ValueError(
                f"Input dimension {len(x)} != tensor columns {self._shape[1]}"
            )

        indices = self.fetch_block(block_id)
        return cdna2_matmul_block(indices, self._codebook, x)

    def stream_matmul(
        self,
        x: np.ndarray,
        verify: bool = True,
        detailed_verification: bool = False,
    ) -> Tuple[np.ndarray, KernelReceipt]:
        """
        Compute y = W @ x via streaming blocks.

        Memory-efficient matmul that never materializes full W.

        Args:
            x: Input vector [K]
            verify: If True, verify each block before computing
            detailed_verification: If True, include per-block SHA256 in receipt

        Returns:
            (y [M], receipt)
        """
        t0 = time.perf_counter()

        x = np.asarray(x, dtype=np.float32).ravel()
        if len(x) != self._shape[1]:
            raise ValueError(
                f"Input dimension {len(x)} != tensor columns {self._shape[1]}"
            )

        y_chunks = []
        blocks_processed = 0
        blocks_verified = 0
        block_verification_detail = {}

        for block_id in range(self._num_blocks):
            if verify:
                if detailed_verification:
                    match, stored, computed = self.verify_block_with_receipt(block_id)
                    block_verification_detail[str(block_id)] = {
                        "stored": stored,
                        "computed": computed,
                        "match": match,
                    }
                    if not match:
                        raise ValueError(f"Block {block_id} failed SHA256 verification")
                else:
                    if not self.verify_block(block_id):
                        raise ValueError(f"Block {block_id} failed SHA256 verification")
                blocks_verified += 1

            y_chunk = self.matmul_block(block_id, x)
            y_chunks.append(y_chunk)
            blocks_processed += 1

        y = np.concatenate(y_chunks)
        duration_ms = (time.perf_counter() - t0) * 1000

        # Build enriched receipt
        extra = {
            "blocks_processed": blocks_processed,
            "blocks_verified": blocks_verified,
            "output_shape": list(y.shape),
            "input_shape": list(x.shape),
        }
        if detailed_verification and block_verification_detail:
            extra["block_verification"] = block_verification_detail

        receipt = KernelReceipt(
            operation="stream_matmul",
            duration_ms=duration_ms,
            file_path=str(self.path.resolve()),
            tensor_name=self.tensor_name,
            tensor_shape=self._shape,
            block_rows=self._block_rows,
            num_blocks=self._num_blocks,
            blocks_touched=list(range(self._num_blocks)),
            blocks_verified_count=blocks_verified,
            verified=verify,
            extra=extra,
        )

        return y, receipt

    def stream_matmul_row_range(
        self,
        x: np.ndarray,
        start_row: int,
        end_row: int,
        verify: bool = True,
        detailed_verification: bool = False,
        sidecar_path: Optional[Path] = None,
    ) -> Tuple[np.ndarray, KernelReceipt]:
        """
        Compute W[start_row:end_row] @ x using only needed blocks.

        Uses address() primitive to map row ranges to blocks, then
        fetches and verifies only those blocks before computing.

        This is the "holy grail" primitive: proves which tensor rows were
        used, which blocks were decompressed, and provides cryptographic
        auditability of the computation.

        Args:
            x: Input vector [K]
            start_row: First row (inclusive)
            end_row: Last row (exclusive)
            verify: If True, verify SHA256 of each touched block
            detailed_verification: If True, include per-block SHA256 in receipt
            sidecar_path: Optional path to HXZO sidecar for outlier correction

        Returns:
            (y [end_row - start_row], receipt)

        Example:
            >>> k = CDNAKernel("ffn_down.cdna2.hxz")
            >>> x = np.random.randn(14336).astype(np.float32)
            >>> y_partial, receipt = k.stream_matmul_row_range(x, 100, 200)
            >>> print(y_partial.shape)  # (100,)
            >>> print(receipt.extra["savings_factor"])  # e.g., 40.96x
        """
        import hashlib as _hashlib

        t0 = time.perf_counter()

        x = np.asarray(x, dtype=np.float32).ravel()
        if len(x) != self._shape[1]:
            raise ValueError(
                f"Input dimension {len(x)} != tensor columns {self._shape[1]}"
            )

        if start_row < 0 or end_row > self._shape[0] or start_row >= end_row:
            raise ValueError(
                f"Invalid row range [{start_row}, {end_row}) for tensor with {self._shape[0]} rows"
            )

        # Use address() to find needed blocks
        blocks_needed = self.address(start_row, end_row)

        # Load sidecar data if provided
        sidecar_positions = None
        sidecar_values = None
        sidecar_corrections = 0

        if sidecar_path is not None:
            from helix_cdc.regrow.hxzo_sidecar import read_outlier_sidecar
            sidecar_positions, sidecar_values, _ = read_outlier_sidecar(str(sidecar_path))

        y_chunks = []
        blocks_verified = 0
        block_verification_detail = {}
        bytes_decompressed = 0

        for block_id, local_start, local_end in blocks_needed:
            # Verify block if requested
            if verify:
                if detailed_verification:
                    match, stored, computed = self.verify_block_with_receipt(block_id)
                    block_verification_detail[str(block_id)] = {
                        "stored": stored[:16] + "...",
                        "computed": computed[:16] + "...",
                        "match": match,
                    }
                    if not match:
                        raise ValueError(f"Block {block_id} failed SHA256 verification")
                else:
                    if not self.verify_block(block_id):
                        raise ValueError(f"Block {block_id} failed SHA256 verification")
                blocks_verified += 1

            # Fetch block indices
            indices = self.fetch_block(block_id)
            bytes_decompressed += indices.nbytes

            # Get the needed slice
            row_indices = indices[local_start:local_end]

            # Dequantize: W_slice[i,j] = codebook[row_indices[i,j]]
            weights_slice = self._codebook[row_indices]

            # Apply sidecar corrections if available
            if sidecar_positions is not None:
                # Calculate global row range for this chunk
                block_start_row = block_id * self._block_rows
                global_start_row = block_start_row + local_start
                global_end_row = block_start_row + local_end

                for i in range(len(sidecar_positions)):
                    pos = sidecar_positions[i]
                    row = pos // self._shape[1]
                    col = pos % self._shape[1]

                    if global_start_row <= row < global_end_row:
                        local_row = row - global_start_row
                        weights_slice[local_row, col] = float(sidecar_values[i])
                        sidecar_corrections += 1

            # Compute y_chunk = W_slice @ x
            y_chunk = weights_slice @ x
            y_chunks.append(y_chunk)

        y = np.concatenate(y_chunks)
        duration_ms = (time.perf_counter() - t0) * 1000

        # Compute proof hashes
        input_x_sha256 = _hashlib.sha256(x.tobytes()).hexdigest()
        output_y_sha256 = _hashlib.sha256(y.tobytes()).hexdigest()

        # Memory stats
        bytes_if_full = self._shape[0] * self._shape[1]  # Full tensor as uint8
        full_tensor_float32 = self._shape[0] * self._shape[1] * 4
        row_range_float32 = (end_row - start_row) * self._shape[1] * 4
        savings_factor = round(bytes_if_full / bytes_decompressed, 2) if bytes_decompressed > 0 else 0

        # Build extra context for receipt
        extra = {
            "blocks_needed": [(b, ls, le) for b, ls, le in blocks_needed],
            "blocks_verified": blocks_verified,
            "bytes_decompressed": bytes_decompressed,
            "bytes_if_full_load": bytes_if_full,
            "full_tensor_float32_bytes": full_tensor_float32,
            "row_range_float32_bytes": row_range_float32,
            "savings_factor": savings_factor,
            "sidecar_applied": sidecar_path is not None,
            "sidecar_corrections": sidecar_corrections,
            "compute_proof": {
                "input_x_sha256": input_x_sha256,
                "output_y_sha256": output_y_sha256,
            },
            "output_shape": list(y.shape),
            "input_shape": list(x.shape),
        }

        if detailed_verification and block_verification_detail:
            extra["block_verification"] = block_verification_detail

        if sidecar_path is not None:
            extra["sidecar_path"] = str(sidecar_path)

        blocks_touched = [b for b, _, _ in blocks_needed]

        receipt = KernelReceipt(
            operation="stream_matmul_row_range",
            duration_ms=duration_ms,
            file_path=str(self.path.resolve()),
            tensor_name=self.tensor_name,
            tensor_shape=self._shape,
            block_rows=self._block_rows,
            num_blocks=self._num_blocks,
            blocks_touched=blocks_touched,
            blocks_verified_count=blocks_verified,
            row_range=(start_row, end_row),
            verified=verify,
            extra=extra,
        )

        return y, receipt

    # ========== Properties ==========

    @property
    def shape(self) -> Tuple[int, int]:
        """Tensor shape (rows, cols)."""
        return self._shape

    @property
    def rows(self) -> int:
        """Number of rows (M)."""
        return self._shape[0]

    @property
    def cols(self) -> int:
        """Number of columns (K)."""
        return self._shape[1]

    @property
    def codebook(self) -> np.ndarray:
        """Codebook [256] float32."""
        return self._codebook.copy()

    @property
    def num_blocks(self) -> int:
        """Number of blocks in the file."""
        return self._num_blocks

    @property
    def block_rows(self) -> int:
        """Rows per block."""
        return self._block_rows

    @property
    def tensor_name(self) -> str:
        """Tensor name from header."""
        return self._reader.tensor_name

    @property
    def memory_stats(self) -> Dict[str, Any]:
        """Get memory statistics."""
        return self._reader.memory_stats

    # ========== Cache Control ==========

    def enable_cache(self) -> None:
        """Enable block caching (agent should call this)."""
        self._cache_enabled = True

    def disable_cache(self) -> None:
        """Disable block caching."""
        self._cache_enabled = False
        self._block_cache.clear()

    def clear_cache(self) -> None:
        """Clear block cache."""
        self._block_cache.clear()

    @property
    def cache_info(self) -> Dict[str, Any]:
        """Get cache statistics."""
        return {
            "enabled": self._cache_enabled,
            "blocks_cached": len(self._block_cache),
            "cached_block_ids": list(self._block_cache.keys()),
        }

    # ========== Verification ==========

    def verify_all_blocks(self) -> Tuple[int, int, List[int]]:
        """
        Verify all blocks.

        Returns:
            (ok_count, fail_count, failed_block_indices)
        """
        return self._reader.verify_all_blocks()

    def get_block_sha256(self, block_id: int) -> str:
        """Get stored SHA256 for a block (from index, not computed)."""
        if block_id < 0 or block_id >= self._num_blocks:
            raise ValueError(f"Block {block_id} out of range")
        entry = self._reader.block_index[block_id]
        return entry.sha256.hex()

    # ========== String Representation ==========

    def __repr__(self) -> str:
        return (
            f"CDNAKernel({self.path.name!r}, "
            f"shape={self.shape}, "
            f"blocks={self.num_blocks}, "
            f"block_rows={self.block_rows})"
        )


# CLI moved to helix_cdc/__main__.py (cdna-fetch-rows, cdna-matmul, cdna-kernel-kat)
# Run: python -m helix_cdc cdna-kernel-kat <path.cdna2.hxz>
