import streamlit as st
import time
import base64
import io
from PIL import Image
from together import Together
import os

# Funzione per generare le immagini, con gestione errori e retry dopo 10 secondi
def generate_image(prompt, max_retries=5):
    api_key = os.getenv('API_REPLIACATE')
    client = Together(api_key = api_key)
    retries = 0
    while retries < max_retries:
        try:
            response = client.images.generate(
                prompt=prompt,
                model="black-forest-labs/FLUX.1-schnell-Free",
                width=960,
                height=1440,
                steps=4,
                n=1,
                response_format="b64_json"
            )
            return response.data  # Una lista di oggetti con attributo b64_json
        except Exception as e:
            st.error(f"Errore durante la generazione delle immagini: {e}. Riprovo tra 10 secondi...")
            time.sleep(10)
            retries += 1
    st.error("Numero massimo di tentativi raggiunto. Impossibile generare le immagini.")
    return None

def generate_images(prompt, num_immagini):
    for numero in range(num_immagini):
        images_data = generate_image(prompt)
        if images_data is not None:
            for i, img_obj in enumerate(images_data):
                try:
                    image_bytes = base64.b64decode(img_obj.b64_json)
                    image = Image.open(io.BytesIO(image_bytes))
                    st.image(image, caption="")
                except Exception as e:
                    st.error(f"Errore nella visualizzazione dell'immagine {i+1}: {e}")
        else:
            st.error("Non è stato possibile generare le immagini. Riprova più tardi.")
        time.sleep(5)
    st.success("Immagini generate con successo!")

def main():
    st.title("AI Imaging")
    st.sidebar.header("Impostazioni")
    stile_default = "Highly detailed, painterly style with a historical yet stylized aesthetic. Rich textures, ornate patterns, and a color palette dominated by imperial gold, deep red, and aged marble tones. Inspired by ancient Roman mosaics, frescoes, and classical sculpture, with a balanced mix of realism and stylization. Elegant, decorative card borders with intricate engravings and antique flourishes. Designed for a tabletop card game, ensuring clarity, readability, and a visually immersive experience."
    stile_immagine = st.sidebar.text_area("Stile Immagine", stile_default, disabled=False)
    prompt_input = st.sidebar.text_area("Prompt Immagine")
    num_immagini = st.sidebar.slider("Variazioni", min_value=1, max_value=6, value=2)
    submit_button = st.sidebar.button(label="Genera Immagine", type = "primary", use_container_width=True)
    st.write("Inserisci il **Prompt** (che verrà unito allo stile) e clicca su *Genera Immagine*.")
    if submit_button:
        if not prompt_input.strip():
            st.error("Per favore, inserisci un prompt per l'immagine!")
        else:
            # Combiniamo il prompt inserito con lo stile fisso
            prompt_completo = f"{prompt_input}, {stile_immagine}"
            
            st.info("Generazione in corso, attendere...")
            generate_images(prompt_completo, num_immagini)

if __name__ == "__main__":
    main()