import torch
import gradio as gr
from transformers import pipeline

# Use a pipeline as a high-level helper
from transformers import pipeline
# model_path = ("../Models/models--Helsinki-NLP--opus-mt-en-de/snapshots"
#               "/6183067f769a302e3861815543b9f312c71b0ca4")

pipe = pipeline("translation", model="Helsinki-NLP/opus-mt-en-de")
# text_translator = pipeline("translation", model=model_path, tokenizer=model_path)

def translate_text(text, destination_language):
    # German uses your preloaded local snapshot
    if destination_language == "German":
        out = text_translator(text)
        return out[0]["translation_text"]

    # For other targets we load the correct Marian model on demand (EN -> target)
    if destination_language == "French":
        p = pipeline("translation", model="Helsinki-NLP/opus-mt-en-fr")
    elif destination_language == "Hindi":
        p = pipeline("translation", model="Helsinki-NLP/opus-mt-en-hi")
    elif destination_language == "Romanian":
        p = pipeline("translation", model="Helsinki-NLP/opus-mt-en-ro")
    elif destination_language == "Spanish":
        p = pipeline("translation", model="Helsinki-NLP/opus-mt-en-es")
    elif destination_language == "Italian":
        p = pipeline("translation", model="Helsinki-NLP/opus-mt-en-it")
    elif destination_language == "Portuguese":
        p = pipeline("translation", model="Helsinki-NLP/opus-mt-en-pt")
    elif destination_language == "Russian":
        p = pipeline("translation", model="Helsinki-NLP/opus-mt-en-ru")
    elif destination_language == "Japanese":
        p = pipeline("translation", model="Helsinki-NLP/opus-mt-en-ja")
    elif destination_language == "Korean":
        p = pipeline("translation", model="Helsinki-NLP/opus-mt-en-ko")
    elif destination_language == "Chinese":
        p = pipeline("translation", model="Helsinki-NLP/opus-mt-en-zh")
    elif destination_language == "Arabic":
        p = pipeline("translation", model="Helsinki-NLP/opus-mt-en-ar")
    elif destination_language == "Turkish":
        p = pipeline("translation", model="Helsinki-NLP/opus-mt-en-tr")
    elif destination_language == "Dutch":
        p = pipeline("translation", model="Helsinki-NLP/opus-mt-en-nl")
    elif destination_language == "Polish":
        p = pipeline("translation", model="Helsinki-NLP/opus-mt-en-pl")
    elif destination_language == "Ukrainian":
        p = pipeline("translation", model="Helsinki-NLP/opus-mt-en-uk")
    elif destination_language == "Czech":
        p = pipeline("translation", model="Helsinki-NLP/opus-mt-en-cs")
    elif destination_language == "Swedish":
        p = pipeline("translation", model="Helsinki-NLP/opus-mt-en-sv")
    else:
        return "Unsupported language. Please choose a listed destination."

    out = p(text)  # Marian pipeline returns a list of dicts
    return out[0]["translation_text"]

gr.close_all()
demo = gr.Interface(
    fn=translate_text,
    inputs=[
        gr.Textbox(label="Input text to translate", lines=6),
        gr.Dropdown(
            [
                "German", "French", "Hindi", "Romanian", "Spanish", "Italian",
                "Portuguese", "Russian", "Japanese", "Korean", "Chinese",
                "Arabic", "Turkish", "Dutch", "Polish", "Ukrainian", "Czech", "Swedish"
            ],
            label="Select Destination Language"
        )
    ],
    outputs=[gr.Textbox(label="Translated text", lines=4)],
    title="@SahibhimGenAI Project 4: Multi language translator",
    description="Translate English text to your selected language (loads the appropriate MarianMT model per language)."
)
demo.launch()