#!/usr/bin/env python3
"""
tokenizer_bridge.py - MistralTokenizerBridge for Native CDNA v2 Chat

WO-TOKENIZER-BRIDGE-01: Bridge between text and native CDNA v2 streaming.

This module provides the "front door" for native CDNA v2 streaming chat:
- Text → token IDs (encode)
- Token IDs → embeddings (lookup_embeddings)
- Hidden states → logits (project_to_logits)
- Token IDs → text (decode)

Architecture:
    Text "Hello"
        ↓ MistralTokenizerBridge.encode()
    Token IDs [1, 15043, ...]
        ↓ MistralTokenizerBridge.lookup_embeddings()
    Embeddings X [seq, 4096]
        ↓ stream_multi_block_forward(X, block_indices)
    Hidden states [seq, 4096]
        ↓ MistralTokenizerBridge.project_to_logits()
    Logits [32000]
        ↓ argmax (greedy)
    Next token → append → repeat
        ↓ MistralTokenizerBridge.decode()
    Generated text

Memory Profile:
    - llama.cpp (vocab only): ~200MB (n_ctx=128, n_gpu_layers=0)
    - token_embd.weight: 128MB (lazy-loaded)
    - output.weight: 128MB or tied (lazy-loaded)
    - CDNA blocks: 8.75MB/block (streaming)
"""

from __future__ import annotations

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

import numpy as np

# Try to import llama_cpp for tokenization
try:
    from llama_cpp import Llama
    HAVE_LLAMA_CPP = True
except ImportError:
    HAVE_LLAMA_CPP = False
    Llama = None


@dataclass
class TokenizerBridgeReceipt:
    """
    Receipt for tokenizer bridge operations.

    Every operation emits a receipt with fingerprints for claim hygiene.
    """
    schema: str = "tokenizer_bridge_receipt_v1"
    work_order: str = "WO-TOKENIZER-BRIDGE-01"

    # Model info
    vocab_size: int = 0
    vocab_hash: str = ""  # SHA256[:16] of first 1000 tokens
    gguf_path: str = ""

    # Operation
    operation: str = ""  # "encode", "decode", "embed_lookup", "lm_head"

    # Token info
    token_ids: Optional[List[int]] = None
    token_count: int = 0
    text_length: int = 0

    # Embedding info (for embed_lookup)
    embedding_shape: Optional[List[int]] = None
    embedding_hash: str = ""  # SHA256[:16] of output

    # Logits info (for lm_head)
    logits_shape: Optional[List[int]] = None
    logits_hash: str = ""
    top_token_id: Optional[int] = None
    top_token_logit: Optional[float] = None

    # Timing
    duration_ms: float = 0.0
    timestamp: str = field(
        default_factory=lambda: datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
    )

    # Status
    status: str = "PASS"
    error_message: Optional[str] = None

    def to_dict(self) -> Dict[str, Any]:
        """Convert to JSON-serializable dictionary."""
        d = {
            "schema": self.schema,
            "work_order": self.work_order,
            "vocab_size": self.vocab_size,
            "vocab_hash": self.vocab_hash,
            "gguf_path": self.gguf_path,
            "operation": self.operation,
            "token_count": self.token_count,
            "text_length": self.text_length,
            "duration_ms": round(self.duration_ms, 3),
            "timestamp": self.timestamp,
            "status": self.status,
        }

        if self.token_ids is not None:
            d["token_ids"] = self.token_ids

        if self.embedding_shape is not None:
            d["embedding_shape"] = self.embedding_shape
            d["embedding_hash"] = self.embedding_hash

        if self.logits_shape is not None:
            d["logits_shape"] = self.logits_shape
            d["logits_hash"] = self.logits_hash
            d["top_token_id"] = self.top_token_id
            d["top_token_logit"] = self.top_token_logit

        if self.error_message is not None:
            d["error_message"] = self.error_message

        return d


def _hash_array(arr: np.ndarray) -> str:
    """Generate SHA256 hash of numpy array (first 16 hex chars)."""
    return hashlib.sha256(arr.tobytes()).hexdigest()[:16]


def _compute_vocab_hash(llm: "Llama", n_tokens: int = 1000) -> str:
    """Compute hash of first N tokens for vocab fingerprinting."""
    tokens = []
    for i in range(min(n_tokens, llm.n_vocab())):
        try:
            token_bytes = llm.detokenize([i])
            tokens.append(token_bytes)
        except Exception:
            tokens.append(b"")

    combined = b"".join(tokens)
    return hashlib.sha256(combined).hexdigest()[:16]


class MistralTokenizerBridge:
    """
    Bridge between text and native CDNA v2 streaming for Mistral models.

    Provides:
    - encode(): text → token IDs
    - decode(): token IDs → text
    - lookup_embeddings(): token IDs → embeddings [seq, d_model]
    - project_to_logits(): hidden states → logits [vocab_size]

    Uses llama_cpp for tokenization (matches model vocab exactly) and
    lazy-loads embedding/output weights from GGUF via GGUFScanner.

    Singleton pattern: one bridge per GGUF path.

    Example:
        bridge = MistralTokenizerBridge.get_bridge("model.gguf")

        # Encode
        ids, receipt = bridge.encode("Hello world")

        # Embed
        embeddings, receipt = bridge.lookup_embeddings(ids)

        # ... run through CDNA streaming blocks ...

        # Project to logits
        logits, receipt = bridge.project_to_logits(hidden_states)

        # Get next token
        next_id = np.argmax(logits[-1])

        # Decode
        text, receipt = bridge.decode([next_id])
    """

    # Singleton cache: one bridge per GGUF path
    _instances: Dict[str, "MistralTokenizerBridge"] = {}

    @classmethod
    def get_bridge(cls, gguf_path: Union[str, Path]) -> "MistralTokenizerBridge":
        """
        Get or create a tokenizer bridge for a GGUF file.

        Args:
            gguf_path: Path to the GGUF file

        Returns:
            MistralTokenizerBridge instance (singleton per path)
        """
        gguf_path = Path(gguf_path).absolute()
        key = str(gguf_path)

        if key not in cls._instances:
            cls._instances[key] = cls(gguf_path)

        return cls._instances[key]

    @classmethod
    def clear_cache(cls) -> None:
        """Clear all cached bridge instances."""
        for bridge in cls._instances.values():
            bridge.close()
        cls._instances.clear()

    def __init__(self, gguf_path: Union[str, Path]):
        """
        Initialize the tokenizer bridge.

        Use get_bridge() instead of direct construction for singleton behavior.

        Args:
            gguf_path: Path to the GGUF file
        """
        if not HAVE_LLAMA_CPP:
            raise ImportError(
                "llama-cpp-python is required for MistralTokenizerBridge. "
                "Install with: pip install llama-cpp-python"
            )

        self.gguf_path = Path(gguf_path).absolute()

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

        # Load llama.cpp with minimal context (only need vocab, not inference)
        # n_ctx=128: minimal context window
        # n_gpu_layers=0: CPU only (we're not doing inference, just tokenization)
        # logits_all=False: don't compute logits
        self._llm = Llama(
            model_path=str(self.gguf_path),
            n_ctx=128,
            n_gpu_layers=0,
            logits_all=False,
            verbose=False,
        )

        # Cache model info
        self.vocab_size = self._llm.n_vocab()
        self.vocab_hash = _compute_vocab_hash(self._llm)

        # Lazy-loaded weights
        self._token_embd: Optional[np.ndarray] = None
        self._output_weight: Optional[np.ndarray] = None
        self._output_norm: Optional[np.ndarray] = None  # WO-CDNA-ROPE-01: Final norm before LM head
        self._gguf_scanner = None

        # Model dimensions (from Mistral 7B)
        self.d_model = 4096
        # WO-CDNA-NORM-EPS-FIX-01: Use correct epsilon from GGUF metadata
        # llama.attention.layer_norm_rms_epsilon: 9.999999747378752e-06 ≈ 1e-5
        self.norm_eps = 1e-5  # RMSNorm epsilon

    def _ensure_gguf_scanner(self):
        """Ensure GGUF scanner is initialized for weight extraction."""
        if self._gguf_scanner is None:
            from helix_cdc.gguf.cache import get_scanner
            # Use global cache to avoid repeated scanning
            self._gguf_scanner = get_scanner(self.gguf_path, quiet=True)

    def _load_token_embeddings(self) -> np.ndarray:
        """Load token embedding weights from GGUF (lazy)."""
        if self._token_embd is not None:
            return self._token_embd

        self._ensure_gguf_scanner()

        # Find token_embd tensor
        tensor_info = self._gguf_scanner.get_tensor_by_name("token_embd.weight")
        if tensor_info is None:
            raise ValueError("token_embd.weight not found in GGUF")

        tensor = self._gguf_scanner.extract_tensor(tensor_info)

        # GGML stores embeddings as [d_model, vocab_size], we need [vocab_size, d_model]
        # for row-indexing by token ID
        if tensor.shape[0] < tensor.shape[1]:
            # Shape is [d_model, vocab_size], transpose to [vocab_size, d_model]
            self._token_embd = tensor.T
        else:
            # Already [vocab_size, d_model]
            self._token_embd = tensor

        # Update d_model from actual shape (now dim 1)
        self.d_model = self._token_embd.shape[1]

        return self._token_embd

    def _load_output_weight(self) -> np.ndarray:
        """Load output (LM head) weights from GGUF (lazy)."""
        if self._output_weight is not None:
            return self._output_weight

        self._ensure_gguf_scanner()

        # Try output.weight first
        tensor_info = self._gguf_scanner.get_tensor_by_name("output.weight")

        if tensor_info is None:
            # Some models tie embeddings: output.weight = token_embd.weight
            # Check for tied weights
            tensor_info = self._gguf_scanner.get_tensor_by_name("token_embd.weight")
            if tensor_info is None:
                raise ValueError("Neither output.weight nor token_embd.weight found in GGUF")
            # Use tied weights
            self._output_weight = self._load_token_embeddings()
            return self._output_weight

        tensor = self._gguf_scanner.extract_tensor(tensor_info)

        # GGML stores output.weight as [d_model, vocab_size], we need [vocab_size, d_model]
        # for logits = hidden @ output.T = hidden @ [d_model, vocab_size] = [vocab_size]
        if tensor.shape[0] < tensor.shape[1]:
            # Shape is [d_model, vocab_size], transpose to [vocab_size, d_model]
            self._output_weight = tensor.T
        else:
            # Already [vocab_size, d_model]
            self._output_weight = tensor

        return self._output_weight

    def _load_output_norm(self) -> np.ndarray:
        """Load output_norm.weight (final RMSNorm before LM head) from GGUF (lazy)."""
        if self._output_norm is not None:
            return self._output_norm

        self._ensure_gguf_scanner()

        tensor_info = self._gguf_scanner.get_tensor_by_name("output_norm.weight")
        if tensor_info is None:
            raise ValueError("output_norm.weight not found in GGUF")

        self._output_norm = self._gguf_scanner.extract_tensor(tensor_info)
        return self._output_norm

    def _rms_norm(self, x: np.ndarray, weight: np.ndarray) -> np.ndarray:
        """Apply RMSNorm: x * weight / sqrt(mean(x^2) + eps)"""
        variance = np.mean(x ** 2, axis=-1, keepdims=True)
        x_normed = x / np.sqrt(variance + self.norm_eps)
        return x_normed * weight

    def encode(
        self,
        text: str,
        add_bos: bool = True,
    ) -> Tuple[List[int], TokenizerBridgeReceipt]:
        """
        Encode text to token IDs.

        Args:
            text: Input text to tokenize
            add_bos: Whether to add beginning-of-sequence token

        Returns:
            (token_ids, receipt)
        """
        t0 = time.perf_counter()

        try:
            # Encode using llama.cpp tokenizer
            text_bytes = text.encode("utf-8")
            token_ids = self._llm.tokenize(text_bytes, add_bos=add_bos)
            token_ids = list(token_ids)  # Convert from ctypes array

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

            receipt = TokenizerBridgeReceipt(
                vocab_size=self.vocab_size,
                vocab_hash=self.vocab_hash,
                gguf_path=str(self.gguf_path),
                operation="encode",
                token_ids=token_ids,
                token_count=len(token_ids),
                text_length=len(text),
                duration_ms=duration_ms,
                status="PASS",
            )

            return token_ids, receipt

        except Exception as e:
            duration_ms = (time.perf_counter() - t0) * 1000
            receipt = TokenizerBridgeReceipt(
                vocab_size=self.vocab_size,
                vocab_hash=self.vocab_hash,
                gguf_path=str(self.gguf_path),
                operation="encode",
                text_length=len(text),
                duration_ms=duration_ms,
                status="FAIL",
                error_message=str(e),
            )
            return [], receipt

    def decode(
        self,
        token_ids: List[int],
    ) -> Tuple[str, TokenizerBridgeReceipt]:
        """
        Decode token IDs to text.

        Args:
            token_ids: List of token IDs

        Returns:
            (text, receipt)
        """
        t0 = time.perf_counter()

        try:
            # Decode using llama.cpp tokenizer
            text_bytes = self._llm.detokenize(token_ids)
            text = text_bytes.decode("utf-8", errors="replace")

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

            receipt = TokenizerBridgeReceipt(
                vocab_size=self.vocab_size,
                vocab_hash=self.vocab_hash,
                gguf_path=str(self.gguf_path),
                operation="decode",
                token_ids=token_ids,
                token_count=len(token_ids),
                text_length=len(text),
                duration_ms=duration_ms,
                status="PASS",
            )

            return text, receipt

        except Exception as e:
            duration_ms = (time.perf_counter() - t0) * 1000
            receipt = TokenizerBridgeReceipt(
                vocab_size=self.vocab_size,
                vocab_hash=self.vocab_hash,
                gguf_path=str(self.gguf_path),
                operation="decode",
                token_count=len(token_ids),
                duration_ms=duration_ms,
                status="FAIL",
                error_message=str(e),
            )
            return "", receipt

    def lookup_embeddings(
        self,
        token_ids: List[int],
    ) -> Tuple[np.ndarray, TokenizerBridgeReceipt]:
        """
        Look up embeddings for token IDs.

        Args:
            token_ids: List of token IDs

        Returns:
            (embeddings [seq, d_model], receipt)
        """
        t0 = time.perf_counter()

        try:
            # Load embedding matrix (lazy)
            embd = self._load_token_embeddings()

            # Look up embeddings for each token
            token_ids_arr = np.array(token_ids, dtype=np.int64)

            # Validate token IDs are in range
            if np.any(token_ids_arr < 0) or np.any(token_ids_arr >= self.vocab_size):
                raise ValueError(
                    f"Token IDs must be in range [0, {self.vocab_size}), "
                    f"got min={token_ids_arr.min()}, max={token_ids_arr.max()}"
                )

            # Index into embedding matrix: [seq, d_model]
            embeddings = embd[token_ids_arr]

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

            receipt = TokenizerBridgeReceipt(
                vocab_size=self.vocab_size,
                vocab_hash=self.vocab_hash,
                gguf_path=str(self.gguf_path),
                operation="embed_lookup",
                token_ids=token_ids,
                token_count=len(token_ids),
                embedding_shape=list(embeddings.shape),
                embedding_hash=_hash_array(embeddings),
                duration_ms=duration_ms,
                status="PASS",
            )

            return embeddings.astype(np.float32), receipt

        except Exception as e:
            duration_ms = (time.perf_counter() - t0) * 1000
            receipt = TokenizerBridgeReceipt(
                vocab_size=self.vocab_size,
                vocab_hash=self.vocab_hash,
                gguf_path=str(self.gguf_path),
                operation="embed_lookup",
                token_count=len(token_ids),
                duration_ms=duration_ms,
                status="FAIL",
                error_message=str(e),
            )
            return np.array([]), receipt

    def project_to_logits(
        self,
        hidden: np.ndarray,
        temperature: float = 1.0,
    ) -> Tuple[np.ndarray, TokenizerBridgeReceipt]:
        """
        Project hidden states to logits using the LM head.

        Computes: logits = hidden @ output_weight.T

        Args:
            hidden: Hidden states [..., d_model]
            temperature: Temperature for scaling (default 1.0)

        Returns:
            (logits [..., vocab_size], receipt)
        """
        t0 = time.perf_counter()

        try:
            # Load output norm and output weight (lazy)
            output_norm = self._load_output_norm()
            output_weight = self._load_output_weight()

            # Apply final RMSNorm before LM head (WO-CDNA-ROPE-01)
            hidden_normed = self._rms_norm(hidden.astype(np.float32), output_norm)

            # hidden_normed: [..., d_model]
            # output_weight: [vocab_size, d_model]
            # logits = hidden_normed @ output_weight.T = [..., vocab_size]
            logits = np.matmul(hidden_normed, output_weight.T)

            # Apply temperature
            if temperature != 1.0:
                logits = logits / temperature

            # Get top token info (from last position if 2D+)
            if logits.ndim == 1:
                last_logits = logits
            else:
                last_logits = logits[..., -1, :]

            top_token_id = int(np.argmax(last_logits))
            top_token_logit = float(last_logits[top_token_id])

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

            receipt = TokenizerBridgeReceipt(
                vocab_size=self.vocab_size,
                vocab_hash=self.vocab_hash,
                gguf_path=str(self.gguf_path),
                operation="lm_head",
                logits_shape=list(logits.shape),
                logits_hash=_hash_array(logits),
                top_token_id=top_token_id,
                top_token_logit=top_token_logit,
                duration_ms=duration_ms,
                status="PASS",
            )

            return logits, receipt

        except Exception as e:
            duration_ms = (time.perf_counter() - t0) * 1000
            receipt = TokenizerBridgeReceipt(
                vocab_size=self.vocab_size,
                vocab_hash=self.vocab_hash,
                gguf_path=str(self.gguf_path),
                operation="lm_head",
                duration_ms=duration_ms,
                status="FAIL",
                error_message=str(e),
            )
            return np.array([]), receipt

    def get_bos_token_id(self) -> int:
        """Get the beginning-of-sequence token ID."""
        return self._llm.token_bos()

    def get_eos_token_id(self) -> int:
        """Get the end-of-sequence token ID."""
        return self._llm.token_eos()

    def close(self) -> None:
        """Close the bridge and release resources."""
        if hasattr(self, '_llm') and self._llm is not None:
            # llama-cpp-python cleanup
            if hasattr(self._llm, 'close'):
                try:
                    self._llm.close()
                except Exception:
                    pass
            self._llm = None

        # Clear cached weights
        self._token_embd = None
        self._output_weight = None
        self._gguf_scanner = None

    def __repr__(self) -> str:
        return (
            f"MistralTokenizerBridge("
            f"vocab_size={self.vocab_size}, "
            f"d_model={self.d_model}, "
            f"vocab_hash={self.vocab_hash})"
        )


def check_tokenizer_bridge_available() -> Tuple[bool, str]:
    """
    Check if tokenizer bridge is available.

    Returns:
        (available, message)
    """
    if not HAVE_LLAMA_CPP:
        return False, "llama-cpp-python not installed"

    return True, "llama-cpp-python available"


if __name__ == "__main__":
    import sys

    print("=== MistralTokenizerBridge Test ===")
    print()

    available, msg = check_tokenizer_bridge_available()
    print(f"Available: {available}")
    print(f"Message: {msg}")

    if not available:
        print("Cannot run tests without llama-cpp-python")
        sys.exit(1)

    if len(sys.argv) < 2:
        print()
        print("Usage: python tokenizer_bridge.py <gguf_path>")
        print()
        print("Example:")
        print("  python tokenizer_bridge.py tmp/mistral_canonical_q4_k_m.gguf")
        sys.exit(0)

    gguf_path = sys.argv[1]
    print(f"\nLoading bridge from: {gguf_path}")

    bridge = MistralTokenizerBridge.get_bridge(gguf_path)
    print(f"Bridge: {bridge}")
    print()

    # Test encode
    test_text = "Hello, world!"
    print(f"Encoding: '{test_text}'")
    token_ids, encode_receipt = bridge.encode(test_text)
    print(f"  Token IDs: {token_ids}")
    print(f"  Receipt: {encode_receipt.operation}, status={encode_receipt.status}")
    print()

    # Test decode
    print(f"Decoding: {token_ids}")
    decoded_text, decode_receipt = bridge.decode(token_ids)
    print(f"  Text: '{decoded_text}'")
    print(f"  Receipt: {decode_receipt.operation}, status={decode_receipt.status}")
    print()

    # Test embedding lookup
    print(f"Looking up embeddings for: {token_ids}")
    embeddings, embed_receipt = bridge.lookup_embeddings(token_ids)
    print(f"  Shape: {embeddings.shape}")
    print(f"  Dtype: {embeddings.dtype}")
    print(f"  Receipt: {embed_receipt.operation}, status={embed_receipt.status}")
    print()

    # Test logits projection
    print(f"Projecting to logits from: {embeddings.shape}")
    logits, logits_receipt = bridge.project_to_logits(embeddings)
    print(f"  Shape: {logits.shape}")
    print(f"  Top token ID: {logits_receipt.top_token_id}")
    print(f"  Receipt: {logits_receipt.operation}, status={logits_receipt.status}")
    print()

    # Decode top token
    top_token_text, _ = bridge.decode([logits_receipt.top_token_id])
    print(f"Top token text: '{top_token_text}'")

    print()
    print("DONE")
