#!/usr/bin/env python3
"""
tensor_index.py - GGUF Tensor Index with JSON Cache
====================================================

Provides O(1) tensor metadata lookup from GGUF files.
Caches the index as JSON for fast subsequent loads.

Features:
- O(1) tensor metadata lookup by name
- JSON caching (instant reload on second access)
- Selective tensor loading (load just one tensor, not all 291)
- SHA256 verification of source GGUF
- Pattern matching for tensor names

Work Order: WO-GGUF-TENSOR-INDEX
"""

from __future__ import annotations

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

import numpy as np

# Import from sibling scanner module
try:
    from helix_cdc.gguf.scanner import GGUFScanner, GGML_TYPE_MAP, HAVE_GGUF
    if HAVE_GGUF:
        from gguf.quants import dequantize
        from gguf.constants import GGMLQuantizationType
except ImportError:
    # Fallback for standalone testing
    from scanner import GGUFScanner, GGML_TYPE_MAP, HAVE_GGUF
    if HAVE_GGUF:
        from gguf.quants import dequantize
        from gguf.constants import GGMLQuantizationType


# Index schema version
INDEX_SCHEMA = "gguf_tensor_index_v1"


class GGUFTensorIndex:
    """
    Index for O(1) tensor metadata lookup from GGUF files.

    Caches index as JSON for fast subsequent loads.

    Usage:
        # From GGUF file (scans file, builds index)
        index = GGUFTensorIndex.from_gguf("model.gguf")

        # O(1) lookup
        info = index.get_tensor("blk.0.attn_k.weight")
        # {'name': ..., 'offset': 12345, 'byte_size': 4194304, 'shape': [1024, 4096], ...}

        # Load just this tensor without scanning entire file
        tensor = index.load_tensor("blk.0.attn_k.weight")

        # Cache for next time
        index.save_cache("model.gguf.index.json")

        # Load from cache (instant)
        index = GGUFTensorIndex.from_cache("model.gguf.index.json")

        # Pattern matching
        attn_tensors = index.list_tensors(pattern="attn_k")
    """

    def __init__(
        self,
        gguf_path: Path,
        tensors: Dict[str, Dict[str, Any]],
        data_offset: int,
        metadata: Dict[str, Any],
        gguf_sha256: Optional[str] = None,
    ):
        """
        Initialize tensor index (use from_gguf or from_cache instead).

        Args:
            gguf_path: Path to source GGUF file
            tensors: Dict mapping tensor name -> tensor info
            data_offset: Offset to tensor data section in GGUF
            metadata: GGUF file metadata
            gguf_sha256: SHA256 of GGUF file header (first 1MB)
        """
        self.gguf_path = Path(gguf_path)
        self.tensors = tensors
        self.data_offset = data_offset
        self.metadata = metadata
        self.gguf_sha256 = gguf_sha256

    @classmethod
    def from_gguf(cls, gguf_path: Union[str, Path], compute_sha256: bool = True) -> "GGUFTensorIndex":
        """
        Build index from GGUF file.

        Scans the GGUF file once to extract all tensor metadata.

        Args:
            gguf_path: Path to GGUF file
            compute_sha256: Whether to compute SHA256 of file header (first 1MB)

        Returns:
            GGUFTensorIndex with O(1) lookup
        """
        gguf_path = Path(gguf_path)

        # Use cached scanner to avoid repeated scanning (WO-TOKENIZER-BRIDGE-01.1)
        from helix_cdc.gguf.cache import get_scanner
        scanner = get_scanner(gguf_path, quiet=True)

        # Build tensor dict for O(1) lookup
        tensors: Dict[str, Dict[str, Any]] = {}
        for t in scanner.tensors:
            tensors[t["name"]] = {
                "name": t["name"],
                "offset": t["offset"],
                "byte_size": t["byte_size"],
                "shape": list(t["dims"]),
                "dtype_int": t["dtype_int"],
                "dtype_name": t["dtype_name"],
                "n_elements": t["n_elements"],
                "block_size": t["block_size"],
                "type_size": t["type_size"],
            }

        # Compute SHA256 of first 1MB for integrity check
        gguf_sha256 = None
        if compute_sha256:
            h = hashlib.sha256()
            with open(gguf_path, "rb") as f:
                h.update(f.read(1024 * 1024))
            gguf_sha256 = h.hexdigest()

        return cls(
            gguf_path=gguf_path,
            tensors=tensors,
            data_offset=scanner.data_offset,
            metadata=scanner.metadata,
            gguf_sha256=gguf_sha256,
        )

    @classmethod
    def from_cache(cls, cache_path: Union[str, Path], verify: bool = True) -> "GGUFTensorIndex":
        """
        Load index from cached JSON file.

        Args:
            cache_path: Path to .index.json file
            verify: Whether to verify GGUF file still exists and matches SHA256

        Returns:
            GGUFTensorIndex with O(1) lookup
        """
        cache_path = Path(cache_path)
        with open(cache_path, "r") as f:
            data = json.load(f)

        if data.get("schema") != INDEX_SCHEMA:
            raise ValueError(f"Unknown index schema: {data.get('schema')}")

        gguf_path = Path(data["gguf_path"])
        gguf_sha256 = data.get("gguf_sha256")

        if verify:
            if not gguf_path.exists():
                raise FileNotFoundError(f"GGUF file not found: {gguf_path}")

            if gguf_sha256:
                h = hashlib.sha256()
                with open(gguf_path, "rb") as f:
                    h.update(f.read(1024 * 1024))
                actual_sha256 = h.hexdigest()
                if actual_sha256 != gguf_sha256:
                    raise ValueError(
                        f"GGUF file changed since index creation. "
                        f"Expected SHA256: {gguf_sha256[:16]}..., "
                        f"actual: {actual_sha256[:16]}..."
                    )

        return cls(
            gguf_path=gguf_path,
            tensors=data["tensors"],
            data_offset=data["data_offset"],
            metadata=data.get("metadata", {}),
            gguf_sha256=gguf_sha256,
        )

    def save_cache(self, cache_path: Union[str, Path]) -> Dict[str, Any]:
        """
        Save index to JSON cache file.

        Args:
            cache_path: Where to write the cache file

        Returns:
            Receipt with save stats
        """
        cache_path = Path(cache_path)

        data = {
            "schema": INDEX_SCHEMA,
            "gguf_path": str(self.gguf_path.absolute()),
            "gguf_sha256": self.gguf_sha256,
            "data_offset": self.data_offset,
            "tensor_count": len(self.tensors),
            "total_bytes": self.total_bytes,
            "created_utc": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
            "metadata": self.metadata,
            "tensors": self.tensors,
        }

        cache_path.parent.mkdir(parents=True, exist_ok=True)
        with open(cache_path, "w") as f:
            json.dump(data, f, indent=2)

        return {
            "schema": "gguf_tensor_index_save_v1",
            "cache_path": str(cache_path),
            "tensor_count": len(self.tensors),
            "cache_bytes": cache_path.stat().st_size,
            "timestamp": data["created_utc"],
        }

    def get_tensor(self, name: str) -> Optional[Dict[str, Any]]:
        """
        O(1) lookup of tensor metadata by name.

        Args:
            name: Tensor name (e.g., "blk.0.attn_k.weight")

        Returns:
            Tensor info dict or None if not found
        """
        return self.tensors.get(name)

    def load_tensor(self, name: str) -> np.ndarray:
        """
        Load a single tensor from GGUF without scanning entire file.

        Args:
            name: Tensor name

        Returns:
            Dequantized tensor as float32 numpy array
        """
        info = self.get_tensor(name)
        if info is None:
            raise KeyError(f"Tensor not found: {name}")

        # Read raw bytes
        with open(self.gguf_path, "rb") as f:
            f.seek(self.data_offset + info["offset"])
            raw = f.read(info["byte_size"])

        data = np.frombuffer(raw, dtype=np.uint8)

        # Get quantization type
        dtype_int = info["dtype_int"]
        dtype_info = GGML_TYPE_MAP.get(dtype_int)

        if dtype_info is None:
            raise ValueError(f"Unknown dtype_int: {dtype_int}")

        dtype_name, qtype, (block_size, type_size) = dtype_info

        # Handle unquantized formats
        if dtype_name == "F32":
            tensor = np.frombuffer(raw, dtype=np.float32)
            return tensor.reshape(info["shape"]).astype(np.float32)
        elif dtype_name == "F16":
            tensor = np.frombuffer(raw, dtype=np.float16)
            return tensor.reshape(info["shape"]).astype(np.float32)

        # Quantized: use gguf dequantize
        if not HAVE_GGUF:
            raise ImportError("gguf package required for quantized tensors")

        if qtype is None:
            raise ValueError(f"Dequantization not supported for {dtype_name}")

        data = data.reshape(-1, type_size)
        dequant = dequantize(data, qtype)
        tensor = dequant.reshape(info["shape"])

        return tensor.astype(np.float32)

    def load_tensor_raw(self, name: str) -> bytes:
        """
        Load raw (quantized) bytes for a tensor.

        Args:
            name: Tensor name

        Returns:
            Raw bytes from GGUF file
        """
        info = self.get_tensor(name)
        if info is None:
            raise KeyError(f"Tensor not found: {name}")

        with open(self.gguf_path, "rb") as f:
            f.seek(self.data_offset + info["offset"])
            return f.read(info["byte_size"])

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

        Args:
            pattern: Optional regex or substring pattern to filter

        Returns:
            List of matching tensor names
        """
        names = list(self.tensors.keys())

        if pattern is None:
            return sorted(names)

        # Try as regex first, fall back to substring
        try:
            regex = re.compile(pattern)
            return sorted([n for n in names if regex.search(n)])
        except re.error:
            # Substring match
            return sorted([n for n in names if pattern in n])

    @property
    def tensor_count(self) -> int:
        """Number of tensors in index."""
        return len(self.tensors)

    @property
    def total_bytes(self) -> int:
        """Total bytes of all tensors (raw, quantized)."""
        return sum(t["byte_size"] for t in self.tensors.values())

    @property
    def total_elements(self) -> int:
        """Total number of elements across all tensors."""
        return sum(t["n_elements"] for t in self.tensors.values())

    def summary(self) -> Dict[str, Any]:
        """Get summary statistics for the index."""
        dtype_counts: Dict[str, int] = {}
        for t in self.tensors.values():
            dtype = t["dtype_name"]
            dtype_counts[dtype] = dtype_counts.get(dtype, 0) + 1

        return {
            "gguf_path": str(self.gguf_path),
            "tensor_count": self.tensor_count,
            "total_bytes": self.total_bytes,
            "total_elements": self.total_elements,
            "data_offset": self.data_offset,
            "dtype_counts": dtype_counts,
            "gguf_sha256": self.gguf_sha256[:16] + "..." if self.gguf_sha256 else None,
        }

    def __repr__(self) -> str:
        return (
            f"GGUFTensorIndex(gguf={self.gguf_path.name}, "
            f"tensors={self.tensor_count}, "
            f"bytes={self.total_bytes:,})"
        )


# Convenience function
def build_tensor_index(
    gguf_path: Union[str, Path],
    cache_path: Optional[Union[str, Path]] = None,
) -> GGUFTensorIndex:
    """
    Build or load tensor index with automatic caching.

    Args:
        gguf_path: Path to GGUF file
        cache_path: Optional cache path (default: {gguf_path}.index.json)

    Returns:
        GGUFTensorIndex
    """
    gguf_path = Path(gguf_path)

    if cache_path is None:
        cache_path = gguf_path.with_suffix(gguf_path.suffix + ".index.json")
    else:
        cache_path = Path(cache_path)

    # Try loading from cache
    if cache_path.exists():
        try:
            return GGUFTensorIndex.from_cache(cache_path, verify=True)
        except (ValueError, FileNotFoundError):
            # Cache invalid, rebuild
            pass

    # Build from GGUF
    index = GGUFTensorIndex.from_gguf(gguf_path)
    index.save_cache(cache_path)

    return index


if __name__ == "__main__":
    import sys

    print("=== GGUF Tensor Index Test ===")
    print()

    if len(sys.argv) < 2:
        print("Usage: python tensor_index.py <gguf_file>")
        print()
        print("Example:")
        print("  python tensor_index.py model.gguf")
        sys.exit(1)

    gguf_path = Path(sys.argv[1])

    print(f"Loading: {gguf_path}")
    print()

    # Build index
    index = GGUFTensorIndex.from_gguf(gguf_path)

    # Summary
    summary = index.summary()
    print("Summary:")
    for k, v in summary.items():
        if k == "total_bytes":
            print(f"  {k}: {v:,} ({v/1024/1024/1024:.2f} GB)")
        elif k == "total_elements":
            print(f"  {k}: {v:,}")
        else:
            print(f"  {k}: {v}")
    print()

    # O(1) lookup test
    if index.tensor_count > 0:
        first_name = index.list_tensors()[0]
        info = index.get_tensor(first_name)
        print(f"O(1) lookup test: {first_name}")
        print(f"  offset: {info['offset']}")
        print(f"  byte_size: {info['byte_size']:,}")
        print(f"  shape: {info['shape']}")
        print(f"  dtype: {info['dtype_name']}")
        print()

    # Cache roundtrip test
    cache_path = gguf_path.with_suffix(".gguf.index.json")
    print(f"Saving cache to: {cache_path}")
    receipt = index.save_cache(cache_path)
    print(f"  cache_bytes: {receipt['cache_bytes']:,}")
    print()

    print("Loading from cache...")
    index2 = GGUFTensorIndex.from_cache(cache_path)
    print(f"  tensors: {index2.tensor_count}")
    print(f"  match: {index.tensor_count == index2.tensor_count}")
    print()

    print("DONE")
