## translation layer for Farmlingua (Milestone2)

import nest_asyncio
import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel
import requests
import os
from models.model import translator


nest_asyncio.apply()

lang_codes = {
    "en": "eng_Latn",
    "ig": "ibo_Latn",
    "ha": "hau_Latn",
    "yo": "yor_Latn",
    "sw": "swh_Latn"
}

# language identifier

app = FastAPI(
    title="Farmlingua Translation API",
    description="Multilingual ability for farmlingua supporting Igbo, Hausa, Yoruba, and Swahili",
    version="0.0.1"
)

class Query(BaseModel):
    text: str
    lang: str  

def translate(text: str, src: str, tgt: str) -> str:
    try:
        out = translator(text, src_lang=lang_codes[src], tgt_lang=lang_codes[tgt], max_length=400)
        return out[0]["translation_text"]
    except Exception as e:
        return f"Translation error: {str(e)}"

def query_milestone1(prompt: str) -> str:
    try:
        response = requests.post(
            "https://remostart-milestone-one-farmlingua-ai.hf.space/ask",
            json={"query": prompt},
            timeout=200
        )
        if response.status_code == 200:
            data = response.json()
            return data.get("answer") or str(data)
        else:
            return f"Base model error ({response.status_code}): {response.text}"
    except Exception as e:
        return f"Failed to connect to Milestone 1: {str(e)}"

# Root endpoint
@app.get("/")
def root():
    return {
        "message": "Welcome to Farmlingua Translation API",
        "version": "3.0",
        "features": {
            "simple_language_identifier": "Specify language with 'lang' parameter",
            "single_endpoint": "/ask - Send text with language identifier",
            "supported_codes": "en, ig, ha, yo, sw"
        },
        "supported_languages": {
            "en": "English",
            "ig": "Igbo", 
            "ha": "Hausa",
            "yo": "Yoruba",
            "sw": "Swahili"
        },
        "documentation": "/docs"
    }

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

@app.post("/ask")
def multilingual_ask(query: Query):
    """Multilingual endpoint with simple language identifier"""
    user_text = query.text
    user_lang = query.lang.lower()
    
    # Validate language code
    if user_lang not in lang_codes:
        return {"error": "Unsupported language. Use one of: en, ig, ha, yo, sw."}
    
    # language name for response
    lang_names = {
        "en": "English",
        "ig": "Igbo", 
        "ha": "Hausa",
        "yo": "Yoruba",
        "sw": "Swahili"
    }
    
    # Translate to English if not English
    if user_lang != "en":
        english_text = translate(user_text, user_lang, "en")
    else:
        english_text = user_text
    
    # Get response from Milestone 1
    english_response = query_milestone1(english_text)
    
    # Translate response back to original language
    if user_lang != "en":
        final_response = translate(english_response, "en", user_lang)
    else:
        final_response = english_response
    
    return {
        #"language": lang_names[user_lang],
        #"language_code": user_lang,
        #"original_text": user_text,
        #"english_translation": english_text,
        "response": final_response
    }

# Run FastAPI server
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)