#!/usr/bin/env python3
"""
Local cross-encoder reranking worker using sentence-transformers.
Supports CUDA (any NVIDIA GPU), MPS (Apple Silicon), and CPU.
BF16 on CUDA when GPU supports it, FP32 otherwise.

Reads JSON from stdin with format:
  {"query": "search text", "passages": [{"index": 0, "text": "...", "original_score": 0.5}, ...]}

Outputs JSON array of reranked results to stdout:
  [{"index": 0, "relevance_score": 0.95, "original_score": 0.5}, ...]

CRITICAL: All logging goes to stderr. Stdout is reserved for JSON output.
"""

import json
import logging
import sys
import time
import warnings

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

import torch
from sentence_transformers import CrossEncoder

try:
    from gpu_utils import GPUNotAvailableError, detect_best_device, startup_gpu
except ImportError:
    from .gpu_utils import GPUNotAvailableError, detect_best_device, startup_gpu

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


# =============================================================================
# DEVICE STARTUP VALIDATION (runs on module load)
# =============================================================================

# The reranker is a lightweight model (~129MB) that MUST work on CPU as a fallback.
# Other workers (OCR, VLM) may require GPU, but the reranker should never fail
# just because a GPU isn't available. Use startup_gpu() for GPU-available systems,
# but catch GPUNotAvailableError and fall back to CPU gracefully.
try:
    startup_gpu("Reranker worker")
except GPUNotAvailableError as e:
    logger.warning("GPU not available for reranker, falling back to CPU: %s", e)
    import os
    os.environ["TORCH_DEVICE"] = "cpu"


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

_cached_model = None
_cached_model_name = None


def get_model(model_name: str):
    """Get or load a cached cross-encoder model. Prefers GPU, falls back to CPU."""
    global _cached_model, _cached_model_name

    if _cached_model is not None and _cached_model_name == model_name:
        logger.info("Using cached reranker model: %s", model_name)
        return _cached_model

    device = detect_best_device()
    logger.info("Loading reranker model to %s: %s", device, model_name)
    start = time.time()

    try:
        _cached_model = CrossEncoder(model_name, max_length=512, device=device)
    except Exception as load_err:
        if device != "cpu":
            # GPU load failed — fall back to CPU. The reranker is small enough
            # to run on CPU with acceptable latency (~50ms for 20 passages).
            logger.warning(
                "Failed to load reranker on %s: %s. Falling back to CPU.", device, load_err
            )
            device = "cpu"
            _cached_model = CrossEncoder(model_name, max_length=512, device=device)
        else:
            raise  # CPU load failed — nothing to fall back to

    _cached_model_name = model_name

    # BF16 on CUDA when supported, FP32 otherwise
    if device.startswith("cuda"):
        try:
            if torch.cuda.is_bf16_supported():
                _cached_model.model = _cached_model.model.to(dtype=torch.bfloat16)
                logger.info("CrossEncoder model converted to BF16")
            else:
                logger.info("GPU does not support BF16, using FP32")
        except Exception as bf16_err:
            logger.warning(
                "BF16 conversion failed (non-fatal, continuing with FP32): %s", bf16_err
            )

    elapsed = time.time() - start
    if device.startswith("cuda") and torch.cuda.is_available():
        vram_used = torch.cuda.memory_allocated(0) / (1024**3)
        logger.info("Reranker model loaded in %.2fs on %s (VRAM used: %.2f GB)", elapsed, device, vram_used)
    else:
        logger.info("Reranker model loaded in %.2fs on %s", elapsed, device)
    return _cached_model


def rerank(query: str, passages: list[dict]) -> list[dict]:
    """Rerank passages using cross-encoder model on the best available device."""
    import os
    model_path = os.environ.get("RERANKER_MODEL_PATH", "cross-encoder/ms-marco-MiniLM-L-12-v2")
    model = get_model(model_path)

    # Truncate passages to 500 chars for efficiency
    pairs = [(query, p["text"][:500]) for p in passages]

    start = time.time()
    scores = model.predict(pairs)
    elapsed = time.time() - start
    logger.info("Reranked %d passages in %.2fs", len(passages), elapsed)

    results = []
    for passage, score in zip(passages, scores, strict=False):
        results.append(
            {
                "index": passage["index"],
                "relevance_score": float(score),
                "original_score": passage.get("original_score", 0),
            }
        )

    results.sort(key=lambda x: x["relevance_score"], reverse=True)
    return results


if __name__ == "__main__":
    try:
        input_data = json.loads(sys.stdin.read())
        query = input_data["query"]
        passages = input_data["passages"]

        if not passages:
            print(json.dumps([]))
            sys.exit(0)

        result = rerank(query, passages)
        print(json.dumps(result))
    except GPUNotAvailableError as e:
        # This should not happen — startup and get_model both fall back to CPU.
        # If we still get here, something is deeply wrong.
        logger.error("FATAL: Device error even after CPU fallback: %s", e)
        print(json.dumps({"error": str(e), "error_type": "GPU_NOT_AVAILABLE"}))
        sys.exit(1)
    except ImportError as e:
        logger.error("FATAL: Required package not installed: %s", e)
        print(json.dumps({"error": str(e), "error_type": "IMPORT_ERROR"}))
        sys.exit(1)
    except Exception as e:
        logger.error("Reranker failed: %s", str(e))
        print(json.dumps({"error": str(e), "error_type": type(e).__name__}))
        sys.exit(1)
