# translate_api.py
# Defines the FastAPI application for the Chiakamo Translator.

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import MarianMTModel, MarianTokenizer, pipeline
import torch
import os # Import os for environment variables (e.g., HF_HOME)

# Initialize the FastAPI application
app = FastAPI(title="ChikaMo Translator API",
              description="API for translating between Tagalog and English, with fallback to Helsinki-NLP models.")

# --- Model Configuration ---
# Define the Hugging Face model repository IDs for your custom models.
# These should be your repositories on Hugging Face Hub (e.g., chikamov1/opus-mt-tl-en-chikamo)
LOCAL_MODELS = {
    "tl-en": "chikamov1/opus-mt-tl-en-chikamo", # Your custom Tagalog-to-English model
    "en-tl": "chikamov1/opus-mt-en-tl-chikamo", # Your custom English-to-Tagalog model
}

# Define HuggingFace fallback models (Helsinki-NLP is a good choice)
FALLBACK_MODELS = {
    "tl-en": "Helsinki-NLP/opus-mt-tl-en",
    "en-tl": "Helsinki-NLP/opus-mt-en-tl",
}

# Dictionaries to store loaded models and tokenizers to avoid reloading on every request
loaded_models = {}
loaded_tokenizers = {}
fallback_pipelines = {}

# --- Model Loading Functions ---

# Function to get custom model and tokenizer (loads and caches)
def get_model_and_tokenizer(pair: str):
    """
    Loads and caches the custom MarianMT model and tokenizer for a given language pair.
    """
    if pair in LOCAL_MODELS:
        if pair not in loaded_models:
            print(f"Attempting to load local model: {LOCAL_MODELS[pair]}")
            try:
                # Load model and tokenizer from your Hugging Face Hub repository
                model = MarianMTModel.from_pretrained(LOCAL_MODELS[pair])
                tokenizer = MarianTokenizer.from_pretrained(LOCAL_MODELS[pair])
                loaded_models[pair] = model
                loaded_tokenizers[pair] = tokenizer
                print(f"Successfully loaded local model: {LOCAL_MODELS[pair]}")
            except Exception as e:
                print(f"Failed to load local model {LOCAL_MODELS[pair]}: {e}")
                return None, None # Return None if loading fails
        return loaded_models.get(pair), loaded_tokenizers.get(pair)
    return None, None

# Function to get fallback pipeline (loads and caches)
def get_fallback_pipeline(pair: str):
    """
    Loads and caches a fallback translation pipeline for a given language pair.
    """
    if pair not in fallback_pipelines:
        print(f"Attempting to load fallback model: {FALLBACK_MODELS[pair]}")
        try:
            # Use the pipeline abstraction for fallback models
            pipe = pipeline("translation", model=FALLBACK_MODELS[pair])
            fallback_pipelines[pair] = pipe
            print(f"Successfully loaded fallback model: {FALLBACK_MODELS[pair]}")
        except Exception as e:
            print(f"Failed to load fallback model {FALLBACK_MODELS[pair]}: {e}")
            return None # Return None if loading fails
    return fallback_pipelines.get(pair)

# --- Pydantic Model for Request Body ---
class TranslationRequest(BaseModel):
    """
    Defines the structure of the incoming JSON request for translation.
    """
    source_lang: str # e.g., "tl" for Tagalog, "en" for English
    target_lang: str # e.g., "en" for English, "tl" for Tagalog
    text: str        # The text to be translated

# --- API Endpoints ---

# Root endpoint for basic API information
@app.get("/")
def read_root():
    """
    Returns a welcome message and API status.
    """
    return {"message": "Welcome to ChikaMo Translator API", "status": "running"}

# Translation endpoint
@app.post("/translate")
def translate_text_endpoint(req: TranslationRequest):
    """
    Translates text between specified source and target languages.
    Prioritizes custom models, falls back to Helsinki-NLP if custom fails or is not found.
    """
    pair = f"{req.source_lang.lower()}-{req.target_lang.lower()}"
    text = req.text.strip()

    if not text:
        raise HTTPException(status_code=400, detail="Input text is empty.")

    translated_output = ""
    fallback_used = False

    # Try to use local (custom) model first
    model, tokenizer = get_model_and_tokenizer(pair)
    if model and tokenizer:
        try:
            inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
            # Ensure model is on CPU if no GPU is available, or move to GPU if present
            device = "cuda" if torch.cuda.is_available() else "cpu"
            model.to(device)
            inputs = {k: v.to(device) for k, v in inputs.items()}

            with torch.no_grad():
                translated = model.generate(**inputs)
            translated_output = tokenizer.decode(translated[0], skip_special_tokens=True)
            fallback_used = False
        except Exception as e:
            print(f"Error during custom model translation for pair {pair}: {e}. Attempting fallback.")
            # Clear to force fallback path
            model, tokenizer = None, None
    
    # If custom model failed or wasn't available, try fallback
    if not model and pair in FALLBACK_MODELS:
        pipe = get_fallback_pipeline(pair)
        if pipe:
            try:
                translated_output = pipe(text)[0]["translation_text"]
                fallback_used = True
            except Exception as e:
                print(f"Error during fallback model translation for pair {pair}: {e}.")
                raise HTTPException(status_code=500, detail=f"Translation failed for pair {pair} with both custom and fallback models.")
        else:
            raise HTTPException(status_code=500, detail=f"Fallback model for pair {pair} could not be loaded.")
    elif not model: # No custom model and no fallback defined
        raise HTTPException(status_code=400, detail=f"Unsupported language pair: {pair}, and no fallback model configured.")


    return {
        "translation": translated_output,
        "source_lang": req.source_lang,
        "target_lang": req.target_lang,
        "fallback_used": fallback_used
    }

# --- Application Startup Event ---
# This ensures models are loaded when the FastAPI app starts up.
# This is crucial for Hugging Face Spaces where the app is started by Uvicorn.
@app.on_event("startup")
async def startup_event():
    # Pre-load all local and fallback models during startup
    print("Pre-loading models during application startup...")
    for pair in LOCAL_MODELS.keys():
        get_model_and_tokenizer(pair) # Attempt to load custom models
    for pair in FALLBACK_MODELS.keys():
        get_fallback_pipeline(pair) # Attempt to load fallback pipelines
    print("Model pre-loading complete.")