import gradio as gr
import requests
import os
import io
from PIL import Image
from dotenv import load_dotenv

load_dotenv()

#Générateur d'images
FLUX1_APIKEY = os.getenv('FLUX1_APIKEY')
#LLM
GEMINI_APIKEY = os.getenv('GEMINI_APIKEY')

# TODO : 
# - Prendre en entrée soit un texte, soit un fichier audio (empecher saisie de l'un si l'autre est rempli)
# - Le LLM résume cette entrée en plusieurs axes avec des mots clés par axes
# - Flux génère une diapo par axé généré par le LLM précédemment 
# - On se retrouve en sortie avec plusieurs diapo et des mots clés les accompagnants

def Main(prompt, inputText):
    return GenerateImageFromText(prompt), GenerateTextLLM(inputText)

def GenerateImageFromText(prompt):
    API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
    headers = {"Authorization": f"Bearer {FLUX1_APIKEY}"}

    def query(payload):
        response = requests.post(API_URL, headers=headers, json=payload)
        return response.content

    image_bytes = query({
        "inputs": prompt,
    })
    image = Image.open(io.BytesIO(image_bytes))

    return image

def GenerateTextLLM(inputText):
    
    url = f'https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash-latest:generateContent?key={GEMINI_APIKEY}'
    headers = {
        'Content-Type': 'application/json'
    }
    data = {
        "contents": [
            {
                "parts": [
                    {
                        "text": inputText
                    }
                ]
            }
        ]
    }

    response = requests.post(url, headers=headers, json=data)

    try:
        return response.json()['candidates'][0]['content']['parts'][0]['text']
    except:
        return 'Error'


outputs_image = [
    gr.components.Image(type="pil", label="Output Image"),
    "text"
]

interface_image = gr.Interface(
    fn=Main,
    inputs=["text", "text"],
    outputs=outputs_image,
    title="Feur",
) 

interface_image.launch()