import requests
import gradio as gr
from dotenv import load_dotenv
import os

# Load environment variables
load_dotenv()
HF_TOKEN = os.getenv("HF_TOKEN")
headers = {"Authorization": f"Bearer {HF_TOKEN}"}

# Language to ISO 639-3 codes (used for NLLB-200)
LANGUAGES = {
    "English → Afrikaans": "afr",
    "English → Xhosa": "xho",
    "English → Zulu": "zul",
    "English → Sesotho": "sot",
    "English → Tswana": "tsn",
    "English → Northern Sotho": "nso",
    "English → Swati": "ssw",
    "English → Tsonga": "tso",
    "English → Venda": "ven",
}

MODEL_NAME = "facebook/nllb-200-distilled-600M"
API_URL = f"https://api-inference.huggingface.co/models/{MODEL_NAME}"


def query(payload):
    response = requests.post(API_URL, headers=headers, json=payload)

    if response.status_code != 200:
        print(f"[ERROR] API failed: {response.status_code} - {response.text}")
        return {"error": f"Request failed with {response.status_code}"}

    try:
        return response.json()
    except requests.exceptions.JSONDecodeError:
        print(f"[ERROR] Failed to parse JSON: {response.text}")
        return {"error": "Invalid JSON from API"}


def translate(input_text, language_label):
    language_code = LANGUAGES[language_label]
    formatted_input = f">>{language_code}<< {input_text}"

    response = query({"inputs": formatted_input, "options": {"wait_for_model": True}})

    if "error" in response:
        return f"Error: {response['error']}"

    return response[0]["translation_text"]


translator = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(label="Input Text", placeholder="Type text here..."),
        gr.Dropdown(list(LANGUAGES.keys()), label="Select Language Target"),
    ],
    outputs=gr.Textbox(label="Translation"),
    title="Translademia",
    description="Translate English text to South African languages using Meta's NLLB-200 model.",
)

translator.launch()