from typing import Dict, List, Any
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch.nn.functional as F

class EndpointHandler:
    def __init__(self, path: str = "BAAI/bge-reranker-v2-m3"):
        # Load tokenizer and model
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        self.model = AutoModelForSequenceClassification.from_pretrained(path)
        self.model.eval()

        # Determine the computation device
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        Expected input format:
        {
            "query": "Your query here",
            "texts": ["Document 1", "Document 2", ...],
            "normalize": true  # Optional; defaults to False
        }
        """
        query = data.get("query")
        texts = data.get("texts", [])
        normalize = data.get("normalize", False)

        if not query or not texts:
            return [{"error": "Both 'query' and 'texts' fields are required."}]

        # Prepare input pairs
        pairs = [[query, text] for text in texts]

        # Tokenize input pairs
        inputs = self.tokenizer(
            pairs,
            padding=True,
            truncation=True,
            return_tensors="pt",
            max_length=512
        ).to(self.device)

        with torch.no_grad():
            # Get model logits
            outputs = self.model(**inputs)
            scores = outputs.logits.view(-1)

            # Apply sigmoid normalization if requested
            if normalize:
                scores = torch.sigmoid(scores)

        # Prepare the response
        results = [
            {"index": idx, "score": score.item()}
            for idx, score in enumerate(scores)
        ]

        # Sort results by descending score
        results.sort(key=lambda x: x["score"], reverse=True)
        return results