from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(
    title="Support Ticket Classifier API",
    description="Classify customer support tickets into 5 categories using BERTweet",
    version="1.0.0",
    docs_url="/docs",
    redoc_url="/redoc"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

logger.info("Loading support ticket classifier model...")
try:
    model = AutoModelForSequenceClassification.from_pretrained("Adedayo2000/BERTweet-Support-Ticket-Classifier")
    tokenizer = AutoTokenizer.from_pretrained("Adedayo2000/BERTweet-Support-Ticket-Classifier")
    logger.info("Model loaded successfully!")
except Exception as e:
    logger.error(f"Failed to load model: {e}")
    model = None
    tokenizer = None

class TicketInput(BaseModel):
    text: str

    class Config:
        json_schema_extra = {
            "example": {
                "text": "I can't login to my account, password reset not working"
            }
        }

@app.get("/")
def home():
    return {
        "message": "Support Ticket Classifier API",
        "status": "live",
        "model": "Adedayo2000/BERTweet-Support-Ticket-Classifier",
        "accuracy": "85%",
        "categories": [
            "Account Access",
            "Billing/Payment",
            "Feature Request",
            "General Inquiry",
            "Technical Support"
        ],
        "endpoints": {
            "GET /": "API information",
            "POST /classify": "Classify support ticket",
            "GET /categories": "List all categories",
            "GET /docs": "Interactive API documentation",
            "GET /health": "Health check"
        }
    }

@app.get("/health")
def health_check():
    return {
        "status": "healthy",
        "model_loaded": model is not None and tokenizer is not None
    }

@app.get("/categories")
def get_categories():
    if model is None:
        return {"error": "Model not loaded"}
    
    return {
        "categories": list(model.config.id2label.values()),
        "count": len(model.config.id2label)
    }

@app.post("/classify")
def classify_ticket(input: TicketInput):
    if not input.text or not input.text.strip():
        return {
            "error": "Ticket text cannot be empty",
            "example": {"text": "I can't login to my account"}
        }
    
    if model is None or tokenizer is None:
        return {
            "error": "Model not loaded",
            "status": "unavailable"
        }
    
    try:
        inputs = tokenizer(
            input.text,
            return_tensors='pt',
            truncation=True,
            max_length=128,
            padding=True
        )
        
        with torch.no_grad():
            outputs = model(**inputs)
        
        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
        pred_idx = torch.argmax(probs, dim=-1).item()
        
        category = model.config.id2label[pred_idx]
        confidence_score = probs[0][pred_idx].item()
        confidence_percentage = f"{confidence_score * 100:.2f}%"
        
        all_predictions = {
            model.config.id2label[i]: f"{probs[0][i].item() * 100:.2f}%"
            for i in range(len(probs[0]))
        }
        
        return {
            "text": input.text,
            "category": category,
            "confidence": confidence_percentage,
            "confidence_score": round(confidence_score, 4),
            "all_categories": all_predictions
        }
        
    except Exception as e:
        logger.error(f"Classification error: {e}")
        return {
            "error": "Classification failed",
            "details": str(e)
        }