"""
CDNA v2 Reader - Wrapper around helix-substrate's CDNAv2Reader

Returns tensors in the same orientation as load_weights_from_gguf() expects.

Interface:
    CDNATensor - single tensor shard
    CDNAModel - collection of tensors for a model
"""

import json
import numpy as np
from pathlib import Path
from typing import Iterator, Tuple, Optional, Dict

# Import the real implementation from helix-substrate
from helix_substrate.cdna_reader import CDNAv2Reader as _CDNAv2Reader


class CDNATensor:
    """
    Single CDNA tensor shard reader.

    Wraps helix-substrate's CDNAv2Reader.
    Returns tensor in the same orientation as GGUF loader.
    """

    def __init__(self, path: str):
        self._path = Path(path)
        self._reader = _CDNAv2Reader(path)
        self._data: Optional[np.ndarray] = None

    @property
    def name(self) -> str:
        """Tensor name in GGUF format (e.g., 'blk.0.attn_q.weight')."""
        return self._reader.tensor_name

    @property
    def shape(self) -> Tuple[int, ...]:
        """Return tensor shape (same as GGUF orientation)."""
        return self._reader.shape

    @property
    def codebook(self) -> np.ndarray:
        """Return codebook for this tensor."""
        return self._reader.codebook

    def load_full(self) -> np.ndarray:
        """Load and return the full dequantized tensor."""
        if self._data is None:
            indices = self._reader.get_all_indices()
            self._data = self._reader.codebook[indices]
        return self._data.copy()

    def iter_row_blocks(self, block_rows: int) -> Iterator[Tuple[int, np.ndarray]]:
        """
        Iterate over row blocks of the dequantized tensor.

        For a tensor with shape [rows, cols], yields:
            (start_row, block) where block.shape = [min(block_rows, remaining), cols]

        This is memory-efficient for streaming matmul.
        """
        n_rows = self.shape[0]
        for start in range(0, n_rows, block_rows):
            end = min(start + block_rows, n_rows)
            indices, _ = self._reader.stream_rows(start, end, emit_receipt=False)
            weights = self._reader.codebook[indices]
            yield start, weights


def _parse_tensor_name(filename: str) -> str:
    """
    Parse GGUF tensor name from filename.

    Examples:
        blk_0_attn_q_weight.cdna2.hxz -> blk.0.attn_q.weight
        token_embd_weight.cdna2.hxz -> token_embd.weight
        output_weight.cdna2.hxz -> output.weight
    """
    # Remove extensions
    stem = filename.replace('.cdna2.hxz', '').replace('.npy', '')
    parts = stem.split('_')

    if parts[0] == 'blk' and len(parts) > 2:
        # blk_0_attn_q_weight -> blk.0.attn_q.weight
        block_num = parts[1]
        middle = '_'.join(parts[2:-1])
        last = parts[-1]
        return f"blk.{block_num}.{middle}.{last}"
    else:
        # token_embd_weight -> token_embd.weight
        middle = '_'.join(parts[:-1])
        last = parts[-1]
        return f"{middle}.{last}"


class CDNAModel:
    """
    Collection of CDNA tensors for a model.

    Loads manifest and provides access to individual tensors.
    """

    def __init__(self, cdna_dir: str):
        self._dir = Path(cdna_dir)
        self._manifest: Optional[Dict] = None
        self._tensors: Dict[str, CDNATensor] = {}
        self._npy_cache: Dict[str, np.ndarray] = {}

        # Load manifest
        manifest_path = self._dir / 'manifest.json'
        if manifest_path.exists():
            self._manifest = json.loads(manifest_path.read_text())

        # Index available tensors
        self._index_tensors()

    def _index_tensors(self) -> None:
        """Build index of available tensors."""
        # Index .cdna2.hxz files
        for hxz_path in self._dir.glob('*.cdna2.hxz'):
            tensor = CDNATensor(str(hxz_path))
            self._tensors[tensor.name] = tensor

        # Index .npy files (1D tensors like norms)
        for npy_path in self._dir.glob('*.npy'):
            name = _parse_tensor_name(npy_path.name)
            self._npy_cache[name] = np.load(npy_path)

    @property
    def manifest(self) -> Optional[Dict]:
        """Return the manifest metadata."""
        return self._manifest

    def tensor_names(self) -> list:
        """List all available tensor names."""
        cdna_names = list(self._tensors.keys())
        npy_names = list(self._npy_cache.keys())
        return sorted(cdna_names + npy_names)

    def has(self, tensor_name: str) -> bool:
        """Check if a tensor exists."""
        return tensor_name in self._tensors or tensor_name in self._npy_cache

    def get(self, tensor_name: str) -> CDNATensor:
        """
        Get a CDNATensor by name.

        For 1D tensors (.npy), raises KeyError - use get_npy() instead.
        """
        if tensor_name not in self._tensors:
            if tensor_name in self._npy_cache:
                raise KeyError(f"{tensor_name} is a 1D tensor, use get_npy()")
            raise KeyError(f"Tensor not found: {tensor_name}")
        return self._tensors[tensor_name]

    def get_npy(self, tensor_name: str) -> np.ndarray:
        """Get a 1D tensor (norm weights) directly."""
        if tensor_name not in self._npy_cache:
            raise KeyError(f"1D tensor not found: {tensor_name}")
        return self._npy_cache[tensor_name]

    def is_npy(self, tensor_name: str) -> bool:
        """Check if tensor is a 1D npy tensor."""
        return tensor_name in self._npy_cache

    def load_tensor(self, tensor_name: str, apply_sidecar: bool = True) -> np.ndarray:
        """
        Load a tensor by name, handling both CDNA and npy formats.

        Args:
            tensor_name: Name of tensor to load
            apply_sidecar: If True, apply outlier corrections from sidecar

        Returns tensor in same orientation as GGUF loader.
        """
        if tensor_name in self._npy_cache:
            return self._npy_cache[tensor_name]
        if tensor_name in self._tensors:
            weights = self._tensors[tensor_name].load_full()

            if apply_sidecar:
                weights = self._apply_sidecar(tensor_name, weights)

            return weights
        raise KeyError(f"Tensor not found: {tensor_name}")

    def _apply_sidecar(self, tensor_name: str, weights: np.ndarray) -> np.ndarray:
        """Apply outlier corrections from sidecar file if it exists."""
        # Look for sidecar: blk.2.ffn_down.weight -> blk_2_ffn_down_weight.sidecar.npz
        safe_name = tensor_name.replace('.', '_')
        sidecar_path = self._dir / f"{safe_name}.sidecar.npz"

        if not sidecar_path.exists():
            return weights

        # Load sidecar
        sidecar = np.load(sidecar_path, allow_pickle=True)
        rows = sidecar['rows']
        cols = sidecar['cols']
        values = sidecar['values']

        # Apply corrections
        weights = weights.copy()  # Don't modify cached data
        weights[rows, cols] = values

        return weights
