import gradio as gr
import requests
from transformers import MarianMTModel, MarianTokenizer, AutoModelForCausalLM, AutoTokenizer
from PIL import Image
import torch
import io
import os
from typing import Tuple

# Load Hugging Face token
HF_API_KEY = os.getenv("HF_API_KEY") or "your_hf_token_here"
if not HF_API_KEY:
    raise ValueError("HF_API_KEY is not set.")

# Hugging Face inference API endpoint
IMAGE_GEN_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
HEADERS = {"Authorization": f"Bearer {HF_API_KEY}"}

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Translation model (Tamil to English)
translator_model = "Helsinki-NLP/opus-mt-mul-en"
translator_tokenizer = MarianTokenizer.from_pretrained(translator_model)
translator = MarianMTModel.from_pretrained(translator_model).to(device)

# Text generation model
text_model = "EleutherAI/gpt-neo-1.3B"
text_tokenizer = AutoTokenizer.from_pretrained(text_model)
text_generator = AutoModelForCausalLM.from_pretrained(text_model).to(device)
text_tokenizer.pad_token = text_tokenizer.eos_token

# Step 1: Tamil to English translation
def translate_tamil_to_english(text: str) -> str:
    inputs = translator_tokenizer(text, return_tensors="pt", padding=True, truncation=True).to(device)
    outputs = translator.generate(**inputs)
    return translator_tokenizer.decode(outputs[0], skip_special_tokens=True)

# Step 2: Generate creative text
def generate_text(prompt: str) -> str:
    inputs = text_tokenizer(prompt, return_tensors="pt", padding=True, truncation=True).to(device)
    outputs = text_generator.generate(**inputs, max_length=100, num_return_sequences=1)
    return text_tokenizer.decode(outputs[0], skip_special_tokens=True)

# Step 3: Generate image
def generate_image(prompt: str) -> Image.Image:
    response = requests.post(IMAGE_GEN_URL, headers=HEADERS, json={"inputs": prompt})
    if response.status_code == 200 and response.headers.get("content-type", "").startswith("image"):
        return Image.open(io.BytesIO(response.content))
    else:
        return Image.new("RGB", (512, 512), color="gray")

# Master function
def process_input(tamil_text: str) -> Tuple[str, str, Image.Image]:
    english = translate_tamil_to_english(tamil_text)
    creative = generate_text(english)
    image = generate_image(english)
    return english, creative, image

# Gradio UI using Blocks API
with gr.Blocks() as demo:
    gr.Markdown("## 🌍 Tamil to English | Text & Image Generator")

    with gr.Row():
        tamil_input = gr.Textbox(label="📝 Enter Tamil Text", placeholder="உங்கள் உரையை இங்கே உள்ளிடவும்...", lines=2)
        generate_btn = gr.Button("Translate & Generate")

    english_output = gr.Textbox(label="🇬🇧 Translated English")
    creative_output = gr.Textbox(label="✨ Generated Text")
    image_output = gr.Image(label="🖼️ Generated Image", type="pil")

    generate_btn.click(fn=process_input, inputs=tamil_input, outputs=[english_output, creative_output, image_output])

demo.launch()