# app.py
import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
from diffusers import StableDiffusionPipeline
import torch
import re

# Initialize models
def load_models():
    # Load translation model (NLLB - best for Tamil)
    translator = pipeline(
        "translation",
        model="facebook/nllb-200-distilled-600M",
        src_lang="tam_Taml",
        tgt_lang="eng_Latn",
        device=0 if torch.cuda.is_available() else -1
    )
    
    # Text generation pipeline
    text_generator = pipeline(
        "text-generation", 
        model="gpt2-medium",
        device=0 if torch.cuda.is_available() else -1
    )
    
    # Image generation pipeline
    if torch.cuda.is_available():
        image_pipe = StableDiffusionPipeline.from_pretrained(
            "runwayml/stable-diffusion-v1-5",
            torch_dtype=torch.float16,
            revision="fp16"
        ).to("cuda")
    else:
        image_pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
    
    return translator, text_generator, image_pipe

# Load models at startup
translator, text_generator, image_pipe = load_models()

def clean_text(text):
    """Remove special characters and truncate to complete sentences"""
    cleaned = re.split(r'(?<=[.!?])\s+', text)[0]
    return re.sub(r'[^a-zA-Z0-9,.!?\'"\- ]+', '', cleaned).strip()

def process_content(tamil_input, creativity_level):
    outputs = {}
    error = ""
    
    try:
        # Translate Tamil to English
        translation_result = translator(tamil_input)
        english_text = translation_result[0]['translation_text']
        outputs["translation"] = english_text
        
        # Generate image
        image = image_pipe(
            english_text, 
            guidance_scale=creativity_level,
            num_inference_steps=30
        ).images[0]
        outputs["image"] = image
        
        # Generate creative text
        creative_output = text_generator(
            f"Create creative content about: {english_text}",
            max_length=150,
            temperature=creativity_level/10,
            num_return_sequences=1
        )
        outputs["creative_text"] = clean_text(creative_output[0]['generated_text'])
        
    except Exception as e:
        error = f"⚠️ Error: {str(e)}"
    
    return outputs, error

# Gradio UI
with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown("# 🌐 தமிழ் உரை முதல் பட உருவாக்கம் (Tamil to Image Generator)")
    gr.Markdown("தமிழில் உள்ளீடு செய்து → ஆங்கில மொழிபெயர்ப்பு + AI உருவம் + படைப்பு உரை பெறவும்")
    
    with gr.Row():
        with gr.Column():
            tamil_input = gr.Textbox(
                label="தமிழ் உள்ளீடு", 
                placeholder="உங்கள் உரையை இங்கே உள்ளிடவும்...",
                lines=3
            )
            creativity = gr.Slider(
                label="படைப்பாற்றல் நிலை", 
                minimum=1, maximum=10, value=7, step=1
            )
            submit_btn = gr.Button("உருவாக்கு")
            
        with gr.Column():
            translation_box = gr.Textbox(label="ஆங்கில மொழிபெயர்ப்பு")
            creative_output = gr.Textbox(label="படைப்பு உரை", lines=3)
            image_output = gr.Image(label="உருவாக்கப்பட்ட படம்")
            error_output = gr.Textbox(label="பிழை செய்திகள்", visible=True)
    
    examples = gr.Examples(
        examples=[
            ["கடலின் அடியில் மறைந்திருக்கும் பழைய நகரம்", 8],
            ["பனி படர்ந்த குளிர்காலத்தில் வெப்பமான காபி குடிக்கும் பழங்குடி பெண்", 7],
            ["வேறு கிரகத்தில் இருந்து வந்த அறிவார்ந்த இயந்திரங்கள்", 9]
        ],
        inputs=[tamil_input, creativity]
    )
    
    # Clear inputs button
    clear_btn = gr.Button("துடைத்து துவக்கவும்")
    
    def clear_all():
        return {"": ""}, "", "", None, ""
    
    submit_btn.click(
        fn=process_content,
        inputs=[tamil_input, creativity],
        outputs=[{"translation": translation_box, "creative_text": creative_output, "image": image_output}, error_output]
    )
    
    clear_btn.click(
        fn=clear_all,
        inputs=[],
        outputs=[tamil_input, translation_box, creative_output, image_output, error_output]
    )

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