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

# Load Hugging Face API key from environment
HF_API_KEY = os.getenv("HF_API_KEY")  # Add this in 'Variables and secrets' on HF Spaces
if not HF_API_KEY:
    raise ValueError("HF_API_KEY is not set. Please add it in Hugging Face 'Variables and secrets'.")

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

# Check if GPU is available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

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

# Load text generation model
generator_model = "EleutherAI/gpt-neo-1.3B"
generator = AutoModelForCausalLM.from_pretrained(generator_model).to(device)
generator_tokenizer = AutoTokenizer.from_pretrained(generator_model)
if generator_tokenizer.pad_token is None:
    generator_tokenizer.pad_token = generator_tokenizer.eos_token
generator.config.pad_token_id = generator_tokenizer.pad_token_id

def translate_tamil_to_english(text):
    """Translate Tamil text to English."""
    try:
        inputs = translator_tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128).to(device)
        output = translator.generate(**inputs)
        return translator_tokenizer.decode(output[0], skip_special_tokens=True)
    except Exception as e:
        return f"Translation error: {str(e)}"

def generate_text(prompt):
    """Generate creative text from English input."""
    try:
        inputs = generator_tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=128).to(device)
        output = generator.generate(**inputs, max_length=100)
        return generator_tokenizer.decode(output[0], skip_special_tokens=True)
    except Exception as e:
        return f"Text generation error: {str(e)}"

def generate_image(prompt):
    """Generate image using Hugging Face inference API."""
    try:
        response = requests.post(IMAGE_GEN_URL, headers=HEADERS, json={"inputs": prompt})
        if response.status_code == 200:
            content_type = response.headers.get("content-type", "")
            if "image" in content_type:
                return Image.open(io.BytesIO(response.content))
            else:
                result = response.json()
                if "image" in result:
                    image_data = base64.b64decode(result["image"])
                    return Image.open(io.BytesIO(image_data))
    except Exception as e:
        print("Error generating image:", e)
    return Image.new("RGB", (300, 300), "red")  # fallback placeholder

def process_input(tamil_text):
    """Pipeline: Translate → Generate Text → Generate Image"""
    english_text = translate_tamil_to_english(tamil_text)
    creative_text = generate_text(english_text)
    image = generate_image(english_text)
    return english_text, creative_text, image

# Create Gradio UI
interface = gr.Interface(
    fn=process_input,
    inputs=gr.Textbox(label="Enter Tamil Text"),
    outputs=[
        gr.Textbox(label="Translated English Text"),
        gr.Textbox(label="Creative Text"),
        gr.Image(label="Generated Image")
    ],
    title="Tamil to English Translator & Image Generator",
    description="Enter Tamil text. This app translates it to English, generates a creative description, and produces an image based on the translated text."
)

# Launch app
interface.launch()