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

# Initialize models (load once at startup)
# Translation pipeline (Tamil to English)
translator = pipeline(
    "translation", 
    model="Helsinki-NLP/opus-mt-ta-en",
    device=0 if torch.cuda.is_available() else -1
)

# Text generation pipeline (English creative content)
text_generator = pipeline(
    "text-generation", 
    model="gpt2-medium",
    device=0 if torch.cuda.is_available() else -1
)

# Image generation pipeline (English text to image)
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")

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

def process_content(tamil_input, creativity_level):
    """Main processing pipeline: Translate → Generate Image → Create Text"""
    outputs = {}
    error = ""
    
    try:
        # Translate Tamil to English
        translation = translator(tamil_input)
        english_text = translation[0]['translation_text']
        outputs["translation"] = english_text
        
        # Generate image from translated text
        image = image_pipe(
            english_text, 
            guidance_scale=creativity_level
        ).images[0]
        outputs["image"] = image
        
        # Generate creative text from translation
        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 Creative Content Generator")
    gr.Markdown("Enter Tamil text → Get English translation + AI-generated image + creative English text")
    
    with gr.Row():
        with gr.Column():
            tamil_input = gr.Textbox(
                label="தமிழ் உள்ளீடு (Tamil Input)", 
                placeholder="உங்கள் உரையை இங்கே உள்ளிடவும்...",
                lines=3
            )
            creativity = gr.Slider(
                label="Creativity Level", 
                minimum=1, maximum=10, value=7, step=1
            )
            submit_btn = gr.Button("Generate Content")
            
        with gr.Column():
            translation_box = gr.Textbox(label="English Translation")
            creative_output = gr.Textbox(label="Creative English Content")
            image_output = gr.Image(label="Generated Image")
            error_output = gr.Textbox(label="System Messages", visible=True)
    
    examples = gr.Examples(
        examples=[
            ["ஒரு மலையின் மீது உள்ள அழகிய கோட்டை", 8],
            ["விண்மீன்கள் நிறைந்த இரவு வானத்தில் பறக்கும் யானை", 9],
            ["தாவரங்கள் வளரும் புதிரான கோளம்", 10]
        ],
        inputs=[tamil_input, creativity]
    )
    
    submit_btn.click(
        fn=process_content,
        inputs=[tamil_input, creativity],
        outputs=[{"translation": translation_box, "creative_text": creative_output, "image": image_output}, error_output]
    )

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