from typing import Dict, List
import torch
import torch.nn as nn
from transformers import BertModel, BertTokenizer

class SentimentBERT(nn.Module):
    def __init__(self, pretrained_model_name="bert-base-uncased"):
        super(SentimentBERT, self).__init__()
        self.bert_model = BertModel.from_pretrained(pretrained_model_name)
        self.tokenizer = BertTokenizer.from_pretrained(pretrained_model_name)
        self.dropout = nn.Dropout(0.3)
        self.fc = nn.Linear(self.bert_model.config.hidden_size, 2)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert_model(input_ids=input_ids, attention_mask=attention_mask)
        pooler_output = outputs.pooler_output
        output = self.dropout(pooler_output)
        output = self.fc(output)
        return output

class EndpointHandler:
    def __init__(self, path=""):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model = SentimentBERT().to(self.device)
        
        # Load the model state dict
        if path:
            self.model.load_state_dict(torch.load(path, map_location=self.device))
        self.model.eval()
        
        self.tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
        
    def __call__(self, data: Dict) -> Dict:
        """
        Args:
            data: Dictionary with "inputs" key containing text for sentiment analysis
        Returns:
            Dictionary with "label" and "score" keys
        """
        # Get the input text
        inputs = data.pop("inputs", data)
        
        # Ensure the input is a list
        if isinstance(inputs, str):
            inputs = [inputs]
            
        # Tokenize the input texts
        encoded = self.tokenizer(
            inputs,
            truncation=True,
            padding=True,
            max_length=512,
            return_tensors="pt"
        ).to(self.device)
        
        # Get model predictions
        with torch.no_grad():
            outputs = self.model(**encoded)
            predictions = torch.softmax(outputs, dim=1)
            
        # Process predictions
        results = []
        for pred in predictions:
            label_id = pred.argmax().item()
            score = pred[label_id].item()
            label = "Positive" if label_id == 1 else "Negative"
            results.append({
                "label": label,
                "score": round(score, 4)
            })
            
        return {"predictions": results}