import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline
import asyncio

# Initialize FastAPI
app = FastAPI()

# Load Hugging Face translation pipeline globally
print("🚀 Loading translation pipeline...")
translator = pipeline("translation", model="facebook/m2m100_1.2B")
print("✅ Model is ready!")

# Request model for batch translation
class BatchTranslationRequest(BaseModel):
    texts: list[str]
    source_lang: str
    target_lang: str

@app.post("/translate")
async def translate(request: BatchTranslationRequest):
    try:
        # Prepare input batch
        batch_inputs = [f"{text} ({request.source_lang})" for text in request.texts]
        
        # Run translation asynchronously
        translated_texts = await asyncio.to_thread(
            translator, batch_inputs, src_lang=request.source_lang, tgt_lang=request.target_lang
        )
        
        # Extract translated text results
        translated_results = [result['translation_text'] for result in translated_texts]
        
        return {"translated_texts": translated_results}
    
    except Exception as e:
        return {"error": str(e)}

@app.get("/")
async def root():
    return {"message": "🚀 M2M100 Translation API is running!"}

@app.get("/health")
async def health_check():
    return {"status": "healthy"}