#!/usr/bin/env python3
"""
tensor_accessor.py - Unified Tensor Access Layer
=================================================

Provides a single API for accessing tensors from multiple formats:
- Direct GGUF: Load tensor from GGUF file
- CDNA v1: Load from .cdna.hxz (full tensor, zlib)
- CDNA v2: Load from .cdna2.hxz (row-range support, brotli)
- Hybrid: CDNA + HXZO sidecar for outlier correction

Features:
- Format auto-detection (GGUF, CDNA v1, CDNA v2)
- Manifest-based access (hybrid_manifest_v2 format)
- Row-range streaming with receipts
- Sidecar integration for fidelity recovery

Work Order: WO-TENSOR-ACCESSOR
"""

from __future__ import annotations

import json
import re
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np


class TensorAccessor:
    """
    Unified tensor access with format auto-detection.

    Supports:
      - Direct GGUF: Load tensor from GGUF file
      - CDNA v1: Load from .cdna.hxz (full tensor)
      - CDNA v2: Load from .cdna2.hxz (row-range support)
      - Hybrid: CDNA + HXZO sidecar

    Usage:
        # From manifest
        accessor = TensorAccessor.from_manifest("hybrid_manifest_v2.json")

        # Load full tensor
        weights = accessor.load_tensor("blk.0.attn_k.weight")

        # Load row range (only decompresses needed blocks)
        rows_100_200, receipt = accessor.load_rows(
            "blk.0.attn_k.weight",
            start_row=100,
            end_row=200
        )

        # With sidecar applied (fidelity recovery)
        rows, receipt = accessor.load_rows_with_sidecar(
            "blk.0.attn_k.weight",
            start_row=100,
            end_row=200
        )

        # From GGUF directly
        accessor = TensorAccessor.from_gguf("model.gguf")
        tensor = accessor.load_tensor("blk.0.attn_k.weight")
    """

    def __init__(
        self,
        manifest: Optional[Dict[str, Any]] = None,
        manifest_path: Optional[Path] = None,
        gguf_index: Optional[Any] = None,  # GGUFTensorIndex
        base_path: Optional[Path] = None,
    ):
        """
        Initialize accessor (use from_manifest or from_gguf instead).

        Args:
            manifest: Parsed manifest dict
            manifest_path: Path to manifest file
            gguf_index: GGUFTensorIndex for direct GGUF access
            base_path: Base path for resolving relative paths
        """
        self.manifest = manifest
        self.manifest_path = manifest_path
        self.gguf_index = gguf_index
        self.base_path = base_path or Path.cwd()

        # Build tensor lookup from manifest
        self._tensor_info: Dict[str, Dict[str, Any]] = {}
        if manifest and "shards" in manifest:
            for shard in manifest["shards"]:
                name = shard["tensor_name"]
                self._tensor_info[name] = shard

        # Cache for CDNA readers
        self._cdna_readers: Dict[str, Any] = {}

    @classmethod
    def from_manifest(cls, manifest_path: Union[str, Path]) -> "TensorAccessor":
        """
        Create accessor from hybrid manifest file.

        Args:
            manifest_path: Path to hybrid_manifest_v2.json or similar

        Returns:
            TensorAccessor configured for manifest-based access
        """
        manifest_path = Path(manifest_path).absolute()
        with open(manifest_path, "r") as f:
            manifest = json.load(f)

        # Determine base_path from manifest structure
        # If paths in manifest start with common directories (seeds/, tmp/),
        # base_path should be the repo root, not the manifest parent
        base_path = manifest_path.parent

        # Check if we need to adjust base_path
        # Look at first shard path to determine path style
        if manifest.get("shards"):
            first_path = manifest["shards"][0].get("path", "")
            # If path already includes parent dir name (e.g., "seeds/..."),
            # walk up to find the correct base
            if first_path.startswith("seeds/") or first_path.startswith("tmp/"):
                # Path is relative to repo root, find it by walking up
                # from manifest parent until we find a valid path
                test_base = manifest_path.parent.parent
                test_path = test_base / first_path
                if test_path.exists():
                    base_path = test_base

        return cls(
            manifest=manifest,
            manifest_path=manifest_path,
            base_path=base_path,
        )

    @classmethod
    def from_gguf(cls, gguf_path: Union[str, Path]) -> "TensorAccessor":
        """
        Create accessor for direct GGUF access.

        Args:
            gguf_path: Path to GGUF file

        Returns:
            TensorAccessor configured for GGUF-based access
        """
        from helix_cdc.gguf import GGUFTensorIndex

        gguf_path = Path(gguf_path)
        index = GGUFTensorIndex.from_gguf(gguf_path)

        return cls(
            gguf_index=index,
            base_path=gguf_path.parent,
        )

    def _resolve_path(self, path_str: str) -> Path:
        """Resolve a path relative to base_path."""
        path = Path(path_str)
        if path.is_absolute():
            return path
        return self.base_path / path

    def _get_cdna_reader(self, cdna_path: Path) -> Any:
        """Get or create CDNA reader for a path."""
        path_str = str(cdna_path)
        if path_str not in self._cdna_readers:
            from helix_cdc.regrow.cdna_stream_v2 import load_cdna_auto
            self._cdna_readers[path_str] = load_cdna_auto(cdna_path)
        return self._cdna_readers[path_str]

    def get_tensor_info(self, name: str) -> Optional[Dict[str, Any]]:
        """
        Get metadata for a tensor.

        Args:
            name: Tensor name

        Returns:
            Info dict or None if not found
        """
        # Check manifest first
        if name in self._tensor_info:
            return self._tensor_info[name]

        # Fall back to GGUF index
        if self.gguf_index is not None:
            return self.gguf_index.get_tensor(name)

        return None

    def list_tensors(self, pattern: Optional[str] = None) -> List[str]:
        """
        List all available tensor names.

        Args:
            pattern: Optional regex or substring pattern

        Returns:
            List of tensor names
        """
        names = set()

        # From manifest
        if self.manifest and "shards" in self.manifest:
            for shard in self.manifest["shards"]:
                names.add(shard["tensor_name"])

        # From GGUF index
        if self.gguf_index is not None:
            names.update(self.gguf_index.list_tensors())

        names_list = sorted(names)

        if pattern is None:
            return names_list

        # Filter by pattern
        try:
            regex = re.compile(pattern)
            return [n for n in names_list if regex.search(n)]
        except re.error:
            return [n for n in names_list if pattern in n]

    def load_tensor(self, name: str) -> np.ndarray:
        """
        Load a full tensor.

        Automatically selects the best source (manifest → GGUF).

        Args:
            name: Tensor name

        Returns:
            Tensor as float32 numpy array
        """
        # Check manifest first
        if name in self._tensor_info:
            info = self._tensor_info[name]
            cdna_path = self._resolve_path(info["path"])

            reader = self._get_cdna_reader(cdna_path)

            # Get full indices based on reader type
            if hasattr(reader, "get_all_indices"):
                # CDNA v2
                indices = reader.get_all_indices()
            elif hasattr(reader, "_ensure_indices_blob"):
                # CDNA v1 - get indices from blob
                blob = reader._ensure_indices_blob()
                indices = np.frombuffer(blob, dtype=np.uint8).reshape(reader.rows, reader.cols)
            else:
                raise ValueError(f"Unknown CDNA reader type for {name}")

            # Dequantize: W[i,j] = codebook[indices[i,j]]
            weights = reader.codebook[indices]

            # Apply sidecar if available
            if info.get("outlier_sidecar_path"):
                sidecar_dir = self.manifest.get("sidecar_dir", "")
                sidecar_path = self._resolve_path(
                    Path(sidecar_dir) / info["outlier_sidecar_path"]
                )
                if sidecar_path.exists():
                    weights = self._apply_sidecar(weights, sidecar_path, reader.shape)

            return weights.astype(np.float32)

        # Fall back to GGUF
        if self.gguf_index is not None:
            return self.gguf_index.load_tensor(name)

        raise KeyError(f"Tensor not found: {name}")

    def _apply_sidecar(
        self,
        weights: np.ndarray,
        sidecar_path: Path,
        shape: Tuple[int, int],
    ) -> np.ndarray:
        """Apply HXZO sidecar corrections to weights."""
        from helix_cdc.regrow.hxzo_sidecar import read_outlier_sidecar

        positions, values, meta = read_outlier_sidecar(str(sidecar_path))

        weights_flat = weights.flatten()
        for i in range(len(positions)):
            pos = positions[i]
            if pos < len(weights_flat):
                weights_flat[pos] = float(values[i])

        return weights_flat.reshape(shape)

    def _apply_sidecar_for_rows(
        self,
        weights: np.ndarray,
        sidecar_path: Path,
        full_shape: Tuple[int, int],
        start_row: int,
        end_row: int,
    ) -> np.ndarray:
        """Apply HXZO sidecar corrections for a specific row range."""
        from helix_cdc.regrow.hxzo_sidecar import read_outlier_sidecar

        positions, values, meta = read_outlier_sidecar(str(sidecar_path))

        rows, cols = full_shape
        corrections = 0

        for i in range(len(positions)):
            pos = positions[i]
            row = pos // cols
            col = pos % cols

            # Check if position is in our row range
            if start_row <= row < end_row:
                local_row = row - start_row
                weights[local_row, col] = float(values[i])
                corrections += 1

        return weights

    def load_rows(
        self,
        name: str,
        start_row: int,
        end_row: int,
        apply_sidecar: bool = False,
        decode_mode: str = "accurate",  # WO-CDNA2-FIDELITY-ROUTER-01
        min_cosine: Optional[float] = None,  # WO-CDNA2-FIDELITY-ROUTER-02
    ) -> Tuple[np.ndarray, Dict[str, Any]]:
        """
        Load a range of rows from a tensor.

        Only decompresses blocks containing the requested rows (CDNA v2).
        For CDNA v1 or GGUF, falls back to full load + slice.

        Args:
            name: Tensor name
            start_row: First row (inclusive)
            end_row: Last row (exclusive)
            apply_sidecar: If True, apply outlier sidecar corrections (fidelity recovery).
                           If False (default), return raw quantized weights.
            decode_mode: "accurate" or "fast" (convenience preset for min_cosine)
                - "accurate": min_cosine default = 0.999 (court-grade)
                - "fast": min_cosine default = 0.998 (speed-leaning)
            min_cosine: Explicit quality contract (WO-CDNA2-FIDELITY-ROUTER-02)
                - If provided, overrides decode_mode
                - Sidecar applied if header.cosine_no_sidecar < min_cosine

        Returns:
            (weights [end_row - start_row, K] float32, receipt)

        Note:
            Raw vs Corrected semantics:
            - apply_sidecar=False: Returns weights as dequantized from codebook indices.
              These have quantization error but are faster to compute.
            - apply_sidecar=True: Returns weights with outlier corrections applied.
              These are fidelity-corrected and match the original tensor more closely.
        """
        # Check manifest first
        if name in self._tensor_info:
            info = self._tensor_info[name]
            cdna_path = self._resolve_path(info["path"])

            # Resolve sidecar path if apply_sidecar is requested
            sidecar_path = None
            if apply_sidecar and info.get("outlier_sidecar_path"):
                sidecar_dir = self.manifest.get("sidecar_dir", "")
                sidecar_path = self._resolve_path(
                    Path(sidecar_dir) / info["outlier_sidecar_path"]
                )
                if not sidecar_path.exists():
                    sidecar_path = None

            reader = self._get_cdna_reader(cdna_path)

            # Use stream_rows_dequant if available (CDNA v2)
            if hasattr(reader, "stream_rows_dequant"):
                weights, receipt = reader.stream_rows_dequant(
                    start_row, end_row,
                    sidecar_path=sidecar_path,
                    decode_mode=decode_mode,  # WO-CDNA2-FIDELITY-ROUTER-01
                    min_cosine=min_cosine,  # WO-CDNA2-FIDELITY-ROUTER-02
                    emit_receipt=True,
                )
                return weights, receipt

            # CDNA v1: Get indices from blob and slice
            if hasattr(reader, "_ensure_indices_blob"):
                blob = reader._ensure_indices_blob()
                all_indices = np.frombuffer(blob, dtype=np.uint8).reshape(reader.rows, reader.cols)
                row_indices = all_indices[start_row:end_row]
                weights = reader.codebook[row_indices]

                sidecar_applied = False
                sidecar_corrections = 0

                # Apply sidecar if requested and available
                if sidecar_path is not None:
                    # Need to apply corrections for rows in our range
                    weights = self._apply_sidecar_for_rows(
                        weights, sidecar_path, reader.shape,
                        start_row, end_row
                    )
                    sidecar_applied = True
                    # Count corrections (simplified - actual count tracked in _apply_sidecar_for_rows)

                receipt = {
                    "schema": "tensor_accessor_row_access_v1",
                    "tensor_name": name,
                    "row_range": [start_row, end_row],
                    "rows_loaded": end_row - start_row,
                    "rows_total": reader.rows,
                    "method": "cdna_v1_full_decompress_slice",
                    "sidecar_applied": sidecar_applied,
                    "timestamp_utc": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
                }

                return weights.astype(np.float32), receipt

            # Unknown format - fall back to full load
            full_tensor = self.load_tensor(name)
            weights = full_tensor[start_row:end_row]

            receipt = {
                "schema": "tensor_accessor_row_access_v1",
                "tensor_name": name,
                "row_range": [start_row, end_row],
                "rows_loaded": end_row - start_row,
                "method": "full_load_slice",
                "sidecar_applied": False,
                "timestamp_utc": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
            }

            return weights, receipt

        # Fall back to GGUF
        if self.gguf_index is not None:
            full_tensor = self.gguf_index.load_tensor(name)
            weights = full_tensor[start_row:end_row]

            receipt = {
                "schema": "tensor_accessor_row_access_v1",
                "tensor_name": name,
                "row_range": [start_row, end_row],
                "rows_loaded": end_row - start_row,
                "method": "gguf_load_slice",
                "sidecar_applied": False,  # GGUF is already full precision
                "timestamp_utc": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
            }

            return weights, receipt

        raise KeyError(f"Tensor not found: {name}")

    def load_rows_with_sidecar(
        self,
        name: str,
        start_row: int,
        end_row: int,
    ) -> Tuple[np.ndarray, Dict[str, Any]]:
        """
        Load rows with sidecar outlier correction applied.

        DEPRECATED: Use load_rows(name, start, end, apply_sidecar=True) instead.
        This method is kept for backward compatibility.

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

        Returns:
            (weights [end_row - start_row, K] float32, receipt)
        """
        return self.load_rows(name, start_row, end_row, apply_sidecar=True)

    def supports_row_range(self, name: str) -> bool:
        """
        Check if a tensor supports efficient row-range access.

        Args:
            name: Tensor name

        Returns:
            True if row-range access is efficient (CDNA v2 with stream_rows)
        """
        if name in self._tensor_info:
            info = self._tensor_info[name]
            cdna_path = self._resolve_path(info["path"])

            try:
                reader = self._get_cdna_reader(cdna_path)
                # Only CDNA v2 has efficient row-range via stream_rows_dequant
                return hasattr(reader, "stream_rows_dequant")
            except Exception:
                return False

        return False

    def close(self) -> None:
        """Close all cached readers."""
        self._cdna_readers.clear()

    def __repr__(self) -> str:
        tensor_count = len(self._tensor_info)
        if self.gguf_index:
            tensor_count = max(tensor_count, self.gguf_index.tensor_count)

        source = "manifest" if self.manifest else "gguf"
        if self.manifest and self.gguf_index:
            source = "manifest+gguf"

        return f"TensorAccessor(source={source}, tensors={tensor_count})"


def create_accessor(
    source: Union[str, Path],
) -> TensorAccessor:
    """
    Create a TensorAccessor from a source file.

    Auto-detects source type:
    - .json → manifest
    - .gguf → GGUF

    Args:
        source: Path to manifest or GGUF file

    Returns:
        TensorAccessor
    """
    source = Path(source)

    if source.suffix == ".json":
        return TensorAccessor.from_manifest(source)
    elif source.suffix == ".gguf":
        return TensorAccessor.from_gguf(source)
    else:
        raise ValueError(f"Unknown source type: {source.suffix}")


if __name__ == "__main__":
    import sys

    print("=== Tensor Accessor Test ===")
    print()

    if len(sys.argv) < 2:
        print("Usage: python tensor_accessor.py <manifest.json|model.gguf>")
        sys.exit(1)

    source = Path(sys.argv[1])
    print(f"Loading: {source}")
    print()

    accessor = create_accessor(source)
    print(f"Accessor: {accessor}")
    print()

    tensors = accessor.list_tensors()
    print(f"Tensors: {len(tensors)}")
    for t in tensors[:5]:
        print(f"  - {t}")
    if len(tensors) > 5:
        print(f"  ... and {len(tensors) - 5} more")
    print()

    if tensors:
        first = tensors[0]
        print(f"Loading first tensor: {first}")
        weights = accessor.load_tensor(first)
        print(f"  Shape: {weights.shape}")
        print(f"  Dtype: {weights.dtype}")
        print(f"  Min: {weights.min():.6f}")
        print(f"  Max: {weights.max():.6f}")
        print()

        if accessor.supports_row_range(first):
            print(f"Testing row-range access (rows 0-10):")
            rows, receipt = accessor.load_rows(first, 0, 10)
            print(f"  Shape: {rows.shape}")
            print(f"  Receipt schema: {receipt.get('schema')}")
            if "savings_factor" in receipt:
                print(f"  Savings factor: {receipt['savings_factor']:.1f}x")
        print()

    print("DONE")
