from typing import Dict, List, Any
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import re

class EndpointHandler:
    def __init__(self, path=""):
        self.model = AutoModelForCausalLM.from_pretrained("EleutherAI/pythia-410m")
        self.tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-410m")
        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]]:
        inputs = data.pop("inputs", data).strip()
        input_tokens = self.tokenizer(inputs, return_tensors="pt").to(self.device)
        with torch.no_grad():
            outputs = self.model(**input_tokens)
            logits = outputs.logits[:, -1, :]
            probs = torch.softmax(logits, dim=-1)
            top_k = torch.topk(probs, k=100)  # Check top 100 tokens
        top_predictions = []
        input_ids = input_tokens["input_ids"]
        for idx, prob in zip(top_k.indices[0], top_k.values[0]):
            # Append token to input and decode
            next_token_id = idx.item()
            test_ids = torch.cat([input_ids, torch.tensor([[next_token_id]], device=self.device)], dim=-1)
            text = self.tokenizer.decode(test_ids[0], skip_special_tokens=True).strip()
            # Extract last word
            last_word = text.split()[-1]
            # Filter for meaningful words (3+ letters, alphabetic)
            if re.match(r"^[a-zA-Z]{3,}$", last_word) and last_word not in [p["text"] for p in top_predictions]:
                top_predictions.append({"text": last_word, "probability": prob.item()})
            if len(top_predictions) >= 5:
                break
        return top_predictions if top_predictions else [{"text": "No valid words found", "probability": 0.0}]