import os
import requests
import streamlit as st
from PIL import Image
from io import BytesIO
from transformers import MBartForConditionalGeneration, MBart50Tokenizer
import time

# Fetch the API keys from Hugging Face Secrets
HUGGINGFACE_TOKEN = os.getenv("Hugging_face_token")
GROQ_API_KEY = os.getenv("Groq_api")

# Hugging Face API endpoint
HF_API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
hf_headers = {"Authorization": f"Bearer {HUGGINGFACE_TOKEN}"}

# Groq API endpoint
groq_url = "https://api.groq.com/openai/v1/chat/completions"
groq_headers = {
    "Authorization": f"Bearer {GROQ_API_KEY}",
    "Content-Type": "application/json"
}

# Function to query Hugging Face model for image generation
def query_huggingface(payload):
    try:
        response = requests.post(HF_API_URL, headers=hf_headers, json=payload)
        if response.status_code != 200:
            return None
        return response.content
    except Exception:
        return None

# Function to generate text using Groq API
def generate_response(prompt):
    try:
        payload = {
            "model": "mixtral-8x7b-32768",
            "messages": [
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": prompt}
            ],
            "max_tokens": 100,
            "temperature": 0.7
        }
        response = requests.post(groq_url, json=payload, headers=groq_headers)
        if response.status_code == 200:
            result = response.json()
            return result['choices'][0]['message']['content']
        else:
            return None
    except Exception:
        return None

# Function to translate Tamil to English using MBart model
def translate_tamil_to_english(tamil_text):
    try:
        model_name = "facebook/mbart-large-50-many-to-one-mmt"
        model = MBartForConditionalGeneration.from_pretrained(model_name)
        tokenizer = MBart50Tokenizer.from_pretrained(model_name, src_lang="ta_IN")

        inputs = tokenizer(tamil_text, return_tensors="pt")
        translated = model.generate(**inputs, forced_bos_token_id=tokenizer.lang_code_to_id["en_XX"])
        translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
        return translated_text
    except Exception:
        return None

# Main function to generate text and image
def generate_image_and_text(user_input):
    with st.spinner("Generating results..."):
        time.sleep(2)  # Simulate some processing time

        # Translate Tamil to English
        english_input = translate_tamil_to_english(user_input)
        if not english_input:
            st.warning("Sorry, the translation model is unavailable right now 😥😥😥. Please try again later.")
        else:
            st.markdown("### Translated English Text:")
            st.write(english_input)

        # Generate text description (100 tokens) using Groq API
        if english_input:
            full_text_description = generate_response(english_input)
            if not full_text_description:
                st.warning("Sorry, the text generation model is unavailable right now 😥😥😥. Please try again later.")
            else:
                st.markdown("### Generated Text Response:")
                st.write(full_text_description)

                # Create image prompt based on the full text description
                image_prompt = generate_response(f"Create a concise image prompt from the following text: {full_text_description}")
                if not image_prompt:
                    st.warning("Sorry, the image prompt model is unavailable right now 😥😥😥. Please try again later.")
                else:
                    # Request an image based on the generated image prompt
                    image_data = query_huggingface({"inputs": image_prompt})
                    if not image_data:
                        st.warning("Sorry, the image generation model is unavailable right now 😥😥😥. Please try again later.")
                    else:
                        try:
                            # Load and display the image
                            image = Image.open(BytesIO(image_data))
                            st.image(image, caption="Generated Image", use_column_width=True)
                        except Exception as e:
                            st.error(f"Failed to display image: {e}")

# Streamlit interface
st.title("Multi-Modal Generator (Tamil to English)")
st.write("Enter a prompt in Tamil to generate both text and an image.")

# Input field for Tamil text
user_input = st.text_input("Enter Tamil text here:")

# Generate results when button is clicked
if st.button("Generate"):
    if user_input:
        generate_image_and_text(user_input)
    else:
        st.error("Please enter a Tamil text.")