#!/usr/bin/env python3
"""
GPU Embedding Worker for OCR Provenance MCP System

Generates 768-dimensional embeddings using nomic-embed-text-v1.5.
Supports CUDA, MPS (Apple Silicon), and CPU. BF16 on CUDA when supported.

CRITICAL CONSTRAINTS:
- CP-004: Local inference ONLY - NEVER fall back to cloud
- AP-008: NO Flash Attention
- Task prefixes: "search_document: " for chunks, "search_query: " for queries
- Device priority: CUDA > MPS > CPU (via gpu_utils.resolve_device)

Usage:
    # Embed chunks (document mode)
    python embedding_worker.py --chunks "text1" "text2" --json

    # Embed query (search mode)
    python embedding_worker.py --query "search text" --json

    # From stdin (for large batches from TypeScript)
    echo '["text1", "text2"]' | python embedding_worker.py --stdin --json

    # Daemon mode (load model once, process requests via stdin JSON lines)
    python embedding_worker.py --daemon
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import sys
import time
import warnings
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Optional

# Suppress pynvml FutureWarning from torch.cuda (emitted to stderr, breaks JSON-RPC)
warnings.filterwarnings("ignore", category=FutureWarning, message=".*pynvml.*")

import numpy as np
import torch
from sentence_transformers import SentenceTransformer

try:
    # When run as a script from python/ directory
    from gpu_utils import (
        EmbeddingModelError,
        GPUNotAvailableError,
        GPUOutOfMemoryError,
        auto_tune_batch_size,
        resolve_device as gpu_resolve_device,
        startup_gpu,
    )
except ImportError:
    # When imported as part of python package
    from .gpu_utils import (
        EmbeddingModelError,
        GPUNotAvailableError,
        GPUOutOfMemoryError,
        auto_tune_batch_size,
        resolve_device as gpu_resolve_device,
        startup_gpu,
    )

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    stream=sys.stderr,
)
logger = logging.getLogger(__name__)


# =============================================================================
# Constants - DO NOT CHANGE
# =============================================================================


def _resolve_model_path() -> Path:
    """Resolve embedding model path from env var or default locations."""
    env_path = os.environ.get("EMBEDDING_MODEL_PATH")
    if env_path:
        return Path(env_path)
    # Default: package-relative (works for dev and global npm install)
    pkg_path = Path(__file__).resolve().parent.parent / "models" / "nomic-embed-text-v1.5"
    if pkg_path.exists():
        return pkg_path
    # Fallback: user home directory
    return Path.home() / ".ocr-provenance-mcp" / "models" / "nomic-embed-text-v1.5"


MODEL_PATH = _resolve_model_path()
EMBEDDING_DIM = 768
MODEL_NAME = "nomic-embed-text-v1.5"
MODEL_VERSION = "1.5.0"

# Task prefixes - REQUIRED by nomic model
PREFIX_DOCUMENT = "search_document: "
PREFIX_QUERY = "search_query: "

# Batch configuration
MIN_BATCH_SIZE = 1  # Must support single-item batches for VLM descriptions

# Device configuration - auto-detect best device (CUDA > MPS > CPU)
DEFAULT_DEVICE = os.environ.get("EMBEDDING_DEVICE", "auto")


# =============================================================================
# Data Classes - MUST match TypeScript interfaces
# =============================================================================


@dataclass
class EmbeddingResult:
    """
    Result from batch embedding generation.
    MUST match src/models/embedding.ts EmbeddingGenerationResult
    """

    success: bool
    embeddings: list[list[float]]  # (n, 768) as nested list for JSON
    count: int
    elapsed_ms: float
    ms_per_chunk: float
    device: str
    batch_size: int
    model: str = MODEL_NAME
    model_version: str = MODEL_VERSION
    vram_used_gb: float = 0.0
    error: Optional[str] = None


@dataclass
class QueryEmbeddingResult:
    """Result from single query embedding."""

    success: bool
    embedding: list[float]  # (768,) as list for JSON
    elapsed_ms: float
    device: str
    model: str = MODEL_NAME
    error: Optional[str] = None


# =============================================================================
# Global Model Singleton
# =============================================================================

_model: Optional[SentenceTransformer] = None
_device: Optional[str] = None

# =============================================================================
# CUDA Startup Validation (runs on module load)
# =============================================================================

# Device startup: configure memory pool, validate environment, log device info
_startup_info = startup_gpu("Embedding worker")

# Auto-tune batch size based on detected device and VRAM
DEFAULT_BATCH_SIZE = auto_tune_batch_size(_startup_info.get("device", "cpu"))


# =============================================================================
# Core Functions
# =============================================================================


def resolve_device(requested: str = DEFAULT_DEVICE) -> str:
    """
    Resolve compute device via gpu_utils.resolve_device().

    Supports 'auto', 'cuda', 'cuda:N', 'mps', and 'cpu'.
    'auto' selects the best available device (CUDA > MPS > CPU).

    Args:
        requested: Requested device string ('auto', 'cuda:0', 'mps', 'cpu')

    Returns:
        Validated device string (e.g., 'cuda:0', 'mps', 'cpu')

    Raises:
        GPUNotAvailableError: If an explicitly requested device is unavailable
    """
    return gpu_resolve_device(requested)


def load_model(device: str = DEFAULT_DEVICE) -> SentenceTransformer:
    """
    Load nomic-embed-text-v1.5 to the best available device.

    Uses BF16 on CUDA when GPU supports it, FP32 on MPS and CPU.

    Args:
        device: Device string ('auto', 'cuda:0', 'mps', 'cpu')

    Returns:
        Loaded SentenceTransformer model

    Raises:
        EmbeddingModelError: Model not found or failed to load
        GPUNotAvailableError: If an explicitly requested device is unavailable
    """
    global _model, _device

    # Resolve device (CUDA > MPS > CPU)
    device = resolve_device(device)

    # Return cached model if same device
    if _model is not None and _device == device:
        return _model

    if not MODEL_PATH.exists():
        raise EmbeddingModelError(
            f"Model not found at {MODEL_PATH}. Download with: "
            "huggingface-cli download nomic-ai/nomic-embed-text-v1.5 "
            "--local-dir models/nomic-embed-text-v1.5",
            model_path=str(MODEL_PATH),
        )

    # Check required files
    required = ["config.json", "model.safetensors", "tokenizer.json"]
    missing = [f for f in required if not (MODEL_PATH / f).exists()]
    if missing:
        raise EmbeddingModelError(f"Missing model files: {missing}", model_path=str(MODEL_PATH))

    # Select dtype: BF16 on CUDA when supported, FP32 otherwise
    model_kwargs: dict = {}
    if device.startswith("cuda"):
        try:
            bf16_ok = torch.cuda.is_bf16_supported()
        except Exception:
            bf16_ok = False
        if bf16_ok:
            model_kwargs["torch_dtype"] = torch.bfloat16
            dtype_label = "bfloat16"
        else:
            dtype_label = "float32"
    else:
        # MPS and CPU: use FP32
        dtype_label = "float32"

    logger.info("Loading model to %s with %s precision...", device, dtype_label)

    try:
        # M-7: Pin revision to prevent supply-chain attacks with trust_remote_code=True
        _model = SentenceTransformer(
            str(MODEL_PATH),
            device=device,
            trust_remote_code=True,
            local_files_only=True,
            revision="0ef289df6812e6e0e532352f5174b161cbfb30db",
            model_kwargs=model_kwargs,
        )
        _device = device

        # Verify dimensions
        dim = _model.get_sentence_embedding_dimension()
        if dim != EMBEDDING_DIM:
            raise EmbeddingModelError(
                f"Wrong embedding dimension: {dim}, expected {EMBEDDING_DIM}",
                model_path=str(MODEL_PATH),
            )

        logger.info(
            "Model loaded: %s, dim=%d, device=%s, dtype=%s",
            MODEL_NAME, EMBEDDING_DIM, device, dtype_label,
        )

        # NOTE: torch.compile intentionally NOT used here. This worker is spawned
        # per-call via PythonShell/child_process, so compilation overhead (~30s)
        # is never amortized across requests. BF16 provides the perf gain on CUDA.

        return _model

    except (GPUNotAvailableError, EmbeddingModelError):
        raise
    except Exception as e:
        raise EmbeddingModelError(
            f"Failed to load model: {e}", model_path=str(MODEL_PATH), cause=e
        ) from e


def embed_chunks(
    chunks: list[str], batch_size: int = DEFAULT_BATCH_SIZE, device: str = DEFAULT_DEVICE
) -> np.ndarray:
    """
    Embed document chunks with "search_document: " prefix.

    Args:
        chunks: Text chunks to embed
        batch_size: GPU batch size (default 512)
        device: CUDA device

    Returns:
        np.ndarray of shape (n_chunks, 768), dtype float32

    Raises:
        GPUNotAvailableError: No GPU
        EmbeddingModelError: Model error
    """
    if not chunks:
        return np.zeros((0, EMBEDDING_DIM), dtype=np.float32)

    # load_model() already calls resolve_device() internally
    model = load_model(device)

    # Add task prefix - REQUIRED by nomic model
    prefixed = [f"{PREFIX_DOCUMENT}{chunk}" for chunk in chunks]

    # Generate embeddings (device already resolved in load_model)
    embeddings = model.encode(
        prefixed,
        batch_size=batch_size,
        convert_to_numpy=True,
        normalize_embeddings=True,  # L2 normalize for cosine similarity
        show_progress_bar=False,
        device=_device,
    )

    return embeddings.astype(np.float32)


def embed_query(query: str, device: str = DEFAULT_DEVICE) -> np.ndarray:
    """
    Embed search query with "search_query: " prefix.

    Args:
        query: Search query text
        device: CUDA device string

    Returns:
        np.ndarray of shape (768,), dtype float32
    """
    # load_model() already calls resolve_device() internally
    model = load_model(device)

    # Add query task prefix
    prefixed = f"{PREFIX_QUERY}{query}"

    embedding = model.encode(
        [prefixed],
        convert_to_numpy=True,
        normalize_embeddings=True,
        show_progress_bar=False,
        device=_device,
    )

    return embedding[0].astype(np.float32)


def embed_with_oom_recovery(
    chunks: list[str],
    initial_batch_size: int = DEFAULT_BATCH_SIZE,
    device: str = DEFAULT_DEVICE,
) -> tuple[np.ndarray, int]:
    """
    Embed with automatic OOM recovery via batch size reduction.

    Halves batch size on OOM until MIN_BATCH_SIZE.
    On CUDA: catches torch.cuda.OutOfMemoryError and clears cache.
    On MPS/CPU: catches MemoryError and RuntimeError with "out of memory".

    Args:
        chunks: Text chunks to embed
        initial_batch_size: Starting batch size
        device: Device string

    Returns:
        Tuple of (embeddings, final_batch_size)

    Raises:
        GPUOutOfMemoryError: OOM at minimum batch size
    """
    resolved = resolve_device(device)
    is_cuda = resolved.startswith("cuda")
    batch_size = initial_batch_size

    while batch_size >= MIN_BATCH_SIZE:
        try:
            if is_cuda:
                torch.cuda.empty_cache()
            embeddings = embed_chunks(chunks, batch_size, device)
            return embeddings, batch_size
        except (torch.cuda.OutOfMemoryError, MemoryError, RuntimeError) as e:
            if isinstance(e, RuntimeError) and "out of memory" not in str(e).lower():
                raise
            if is_cuda:
                torch.cuda.empty_cache()
            batch_size //= 2
            if batch_size >= MIN_BATCH_SIZE:
                logger.warning("OOM on %s: Reducing batch size to %d", resolved, batch_size)

    raise GPUOutOfMemoryError(
        f"OOM with {len(chunks)} chunks on {resolved}. "
        f"Tried batch sizes {initial_batch_size} down to {MIN_BATCH_SIZE}.",
        vram_required=None,
        vram_available=None,
    )


def generate_embeddings(
    chunks: list[str],
    batch_size: int = DEFAULT_BATCH_SIZE,
    device: str = DEFAULT_DEVICE,
) -> EmbeddingResult:
    """
    Generate embeddings with full metrics for TypeScript bridge.

    This is the main entry point for batch embedding.

    Args:
        chunks: Text chunks to embed
        batch_size: Initial batch size
        device: CUDA device

    Returns:
        EmbeddingResult with embeddings and metrics
    """
    start_time = time.perf_counter()
    resolved_device = resolve_device(device)
    is_cuda = resolved_device.startswith("cuda")

    # Reset VRAM tracking (CUDA only)
    if is_cuda:
        torch.cuda.reset_peak_memory_stats()

    try:
        embeddings_np, final_batch_size = embed_with_oom_recovery(chunks, batch_size, device)

        elapsed_ms = (time.perf_counter() - start_time) * 1000
        ms_per_chunk = elapsed_ms / len(chunks) if chunks else 0
        vram_gb = torch.cuda.max_memory_allocated() / (1024**3) if is_cuda else 0.0

        # H-8: Convert to list and delete numpy array to avoid ~7x memory overlap
        # (numpy float32 ~1.5MB vs Python float64 list ~10.7MB for 500x768)
        embeddings_list = embeddings_np.tolist()
        del embeddings_np

        return EmbeddingResult(
            success=True,
            embeddings=embeddings_list,
            count=len(chunks),
            elapsed_ms=round(elapsed_ms, 2),
            ms_per_chunk=round(ms_per_chunk, 4),
            device=resolved_device,
            batch_size=final_batch_size,
            vram_used_gb=round(vram_gb, 3),
            error=None,
        )

    except Exception as e:
        elapsed_ms = (time.perf_counter() - start_time) * 1000
        logger.error("Embedding generation failed: %s", e)
        return EmbeddingResult(
            success=False,
            embeddings=[],
            count=0,
            elapsed_ms=round(elapsed_ms, 2),
            ms_per_chunk=0,
            device=resolved_device,
            batch_size=batch_size,
            error=str(e),
        )


def generate_query_embedding(query: str, device: str = DEFAULT_DEVICE) -> QueryEmbeddingResult:
    """
    Generate query embedding with metrics.

    Args:
        query: Search query text
        device: CUDA device

    Returns:
        QueryEmbeddingResult with embedding and metrics
    """
    start_time = time.perf_counter()
    resolved_device = resolve_device(device)

    try:
        embedding = embed_query(query, device)
        elapsed_ms = (time.perf_counter() - start_time) * 1000

        return QueryEmbeddingResult(
            success=True,
            embedding=embedding.tolist(),
            elapsed_ms=round(elapsed_ms, 2),
            device=resolved_device,
            error=None,
        )

    except Exception as e:
        elapsed_ms = (time.perf_counter() - start_time) * 1000
        logger.error("Query embedding failed: %s", e)
        return QueryEmbeddingResult(
            success=False,
            embedding=[],
            elapsed_ms=round(elapsed_ms, 2),
            device=resolved_device,
            error=str(e),
        )


# =============================================================================
# Daemon Mode
# =============================================================================


def daemon_loop(device: str, batch_size: int) -> None:
    """Run embedding daemon: load model once, process requests from stdin.

    Protocol (JSON lines over stdin/stdout):
      Ready:    {"status": "ready", "device": "cuda:0", "model": "...", ...}
      Request:  {"action": "embed_chunks", "chunks": [...], "batch_size": 64}
      Request:  {"action": "embed_query", "query": "search text"}
      Request:  {"action": "shutdown"}
      Response: EmbeddingResult or QueryEmbeddingResult as JSON (one line)
      Error:    {"success": false, "error": "..."}

    All logging goes to stderr. Only JSON protocol on stdout.
    """
    global _model, _device

    start = time.perf_counter()
    load_model(device)  # Force model into GPU VRAM (cached in global _model)
    load_ms = (time.perf_counter() - start) * 1000

    resolved = resolve_device(device)

    # Send ready signal
    ready = {
        "status": "ready",
        "device": resolved,
        "model": MODEL_NAME,
        "model_version": MODEL_VERSION,
        "load_time_ms": round(load_ms, 2),
    }
    print(json.dumps(ready), flush=True)

    logger.info(
        "Daemon: Model loaded in %.0fms on %s. Waiting for requests on stdin...",
        load_ms, resolved,
    )

    items_processed = 0

    # Process requests from stdin
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue

        try:
            request = json.loads(line)
        except json.JSONDecodeError as e:
            error_resp = {"success": False, "error": f"Invalid JSON: {e}"}
            print(json.dumps(error_resp), flush=True)
            continue

        action = request.get("action", "")

        if action == "shutdown":
            logger.info(
                "Daemon: Shutdown requested after %d items. Freeing GPU memory...",
                items_processed,
            )
            # Free model and VRAM
            if _model is not None:
                del _model
                _model = None
            _device = None
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            print(json.dumps({"status": "shutdown", "success": True}), flush=True)
            break

        elif action == "embed_chunks":
            chunks = request.get("chunks", [])
            req_batch_size = request.get("batch_size", batch_size)
            result = generate_embeddings(chunks, req_batch_size, device)
            print(json.dumps(asdict(result)), flush=True)
            items_processed += 1

        elif action == "embed_query":
            query = request.get("query", "")
            result = generate_query_embedding(query, device)
            print(json.dumps(asdict(result)), flush=True)
            items_processed += 1

        else:
            error_resp = {"success": False, "error": f"Unknown action: {action}"}
            print(json.dumps(error_resp), flush=True)

    # Clean exit: stdin EOF or shutdown action
    logger.info("Daemon: Exiting after %d items processed.", items_processed)


# =============================================================================
# CLI Entry Point
# =============================================================================


def main() -> None:
    """CLI entry point for embedding worker."""
    parser = argparse.ArgumentParser(
        description="GPU Embedding Worker - nomic-embed-text-v1.5",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python embedding_worker.py --chunks "text1" "text2" --json
  python embedding_worker.py --query "search text" --json
  echo '["text1", "text2"]' | python embedding_worker.py --stdin --json
  python embedding_worker.py --daemon
        """,
    )

    # Daemon mode flag
    parser.add_argument(
        "--daemon",
        action="store_true",
        help="Run in daemon mode: load model once, accept requests via stdin",
    )

    # Input modes (mutually exclusive, not required when --daemon is used)
    input_group = parser.add_mutually_exclusive_group(required=False)
    input_group.add_argument("--chunks", nargs="+", help="Texts to embed")
    input_group.add_argument("--query", help="Search query to embed")
    input_group.add_argument("--stdin", action="store_true", help="Read JSON array from stdin")

    # Configuration
    parser.add_argument(
        "--batch-size", type=int, default=DEFAULT_BATCH_SIZE, help="Batch size for GPU"
    )
    parser.add_argument("--device", default=DEFAULT_DEVICE, help="CUDA device")
    parser.add_argument("--model-path", help="Path to embedding model directory")
    parser.add_argument("--json", action="store_true", help="JSON output for TypeScript bridge")

    args = parser.parse_args()

    # Override model path if specified via CLI
    if args.model_path:
        global MODEL_PATH
        MODEL_PATH = Path(args.model_path)

    # ---- DAEMON MODE ----
    if args.daemon:
        try:
            daemon_loop(args.device, args.batch_size)
        except Exception as e:
            logger.exception("Daemon fatal error: %s", e)
            error_resp = {
                "success": False,
                "error": str(e),
                "error_type": type(e).__name__,
            }
            print(json.dumps(error_resp), flush=True)
            sys.exit(1)
        return

    # ---- CLI MODE ----
    # Validate that an input mode is specified in non-daemon mode
    if not args.chunks and not args.query and not args.stdin:
        parser.error("One of --daemon, --chunks, --query, or --stdin is required")

    try:
        if args.query:
            # Query mode
            result = generate_query_embedding(args.query, args.device)

        else:
            # Chunk mode
            if args.stdin:
                chunks = json.load(sys.stdin)
                if not isinstance(chunks, list):
                    raise ValueError("stdin must be a JSON array of strings")
            else:
                chunks = args.chunks

            result = generate_embeddings(chunks, args.batch_size, args.device)

        if args.json:
            result_dict = asdict(result)
            result_dict["device_used"] = str(result.device)
            print(json.dumps(result_dict))
            if not result.success:
                sys.exit(1)
        else:
            # Human readable output
            if isinstance(result, EmbeddingResult):
                if result.success:
                    throughput = 1000 / result.ms_per_chunk if result.ms_per_chunk > 0 else 0
                    print(f"Embedded {result.count} chunks in {result.elapsed_ms}ms")
                    print(f"Throughput: {throughput:.0f} chunks/sec")
                    print(f"VRAM used: {result.vram_used_gb:.3f} GB")
                else:
                    print(f"ERROR: {result.error}", file=sys.stderr)
                    sys.exit(1)
            else:
                if result.success:
                    print(f"Query embedded in {result.elapsed_ms}ms")
                else:
                    print(f"ERROR: {result.error}", file=sys.stderr)
                    sys.exit(1)

    except Exception as e:
        error_result = {
            "success": False,
            "error": str(e),
            "error_type": type(e).__name__,
        }
        if args.json:
            print(json.dumps(error_result))
        else:
            print(f"ERROR: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()
