import gradio as gr
from gtts import gTTS
import os
from PIL import Image
from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer, BlipProcessor, BlipForConditionalGeneration

# Define translation models per language pair
translation_models = {
    "French": "Helsinki-NLP/opus-mt-en-fr",
    "Spanish": "Helsinki-NLP/opus-mt-en-es",
    "Mandarin": "Helsinki-NLP/opus-mt-en-zh",
    "Vietnamese": "Helsinki-NLP/opus-mt-en-vi",
    "Tagalog": "Helsinki-NLP/opus-mt-en-tl"
}

# Directory to store local models
local_model_dir = "./models"
os.makedirs(local_model_dir, exist_ok=True)

# Load translation models and tokenizers locally
translation_pipelines = {}
for lang, model_name in translation_models.items():
    local_path = os.path.join(local_model_dir, model_name.replace('/', '_'))

    if not os.path.exists(local_path):
        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model.save_pretrained(local_path)
        tokenizer.save_pretrained(local_path)
    else:
        model = AutoModelForSeq2SeqLM.from_pretrained(local_path)
        tokenizer = AutoTokenizer.from_pretrained(local_path)

    translation_pipelines[lang] = pipeline("translation", model=model, tokenizer=tokenizer)

# Load image captioning model
caption_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
caption_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

def generate_caption(image):
    inputs = caption_processor(image, return_tensors="pt")
    output = caption_model.generate(**inputs)
    caption = caption_processor.decode(output[0], skip_special_tokens=True)
    return caption

# Function to translate text using locally stored models
def translate_text(sentence, target_lang):
    if target_lang in translation_pipelines:
        model = translation_pipelines[target_lang]
        translation = model(sentence, max_length=100, truncation=True)
        return translation[0]['translation_text']
    return sentence  # Return original if no model found

# Language options
language_options = {
    "English": "en", "French": "fr", "Spanish": "es", "Mandarin": "zh", "Vietnamese": "vi", "Tagalog": "tl"
}

# Main function for Gradio
def process_image(image, l1, l2):
    if not image:
        return "Please upload an image.", None, None

    caption_en = generate_caption(image)
    caption_l2 = translate_text(caption_en, l2) if l2 != "English" else caption_en

    # Generate TTS
    tts = gTTS(text=caption_l2, lang=language_options[l2])
    audio_path = "caption_audio.mp3"
    tts.save(audio_path)

    return caption_l2, audio_path, caption_en

# Gradio Interface
interface = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="pil", sources=["upload"]),  # Only allows upload and drag-and-drop
        gr.Dropdown(list(language_options.keys()), label="Your Native Language (L1)"),
        gr.Dropdown(list(language_options.keys()), label="Target Language (L2)")
    ],
    outputs=[
        gr.Textbox(label="Generated Caption (L2)"),
        gr.Audio(label="Listen to the Caption"),
        gr.Textbox(label="Generated Caption (L1)")
    ],
    title="Image Caption Generator for Language Learning",
    description="Upload an image, select your native language (L1) and target language (L2), and get a generated caption with text-to-speech audio."
)

# Launch Gradio
interface.launch()