import streamlit as st
from gpt_vision_prompt import generate_prompt_with_vision
import tempfile
from dalle_generate_img import generate_img_with_dalle
from stability_generate_img import generate_image_with_stability
from finetune_generate_img import generate_finetuned_img
from midjourney_generate_img import midjourney_generate_img


# Configuración de la página
st.set_page_config(layout="wide")
st.sidebar.title("API Keys")
st.markdown("<h1 style='text-align: center; color: grey;'>Image Generation App</h1>", unsafe_allow_html=True)
st.text("Prepared by joel.orellana@above0.com for fomo.ai")

# Lista de nombres de las API keys
api_key_names = ["OPENAI_API_KEY", "MIDJOURNEY_GOAPI_TOKEN", "REPLICATE_API_TOKEN", "STABILITY_API_KEY"]

# Inicializar el estado de la sesión si no existe
if 'api_keys' not in st.session_state:
    st.session_state['api_keys'] = {key_name: "" for key_name in api_key_names}
if 'editable_prompt' not in st.session_state:
    st.session_state['editable_prompt'] = ""

# Definir una función para solicitar y actualizar las API keys
def request_and_update_api_keys():
    all_keys_entered = True
    for key_name in api_key_names:
        key_value = st.sidebar.text_input(f"Enter {key_name}:", value=st.session_state['api_keys'].get(key_name, ""), type="password", key=key_name)
        st.session_state['api_keys'][key_name] = key_value
        if not key_value:
            all_keys_entered = False
    return all_keys_entered

all_keys_entered = request_and_update_api_keys()

# Revisar si todas las API keys han sido ingresadas
if all_keys_entered:
    # Sección para subir la imagen
    uploaded_file = st.file_uploader("Upload Image to analyze", type=['jpg', 'jpeg', 'png'])
    if uploaded_file is not None:
        left_co, cent_co,last_co = st.columns(3)
        with cent_co:
            st.image(uploaded_file, caption="Uploaded Image")
        # Botón para generar el prompt solo si hay una imagen subida
        if st.button("Generate Prompt"):
            with st.spinner("Generating Prompt..."):
                with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as temp_file:
                    temp_file.write(uploaded_file.getvalue())
                    temp_path = temp_file.name
                api_key = st.session_state['api_keys']['OPENAI_API_KEY']
                prompt = generate_prompt_with_vision(temp_path, api_key=api_key)
                st.success("Done!")
                st.session_state['editable_prompt'] = prompt  # Actualizar el prompt en el estado de la sesión

    editable_prompt = st.text_area("Edit the prompt as needed:", placeholder="Enter your prompt here...", height=150, key='editable_prompt', label_visibility='hidden')

    col1, col2, col3, col4 = st.columns(4)
    
    if st.session_state['editable_prompt'] and st.button("Generate New Image"):
        with col1:
            with st.spinner("Generating DALL·E Image..."):
                result_path_1 = generate_img_with_dalle(st.session_state['editable_prompt'], api_key=st.session_state['api_keys']['OPENAI_API_KEY'])
                st.success("Generated DALL·E Image!")
                st.image(result_path_1, caption="DALL·E Image")
                
        with col2:
            with st.spinner("Generating Stable Diffusion Image..."):
                result_path_2 = generate_image_with_stability(st.session_state['editable_prompt'], api_key=st.session_state['api_keys']['STABILITY_API_KEY'])
                st.success("Generated Stable Diffusion Image!")
                st.image(result_path_2, caption="Stable Diffusion Image")
        
        with col3:
            with st.spinner("Generating Finetuning Image..."):
                result_path_3 = generate_finetuned_img(st.session_state['editable_prompt'], api_key=st.session_state['api_keys']['REPLICATE_API_TOKEN'])
                st.success("Generated Image using a finetuned model!")
                st.image(result_path_3, caption="Finetuned SDXL Image")
        
        with col4:
            with st.spinner("Generating Midjourney Image..."):
                result_path_4 = midjourney_generate_img(st.session_state['editable_prompt'], api_key=st.session_state['api_keys']['MIDJOURNEY_GOAPI_TOKEN'])
                st.success("Generated Midjourney Image!")
                st.image(result_path_4, caption="Midjourney Image")

else:
    st.warning('Please enter all required API keys to proceed.', icon="⚠️")