import gradio as gr
from transformers import pipeline
import os

# 1. YOUR CUSTOM DICTIONARY (Add more words here as you find them)
# Format: "English/Spanish Word": "Correct Garifuna"
CUSTOM_MAPPINGS = {
    "morbid": "wiyeti samina",
    "mórbido": "wiyeti samina",
    "punching bag": "düsü saki", # Example fix
    "hola": "hürüti",
}

hf_token = os.environ.get("HF_TOKEN")
model_id = "arnoldramo/garifuna-nllb"

translator = pipeline("translation", model=model_id, token=hf_token)

def respond(message, history, source_lang):
    # Normalize the input (lowercase and remove extra spaces)
    clean_message = message.lower().strip()
    
    # 2. CHECK CUSTOM DICTIONARY FIRST
    if clean_message in CUSTOM_MAPPINGS:
        return CUSTOM_MAPPINGS[clean_message]
    
    # 3. OTHERWISE, ASK THE AI
    lang_map = {"English": "eng_Latn", "Spanish": "spa_Latn"}
    src_code = lang_map.get(source_lang, "eng_Latn")
    
    result = translator(message, src_lang=src_code, tgt_lang="cab_Latn")
    return result[0]['translation_text']

demo = gr.ChatInterface(
    fn=respond,
    title="Garifuna Language Assistant",
    additional_inputs=[
        gr.Dropdown(label="Source Language", choices=["English", "Spanish"], value="English")
    ]
)

if __name__ == "__main__":
    demo.launch()