"""FastAPI application for Myanmar Ghost model."""

import logging
from pathlib import Path
from typing import Any, Dict, List, Optional

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import torch

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

app = FastAPI(
    title="Myanmar Ghost API",
    description="Advanced Myanmar Language Understanding Model",
    version="1.0.0",
)

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

# Global model reference
model = None
tokenizer = None


class TextInput(BaseModel):
    text: str = Field(..., description="Myanmar text to analyze")
    include_prosody: bool = Field(False, description="Include prosody features")


class SentimentResponse(BaseModel):
    text: str
    sentiment: str
    confidence: float
    probabilities: Dict[str, float]


class BatchTextInput(BaseModel):
    texts: List[str] = Field(..., description="List of Myanmar texts")


class BatchSentimentResponse(BaseModel):
    results: List[SentimentResponse]


@app.on_event("startup")
async def startup_event():
    """Load model on startup."""
    global model, tokenizer
    
    logger.info("Loading Myanmar Ghost model...")
    
    try:
        from transformers import AutoModelForSequenceClassification, AutoTokenizer
        
        model_name = "amkyawdev/Myanmar-Ghost-Instruct"
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForSequenceClassification.from_pretrained(model_name)
        model.eval()
        
        logger.info(f"Model loaded: {model_name}")
    except Exception as e:
        logger.warning(f"Could not load model from HuggingFace: {e}")
        logger.info("Using placeholder for demonstration")


@app.get("/")
async def root():
    """Root endpoint."""
    return {
        "name": "Myanmar Ghost API",
        "version": "1.0.0",
        "status": "online",
    }


@app.get("/health")
async def health():
    """Health check endpoint."""
    return {
        "status": "healthy",
        "model_loaded": model is not None,
    }


@app.post("/predict", response_model=SentimentResponse)
async def predict(input_data: TextInput) -> SentimentResponse:
    """Predict sentiment for a single text."""
    if model is None or tokenizer is None:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    try:
        # Tokenize
        inputs = tokenizer(
            input_data.text,
            return_tensors="pt",
            truncation=True,
            max_length=512,
        )
        
        # Predict
        with torch.no_grad():
            outputs = model(**inputs)
            probs = torch.softmax(outputs.logits, dim=-1)[0]
        
        # Get prediction
        sentiment_idx = probs.argmax().item()
        confidence = probs[sentiment_idx].item()
        
        sentiment_labels = ["negative", "neutral", "positive", "sarcastic"]
        sentiment = sentiment_labels[sentiment_idx]
        
        probabilities = {
            label: probs[i].item()
            for i, label in enumerate(sentiment_labels)
        }
        
        return SentimentResponse(
            text=input_data.text,
            sentiment=sentiment,
            confidence=confidence,
            probabilities=probabilities,
        )
    
    except Exception as e:
        logger.error(f"Prediction error: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/predict_batch", response_model=BatchSentimentResponse)
async def predict_batch(input_data: BatchTextInput) -> BatchSentimentResponse:
    """Predict sentiment for multiple texts."""
    if model is None or tokenizer is None:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    results = []
    
    try:
        for text in input_data.texts:
            # Tokenize
            inputs = tokenizer(
                text,
                return_tensors="pt",
                truncation=True,
                max_length=512,
            )
            
            # Predict
            with torch.no_grad():
                outputs = model(**inputs)
                probs = torch.softmax(outputs.logits, dim=-1)[0]
            
            # Get prediction
            sentiment_idx = probs.argmax().item()
            confidence = probs[sentiment_idx].item()
            
            sentiment_labels = ["negative", "neutral", "positive", "sarcastic"]
            sentiment = sentiment_labels[sentiment_idx]
            
            probabilities = {
                label: probs[i].item()
                for i, label in enumerate(sentiment_labels)
            }
            
            results.append(SentimentResponse(
                text=text,
                sentiment=sentiment,
                confidence=confidence,
                probabilities=probabilities,
            ))
        
        return BatchSentimentResponse(results=results)
    
    except Exception as e:
        logger.error(f"Batch prediction error: {e}")
        raise HTTPException(status_code=500, detail=str(e))


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)