import gradio as gr

# import from your src package
from src.inference import load_model_and_vocab, transliterate_sentence

# load model & vocab once at startup
model, src_char2idx, tgt_char2idx, idx2tgt_char, word_map = load_model_and_vocab()


def respond(message: str, history):
    """
    Gradio chat callback.
    message: current user input (Teluguish / mixed)
    history: previous chat turns (ignored for transliteration)
    Returns: Telugu script string
    """
    message = message.strip()
    if not message:
        return ""

    telugu = transliterate_sentence(
        message,
        model,
        src_char2idx,
        idx2tgt_char,
        word_map,
    )
    return telugu


chatbot = gr.ChatInterface(
    fn=respond,
    type="tuples",  # message is a plain string
    title="Teluguish ➜ Telugu Transliteration",
    description=(
        "Type Telugu in English letters (Teluguish) or mixed Telugu/English, "
        "and the model will transliterate it into proper Telugu script."
    ),
    examples=[
        ["ela unnavu mawa"],
        ["nuvvu emi chestunnavu"],
        ["manamu emi chestunnamu"],
        ["trivikram chala manchi director"],
        ["మీరు emi chestunnaru"],
    ],
)

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