"""
Ray Serve deployment for sentence-transformers CrossEncoder reranking.
Runs on: drizzt (Radeon 680M iGPU, ROCm) or danilo (Intel i915 iGPU, OpenVINO/IPEX)
"""

import os
import time
import uuid
from typing import Any, Dict, List, Tuple

from ray import serve


@serve.deployment(name="RerankerDeployment", num_replicas=1)
class RerankerDeployment:
    def __init__(self):
        from sentence_transformers import CrossEncoder
        import torch
        
        self.model_id = os.environ.get("MODEL_ID", "BAAI/bge-reranker-v2-m3")
        self.use_ipex = False
        self.device = "cpu"
        
        # Detect device - check for Intel GPU first via IPEX
        try:
            import intel_extension_for_pytorch as ipex
            self.use_ipex = True
            if hasattr(torch, "xpu") and torch.xpu.is_available():
                self.device = "xpu"
                print("Intel GPU detected via IPEX, using XPU device")
            else:
                print("IPEX available, will use CPU optimization")
        except ImportError:
            print("IPEX not available, checking for other GPUs")
        
        # Check for CUDA/ROCm if not using Intel
        if not self.use_ipex:
            if torch.cuda.is_available():
                self.device = "cuda"
                print(f"Using CUDA/ROCm device")
            else:
                print("No GPU detected, using CPU")
        
        print(f"Loading reranker model: {self.model_id}")
        print(f"Using device: {self.device}")
        
        # Load model
        self.model = CrossEncoder(self.model_id, device=self.device)
        
        # Apply IPEX optimization if available
        if self.use_ipex and self.device == "cpu":
            try:
                import intel_extension_for_pytorch as ipex
                self.model.model = ipex.optimize(self.model.model)
                print("IPEX CPU optimization applied")
            except Exception as e:
                print(f"IPEX optimization failed: {e}")
        
        print(f"Reranker model loaded successfully")

    async def __call__(self, request: Dict[str, Any]) -> Dict[str, Any]:
        """
        Handle reranking requests.
        
        Expected request format:
        {
            "query": "search query",
            "documents": ["doc1", "doc2", "doc3"],
            "top_k": 3,
            "return_documents": true
        }
        
        Alternative format (pairs):
        {
            "pairs": [["query", "doc1"], ["query", "doc2"]]
        }
        """
        # Handle pairs format
        if "pairs" in request:
            pairs = request["pairs"]
            scores = self.model.predict(pairs)
            
            results = []
            for i, (pair, score) in enumerate(zip(pairs, scores)):
                results.append({
                    "index": i,
                    "score": float(score),
                })
            
            return {
                "object": "list",
                "results": results,
                "model": self.model_id,
            }
        
        # Handle query + documents format
        query = request.get("query", "")
        documents = request.get("documents", [])
        top_k = request.get("top_k", len(documents))
        return_documents = request.get("return_documents", True)
        
        if not documents:
            return {
                "object": "list",
                "results": [],
                "model": self.model_id,
            }
        
        # Create query-document pairs
        pairs = [[query, doc] for doc in documents]
        
        # Get scores
        scores = self.model.predict(pairs)
        
        # Create results with indices and scores
        results = []
        for i, (doc, score) in enumerate(zip(documents, scores)):
            result = {
                "index": i,
                "score": float(score),
            }
            if return_documents:
                result["document"] = doc
            results.append(result)
        
        # Sort by score descending
        results.sort(key=lambda x: x["score"], reverse=True)
        
        # Apply top_k
        results = results[:top_k]
        
        return {
            "object": "list",
            "results": results,
            "model": self.model_id,
            "usage": {
                "total_pairs": len(pairs),
            },
        }


app = RerankerDeployment.bind()