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 = "netandreus/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:
        {
            "inputs": {
                "source_sentence": "Your query here",
                "sentences": ["Document 1", "Document 2", ...]
            },
            "normalize": true  # Optional; defaults to False
        }
        """
        inputs = data.get("inputs", {})
        source_sentence = inputs.get("source_sentence")
        sentences = inputs.get("sentences", [])
        normalize = data.get("normalize", False)

        if not source_sentence or not sentences:
            return [{"error": "Both 'source_sentence' and 'sentences' fields are required inside 'inputs'."}]

        # Prepare input pairs
        pairs = [[source_sentence, text] for text in sentences]

        # Tokenize input pairs
        tokenizer_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(**tokenizer_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