import pickle
import json
import torch
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import BertTokenizer
import io

# 1. Initialize App and Load Assets
# ==================================
app = FastAPI(
    title="TEXI'OR Sentiment Analysis API",
    description="An API to classify text into different sentiments using a fine-tuned BERT model.",
    version="1.0"
)
class TextInput(BaseModel):
    text: str

device = torch.device('cpu')
print(f"✅ Using device: {device}")

# Custom Unpickler to load a GPU-trained model onto a CPU
class CPU_Unpickler(pickle.Unpickler):
    def find_class(self, module, name):
        if module == 'torch.storage' and name == '_load_from_bytes':
            return lambda b: torch.load(io.BytesIO(b), map_location='cpu')
        else:
            return super().find_class(module, name)

# Load the fine-tuned model and tokenizer
try:
    with open('bert_sentiment_model.pkl', 'rb') as f:
        model = CPU_Unpickler(f).load() 
    model.to(device)
    model.eval()
    print("✅ Model 'bert_sentiment_model.pkl' loaded successfully onto CPU.")
except FileNotFoundError:
    print("❌ Model file not found. Make sure 'bert_sentiment_model.pkl' is uploaded.")
    model = None

try:
    with open('Label_Dict.json', 'r') as f:
        label_dict = json.load(f)
    idx2label = {int(v): k for k, v in label_dict.items()}
    print("✅ Label dictionary 'Label_Dict.json' loaded successfully.")
except FileNotFoundError:
    print("❌ Label dictionary 'Label_Dict.json' not found.")
    idx2label = None

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
print("✅ Tokenizer loaded successfully.")


# 2. Define API Endpoints
# ========================
@app.get("/")
def read_root():
    return {"message": "Welcome to the TEXI'OR API. Use the /predict endpoint to get a sentiment."}

@app.post("/predict")
def predict_sentiment(text_input: TextInput):
    if not all([model, idx2label, tokenizer]):
        return {"error": "API is not ready. A model, label, or tokenizer component failed to load."}

    text = text_input.text

    inputs = tokenizer(
        text, padding=True, truncation=True, max_length=150, return_tensors='pt'
    )

    input_ids = inputs['input_ids'].to(device)
    attention_mask = inputs['attention_mask'].to(device)

    with torch.no_grad():
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)

    logits = outputs.logits
    probabilities = torch.nn.functional.softmax(logits, dim=-1)
    confidence, predicted_class_idx = torch.max(probabilities, dim=1)

    predicted_label = idx2label.get(predicted_class_idx.item(), "Unknown Label")
    confidence_score = confidence.item()

    return {
        "sentiment": predicted_label,
        "confidence": round(confidence_score, 4)
    }