# echo_cog/backend_llama_cpp.py
"""
Warm llama_cpp backend for Echo Cog.

Uses llama-cpp-python bindings with a singleton model cache to avoid
the cold-start overhead of spawning llama-cli for each question.

Historical performance (July 2025): ~4.51s/question with warm model
Current subprocess approach: ~20s/question due to model reload overhead

This module provides the warm-model path to match historical latency.
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional

from llama_cpp import Llama  # type: ignore


# Singleton model cache - keeps models warm in memory
_MODEL_CACHE: Dict[str, Llama] = {}


@dataclass
class LlamaCppBackendConfig:
    """Configuration for the llama_cpp warm backend."""
    model_path: str
    ctx_size: int = 2048
    threads: int = 8
    seed: int = 1234
    n_gpu_layers: int = -1  # -1 = auto (use GPU if available), 0 = CPU-only


def get_llama(cfg: LlamaCppBackendConfig) -> Llama:
    """
    Get or create a Llama model instance (singleton pattern).

    The model stays warm in memory after first load, giving us
    ~4x speedup over subprocess-based llama-cli calls.
    """
    global _MODEL_CACHE

    # Cache key includes params that affect model loading (NOT seed - seed only affects generation)
    # With greedy decoding (top_k=1, temperature=0), seed doesn't affect output anyway
    key = f"{cfg.model_path}|ctx={cfg.ctx_size}|threads={cfg.threads}|gpu={cfg.n_gpu_layers}"

    if key in _MODEL_CACHE:
        return _MODEL_CACHE[key]

    # First load - this is the slow part (~5-10s), but only happens once
    model_path = str(Path(cfg.model_path).expanduser().resolve())

    # Log GPU configuration for diagnostics
    try:
        import torch
        cuda_available = torch.cuda.is_available()
        cuda_device = torch.cuda.get_device_name(0) if cuda_available else "N/A"
    except ImportError:
        cuda_available = False
        cuda_device = "PyTorch not installed"

    print(f"[llama_cpp] Loading model: {model_path}")
    print(f"[llama_cpp] Config: ctx={cfg.ctx_size}, threads={cfg.threads}, n_gpu_layers={cfg.n_gpu_layers}")
    print(f"[llama_cpp] CUDA available: {cuda_available}, device: {cuda_device}")

    load_start = time.perf_counter()

    llm = Llama(
        model_path=model_path,
        n_ctx=cfg.ctx_size,
        n_threads=cfg.threads,
        n_gpu_layers=cfg.n_gpu_layers,
        seed=cfg.seed,
        logits_all=False,
        verbose=False,
    )

    load_ms = (time.perf_counter() - load_start) * 1000
    print(f"[llama_cpp] Model loaded in {load_ms:.1f}ms")

    _MODEL_CACHE[key] = llm
    return llm


def run_llama_cpp(
    cfg: LlamaCppBackendConfig,
    prompt: str,
    max_tokens: int,
    top_k: int = 1,  # Greedy-ish to match llama-cli behavior
) -> str:
    """
    Run inference using the warm llama_cpp backend.

    Args:
        cfg: Backend configuration
        prompt: The full prompt text
        max_tokens: Maximum tokens to generate
        top_k: Top-k sampling (1 = greedy)

    Returns:
        Generated text (answer portion only)
    """
    llm = get_llama(cfg)

    out = llm(
        prompt,
        max_tokens=max_tokens,
        top_k=top_k,
        stop=None,
        echo=False,
    )

    # Extract the generated text from llama_cpp output
    if out and "choices" in out and len(out["choices"]) > 0:
        return out["choices"][0].get("text", "")

    return ""


def clear_model_cache() -> None:
    """Clear the model cache (useful for testing or memory management)."""
    global _MODEL_CACHE
    _MODEL_CACHE.clear()


def close_llama() -> None:
    """
    Close all cached Llama instances and clear the cache.

    This properly releases GPU memory and file handles. Call this on
    server shutdown to avoid the 'LlamaModel was not closed' warning.
    """
    global _MODEL_CACHE
    for key, llm in list(_MODEL_CACHE.items()):
        try:
            # llama-cpp-python Llama instances have a close() method
            if hasattr(llm, 'close'):
                llm.close()
            elif hasattr(llm, '__del__'):
                llm.__del__()
        except Exception:
            pass  # Best-effort cleanup
    _MODEL_CACHE.clear()


def get_cache_info() -> Dict[str, int]:
    """Return info about cached models."""
    return {
        "num_cached_models": len(_MODEL_CACHE),
        "cached_keys": list(_MODEL_CACHE.keys()),
    }


@dataclass
class TTFTTimings:
    """First-token latency and breakdown timings."""
    ttft_ms: float = 0.0          # Time to first token (prompt submit → first token)
    prompt_process_ms: float = 0.0  # Time to process prompt (evaluate)
    decode_ms: float = 0.0         # Time for all tokens after first
    total_ms: float = 0.0          # Total generation time
    ms_per_token: float = 0.0      # Average per-token latency (decode phase)
    num_tokens: int = 0            # Number of tokens generated
    prompt_tokens: int = 0         # Number of prompt tokens

    def to_dict(self) -> Dict:
        """Convert to dictionary."""
        return {
            "ttft_ms": round(self.ttft_ms, 3),
            "prompt_process_ms": round(self.prompt_process_ms, 3),
            "decode_ms": round(self.decode_ms, 3),
            "total_ms": round(self.total_ms, 3),
            "ms_per_token": round(self.ms_per_token, 3),
            "num_tokens": self.num_tokens,
            "prompt_tokens": self.prompt_tokens,
        }


def run_llama_cpp_timed(
    cfg: LlamaCppBackendConfig,
    prompt: str,
    max_tokens: int,
    top_k: int = 1,
) -> Dict:
    """
    Run inference with detailed timing instrumentation.

    Returns a dict with:
        - text: Generated text
        - timing: TTFTTimings as dict
        - raw_output: Full llama_cpp output

    Metrics captured:
        - ttft_ms: Time to first token (prompt + first decode)
        - prompt_process_ms: Time to evaluate prompt
        - decode_ms: Time for remaining tokens
        - ms_per_token: Average decode latency
    """
    llm = get_llama(cfg)

    # Phase 1: Prompt evaluation timing
    # We use a streaming approach to capture TTFT accurately
    start_total = time.perf_counter()

    # Tokenize prompt to get token count
    prompt_tokens = llm.tokenize(prompt.encode())
    num_prompt_tokens = len(prompt_tokens)

    # Track timings during generation
    first_token_time = None
    token_times: List[float] = []
    generated_tokens: List[str] = []

    # Use streaming to capture per-token timing
    prompt_start = time.perf_counter()

    stream = llm(
        prompt,
        max_tokens=max_tokens,
        top_k=top_k,
        stop=None,
        echo=False,
        stream=True,
    )

    for chunk in stream:
        now = time.perf_counter()

        if first_token_time is None:
            first_token_time = now
            ttft_ms = (first_token_time - start_total) * 1000
            prompt_process_ms = (first_token_time - prompt_start) * 1000

        token_times.append(now)

        # Extract token text
        if "choices" in chunk and chunk["choices"]:
            token_text = chunk["choices"][0].get("text", "")
            generated_tokens.append(token_text)

    end_total = time.perf_counter()

    # Calculate metrics
    total_ms = (end_total - start_total) * 1000
    num_tokens = len(token_times)

    if first_token_time is None:
        # No tokens generated
        timing = TTFTTimings(
            ttft_ms=0.0,
            prompt_process_ms=total_ms,
            decode_ms=0.0,
            total_ms=total_ms,
            ms_per_token=0.0,
            num_tokens=0,
            prompt_tokens=num_prompt_tokens,
        )
    else:
        # Calculate decode timing
        if num_tokens > 1:
            decode_ms = (end_total - first_token_time) * 1000
            ms_per_token = decode_ms / (num_tokens - 1) if num_tokens > 1 else 0.0
        else:
            decode_ms = 0.0
            ms_per_token = 0.0

        timing = TTFTTimings(
            ttft_ms=ttft_ms,
            prompt_process_ms=prompt_process_ms,
            decode_ms=decode_ms,
            total_ms=total_ms,
            ms_per_token=ms_per_token,
            num_tokens=num_tokens,
            prompt_tokens=num_prompt_tokens,
        )

    # Combine generated text
    generated_text = "".join(generated_tokens)

    return {
        "text": generated_text,
        "timing": timing.to_dict(),
    }


def run_llama_cpp_batch_timed(
    cfg: LlamaCppBackendConfig,
    prompts: List[str],
    max_tokens: int,
    top_k: int = 1,
) -> List[Dict]:
    """
    Run inference on multiple prompts with timing.

    Returns list of results, each with text and timing.
    """
    results = []
    for prompt in prompts:
        result = run_llama_cpp_timed(cfg, prompt, max_tokens, top_k)
        results.append(result)
    return results
