#!/usr/bin/env python3
"""
fibpi_chunk_agent.py - FibPi Chunk Agent

Policy brain for CDNA block routing and caching.
Sits above CDNAKernel and makes routing/verification decisions.

Responsibilities:
- Choose verify policy (always, sampled, trust_cached)
- Maintain block cache keyed by (path, block_id, sha256)
- Route between CDNAv2 blocks vs regen seeds (future)
- Emit access pattern receipts

The agent decides WHEN to verify and WHAT to cache.
The kernel executes the actual operations.

Work Order: WO-FIBPI-CHUNK-AGENT-V1
"""

from __future__ import annotations

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

import numpy as np

from helix_cdc.kernel.cdna_kernel import CDNAKernel, KernelReceipt


VerifyPolicy = Literal["always", "sampled", "trust_cached"]


@dataclass
class CacheEntry:
    """Single entry in the block cache."""
    path: str
    block_id: int
    sha256: str  # Expected SHA256
    data: np.ndarray
    verified: bool
    last_access: float = field(default_factory=time.time)
    access_count: int = 1


@dataclass
class AccessReceipt:
    """Receipt for agent access operations."""
    operation: str
    timestamp: str = field(default_factory=lambda: datetime.utcnow().isoformat() + "Z")
    duration_ms: float = 0.0
    cache_hit: bool = False
    verified: bool = False
    blocks_accessed: int = 0
    blocks_verified: int = 0
    cache_size_mb: float = 0.0
    extra: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "schema": "access_receipt_v1",
            "operation": self.operation,
            "timestamp": self.timestamp,
            "duration_ms": self.duration_ms,
            "cache_hit": self.cache_hit,
            "verified": self.verified,
            "blocks_accessed": self.blocks_accessed,
            "blocks_verified": self.blocks_verified,
            "cache_size_mb": self.cache_size_mb,
            **self.extra,
        }


class LRUCache:
    """Simple LRU cache with size limit."""

    def __init__(self, max_size_mb: float = 256.0):
        self.max_size_bytes = max_size_mb * 1024 * 1024
        self.entries: OrderedDict[str, CacheEntry] = OrderedDict()
        self.current_size_bytes = 0

    def _make_key(self, path: str, block_id: int) -> str:
        return f"{path}:{block_id}"

    def get(self, path: str, block_id: int) -> Optional[CacheEntry]:
        key = self._make_key(path, block_id)
        if key not in self.entries:
            return None

        # Move to end (most recently used)
        entry = self.entries.pop(key)
        entry.last_access = time.time()
        entry.access_count += 1
        self.entries[key] = entry
        return entry

    def put(
        self,
        path: str,
        block_id: int,
        sha256: str,
        data: np.ndarray,
        verified: bool,
    ) -> None:
        key = self._make_key(path, block_id)
        data_size = data.nbytes

        # Evict entries until we have room
        while self.current_size_bytes + data_size > self.max_size_bytes:
            if not self.entries:
                break
            # Pop oldest (first) entry
            oldest_key, oldest_entry = self.entries.popitem(last=False)
            self.current_size_bytes -= oldest_entry.data.nbytes

        # Add new entry
        entry = CacheEntry(
            path=path,
            block_id=block_id,
            sha256=sha256,
            data=data.copy(),
            verified=verified,
        )
        self.entries[key] = entry
        self.current_size_bytes += data_size

    def clear(self) -> None:
        self.entries.clear()
        self.current_size_bytes = 0

    def stats(self) -> Dict[str, Any]:
        return {
            "entries": len(self.entries),
            "size_bytes": self.current_size_bytes,
            "size_mb": self.current_size_bytes / (1024 * 1024),
            "max_size_mb": self.max_size_bytes / (1024 * 1024),
            "utilization": self.current_size_bytes / self.max_size_bytes if self.max_size_bytes > 0 else 0,
        }


class FibPiChunkAgent:
    """
    Policy brain for CDNA block routing.

    Responsibilities:
    - Choose verify policy (always, sampled, trust_cached)
    - Maintain block cache keyed by (path, block_id, sha256)
    - Route between CDNAv2 blocks vs regen seeds (future)
    - Emit access pattern receipts

    Example:
        agent = FibPiChunkAgent(verify_policy="trust_cached")
        kernel = agent.open_capsule("tensor.cdna2.hxz")
        rows = agent.fetch_rows(kernel, 0, 100)
        print(agent.cache_stats())

    Verify policies:
    - "always": Verify every block on every access (safest, slowest)
    - "sampled": Verify random subset (~10%) of blocks
    - "trust_cached": Verify on first access, trust cache thereafter
    """

    def __init__(
        self,
        verify_policy: VerifyPolicy = "trust_cached",
        cache_size_mb: float = 256.0,
        sample_rate: float = 0.1,
    ):
        """
        Initialize agent.

        Args:
            verify_policy: How to verify blocks ("always", "sampled", "trust_cached")
            cache_size_mb: Maximum cache size in MB
            sample_rate: For "sampled" policy, fraction of blocks to verify (0-1)
        """
        self.verify_policy = verify_policy
        self.sample_rate = sample_rate
        self._cache = LRUCache(max_size_mb=cache_size_mb)

        # Track access patterns
        self._access_count = 0
        self._cache_hits = 0
        self._cache_misses = 0
        self._blocks_verified = 0
        self._verification_failures = 0

        # Open kernels (path -> kernel)
        self._kernels: Dict[str, CDNAKernel] = {}

    # ========== Kernel Management ==========

    def open_capsule(self, path: Union[str, Path]) -> CDNAKernel:
        """
        Open CDNA file and return kernel handle.

        Caches kernel instances by CANONICAL path (resolved realpath).
        This prevents "same file via two paths" double-caching and
        weird trust boundaries.

        Args:
            path: Path to .cdna2.hxz file

        Returns:
            CDNAKernel instance
        """
        # CRITICAL: Use canonical path to avoid double-caching
        # same file via different paths (symlinks, relative, etc.)
        canonical_path = str(Path(path).resolve())

        if canonical_path not in self._kernels:
            self._kernels[canonical_path] = CDNAKernel(canonical_path)

        return self._kernels[canonical_path]

    def close_capsule(self, path: Union[str, Path]) -> None:
        """Close and remove kernel from cache."""
        path = str(Path(path).resolve())
        if path in self._kernels:
            del self._kernels[path]

    # ========== Policy-Driven Access ==========

    def _should_verify(self, block_id: int, cached: bool) -> bool:
        """Determine if block should be verified based on policy."""
        if self.verify_policy == "always":
            return True
        elif self.verify_policy == "sampled":
            return random.random() < self.sample_rate
        elif self.verify_policy == "trust_cached":
            return not cached
        else:
            return True

    def fetch_block(
        self,
        kernel: CDNAKernel,
        block_id: int,
    ) -> Tuple[np.ndarray, AccessReceipt]:
        """
        Fetch single block with policy-driven caching and verification.

        Args:
            kernel: CDNAKernel instance
            block_id: Block to fetch

        Returns:
            (indices [B, K] uint8, receipt)
        """
        t0 = time.perf_counter()
        path = str(kernel.path)

        # Check cache
        cache_entry = self._cache.get(path, block_id)
        cache_hit = cache_entry is not None

        if cache_hit:
            self._cache_hits += 1
        else:
            self._cache_misses += 1

        # Determine if we should verify
        should_verify = self._should_verify(block_id, cache_hit)
        verified = False

        if cache_hit and not should_verify:
            # Trust cache
            indices = cache_entry.data.copy()
            verified = cache_entry.verified
        else:
            # Fetch from kernel
            if should_verify:
                if not kernel.verify_block(block_id):
                    self._verification_failures += 1
                    raise ValueError(f"Block {block_id} failed SHA256 verification")
                self._blocks_verified += 1
                verified = True

            indices = kernel.fetch_block(block_id)

            # Cache result
            sha256 = kernel.get_block_sha256(block_id)
            self._cache.put(path, block_id, sha256, indices, verified)

        self._access_count += 1
        duration_ms = (time.perf_counter() - t0) * 1000

        receipt = AccessReceipt(
            operation="fetch_block",
            duration_ms=duration_ms,
            cache_hit=cache_hit,
            verified=verified,
            blocks_accessed=1,
            blocks_verified=1 if verified and should_verify else 0,
            cache_size_mb=self._cache.current_size_bytes / (1024 * 1024),
            extra={"block_id": block_id},
        )

        return indices, receipt

    def fetch_rows(
        self,
        kernel: CDNAKernel,
        start_row: int,
        end_row: int,
    ) -> Tuple[np.ndarray, AccessReceipt]:
        """
        Fetch rows with policy-driven caching and verification.

        Args:
            kernel: CDNAKernel instance
            start_row: First row (inclusive)
            end_row: Last row (exclusive)

        Returns:
            (indices [N, K] uint8, receipt)
        """
        t0 = time.perf_counter()

        # Get block mapping
        blocks_needed = kernel.address(start_row, end_row)

        result_rows = []
        total_cache_hits = 0
        total_verified = 0

        for block_id, local_start, local_end in blocks_needed:
            block_indices, block_receipt = self.fetch_block(kernel, block_id)
            result_rows.append(block_indices[local_start:local_end])

            if block_receipt.cache_hit:
                total_cache_hits += 1
            if block_receipt.verified:
                total_verified += 1

        indices = np.vstack(result_rows) if result_rows else np.array([], dtype=np.uint8)
        duration_ms = (time.perf_counter() - t0) * 1000

        receipt = AccessReceipt(
            operation="fetch_rows",
            duration_ms=duration_ms,
            cache_hit=total_cache_hits == len(blocks_needed),
            verified=total_verified > 0,
            blocks_accessed=len(blocks_needed),
            blocks_verified=total_verified,
            cache_size_mb=self._cache.current_size_bytes / (1024 * 1024),
            extra={
                "row_range": [start_row, end_row],
                "blocks_from_cache": total_cache_hits,
            },
        )

        return indices, receipt

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

    def stream_matmul(
        self,
        kernel: CDNAKernel,
        x: np.ndarray,
    ) -> Tuple[np.ndarray, AccessReceipt]:
        """
        Compute y = W @ x using kernel with agent's policy.

        Args:
            kernel: CDNAKernel instance
            x: Input vector [K]

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

        x = np.asarray(x, dtype=np.float32).ravel()
        if len(x) != kernel.cols:
            raise ValueError(
                f"Input dimension {len(x)} != tensor columns {kernel.cols}"
            )

        y_chunks = []
        total_cache_hits = 0
        total_verified = 0

        for block_id in range(kernel.num_blocks):
            block_indices, block_receipt = self.fetch_block(kernel, block_id)

            # Compute matmul for this block
            from helix_cdc.regrow.cdna_stream_v2 import cdna2_matmul_block
            y_chunk = cdna2_matmul_block(block_indices, kernel.codebook, x)
            y_chunks.append(y_chunk)

            if block_receipt.cache_hit:
                total_cache_hits += 1
            if block_receipt.verified:
                total_verified += 1

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

        receipt = AccessReceipt(
            operation="stream_matmul",
            duration_ms=duration_ms,
            cache_hit=total_cache_hits == kernel.num_blocks,
            verified=total_verified > 0,
            blocks_accessed=kernel.num_blocks,
            blocks_verified=total_verified,
            cache_size_mb=self._cache.current_size_bytes / (1024 * 1024),
            extra={
                "blocks_from_cache": total_cache_hits,
                "output_shape": list(y.shape),
            },
        )

        return y, receipt

    # ========== Routing (Future) ==========

    def route_matmul(
        self,
        tensor_id: str,
        x: np.ndarray,
        mode: str = "cdna",  # "cdna" | "regen" | "hybrid"
    ) -> Tuple[np.ndarray, AccessReceipt]:
        """
        Route matmul to appropriate backend.

        Currently only supports CDNA mode. Regen and hybrid modes
        are placeholders for future expansion.

        Args:
            tensor_id: Path to tensor file or tensor identifier
            x: Input vector [K]
            mode: Routing mode ("cdna", "regen", "hybrid")

        Returns:
            (y [M], receipt)
        """
        if mode != "cdna":
            raise NotImplementedError(
                f"Mode '{mode}' not yet implemented. Use 'cdna'."
            )

        kernel = self.open_capsule(tensor_id)
        return self.stream_matmul(kernel, x)

    # ========== Cache Management ==========

    def cache_stats(self) -> Dict[str, Any]:
        """Return cache statistics."""
        return {
            "policy": self.verify_policy,
            "sample_rate": self.sample_rate if self.verify_policy == "sampled" else None,
            "cache": self._cache.stats(),
            "access_count": self._access_count,
            "cache_hits": self._cache_hits,
            "cache_misses": self._cache_misses,
            "hit_rate": self._cache_hits / max(1, self._access_count),
            "blocks_verified": self._blocks_verified,
            "verification_failures": self._verification_failures,
            "open_kernels": len(self._kernels),
        }

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

    def get_access_receipt(self) -> Dict[str, Any]:
        """Return access pattern summary."""
        return {
            "schema": "access_summary_v1",
            "timestamp": datetime.utcnow().isoformat() + "Z",
            "verify_policy": self.verify_policy,
            "total_accesses": self._access_count,
            "cache_hits": self._cache_hits,
            "cache_misses": self._cache_misses,
            "hit_rate": self._cache_hits / max(1, self._access_count),
            "blocks_verified": self._blocks_verified,
            "verification_failures": self._verification_failures,
            "cache": self._cache.stats(),
        }

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

    def __repr__(self) -> str:
        return (
            f"FibPiChunkAgent(policy={self.verify_policy!r}, "
            f"cache_mb={self._cache.max_size_bytes / (1024 * 1024):.1f}, "
            f"kernels={len(self._kernels)})"
        )


    # ========== Attention Layer Integration (WO-AGENT-ATTN-01) ==========

    def stream_attention_forward(
        self,
        X: np.ndarray,
        block_index: int,
        manifest_path: Union[str, Path],
        d_head: int = 128,
    ) -> Tuple[np.ndarray, List[AccessReceipt]]:
        """
        Full attention forward pass with policy-driven verification.

        Computes: softmax(X @ Wq @ (X @ Wk)^T / sqrt(d)) @ (X @ Wv) @ Wo

        All weight matrices (Wq, Wk, Wv, Wo) are streamed from CDNA - never
        fully loaded. Verification follows agent's policy.

        Args:
            X: Input activations [batch, seq, d_model] or [seq, d_model].
            block_index: Transformer block index (0-31 for Mistral).
            manifest_path: Path to hybrid_manifest_v2.json.
            d_head: Head dimension for scaling (default: 128 for Mistral).

        Returns:
            (output [batch, seq, d_model], list of AccessReceipts)

        Note:
            This delegates to stream_attention_forward from stream_attention_layer.py
            but wraps results in agent-level AccessReceipts for policy tracking.
        """
        from helix_cdc.regrow.stream_attention_layer import stream_attention_forward as _stream_attn

        t0 = time.perf_counter()

        # Map agent verify_policy to stream_attention_layer format
        verify_policy = self.verify_policy

        # Call the core implementation
        output, layer_receipt, _ = _stream_attn(
            X=X,
            manifest_path=manifest_path,
            block_index=block_index,
            d_head=d_head,
            verify_policy=verify_policy,
            return_intermediates=False,
        )

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

        # Extract per-projection receipts and convert to AccessReceipts
        receipts = []

        for proj_name in ["q", "k", "v", "o"]:
            proj_receipt = layer_receipt.to_dict()["projections"].get(proj_name, {})
            if proj_receipt:
                receipt = AccessReceipt(
                    operation=f"stream_attention_{proj_name}_projection",
                    duration_ms=proj_receipt.get("duration_ms", 0),
                    cache_hit=False,  # Streaming doesn't use block cache
                    verified=proj_receipt.get("blocks_verified", 0) > 0,
                    blocks_accessed=len(proj_receipt.get("blocks_touched", [])),
                    blocks_verified=proj_receipt.get("blocks_verified", 0),
                    cache_size_mb=self._cache.current_size_bytes / (1024 * 1024),
                    extra={
                        "tensor_name": proj_receipt.get("tensor_name"),
                        "cdna_path": proj_receipt.get("cdna_path"),
                        "sidecar_applied": proj_receipt.get("sidecar_applied", False),
                    },
                )
                receipts.append(receipt)

        # Add summary receipt
        summary = AccessReceipt(
            operation="stream_attention_forward",
            duration_ms=duration_ms,
            cache_hit=False,
            verified=any(r.verified for r in receipts),
            blocks_accessed=sum(r.blocks_accessed for r in receipts),
            blocks_verified=sum(r.blocks_verified for r in receipts),
            cache_size_mb=self._cache.current_size_bytes / (1024 * 1024),
            extra={
                "block_index": block_index,
                "softmax_stable": layer_receipt.softmax_stats.get("numerically_stable", False),
                "memory_savings_factor": layer_receipt.memory_audit.get("savings_factor", 0),
                "full_weights_mb": layer_receipt.memory_audit.get("full_weights_mb", 0),
                "rss_delta_mb": layer_receipt.memory_audit.get("rss_delta_mb", 0),
            },
        )
        receipts.append(summary)

        return output, receipts


# CLI moved to helix_cdc/__main__.py (cdna-cache-stats)
# Run: python -m helix_cdc cdna-cache-stats <path.cdna2.hxz>
