from fastapi import FastAPI
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware
from transformers import pipeline
import os

# -------------------------------
# Initialize FastAPI app
# -------------------------------
app = FastAPI(title="Language Translator API")

# -------------------------------
# Enable CORS for your extension
# -------------------------------
# Replace with your Chrome extension ID for security, or use "*" for testing
app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        #"*",  # ✅ recommended
        "*"  # ⚠️ uncomment for open access during testing
    ],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# -------------------------------
# Model Loading
# -------------------------------
MODEL_PATH = "./model"

if not os.path.exists(MODEL_PATH):
    raise RuntimeError("Model folder not found! Make sure '/model' is uploaded to your Space.")

print("📁 Model files found:", os.listdir(MODEL_PATH))

# Load model from local folder (pre-downloaded)
translator = pipeline("translation", model=MODEL_PATH)

# -------------------------------
# Request/Response Models
# -------------------------------
class TranslationRequest(BaseModel):
    text: str
    src_lang: str
    tgt_lang: str

# -------------------------------
# API Route
# -------------------------------
@app.post("/translate")
def translate(req: TranslationRequest):
    """
    Translate text from source language to target language.
    """
    try:
        result = translator(
            req.text,
            src_lang=req.src_lang,
            tgt_lang=req.tgt_lang
        )
        return {"translated_text": result[0]['translation_text']}
    except Exception as e:
        return {"error": str(e)}

# -------------------------------
# Root route
# -------------------------------
@app.get("/")
def root():
    return {"message": "✅ Language Translator API is running successfully!"}