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

# Page configuration
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")

# List of API key names
api_key_names = ["OPENAI_API_KEY", "MIDJOURNEY_GOAPI_TOKEN", "REPLICATE_API_TOKEN", "STABILITY_API_KEY"]

# Initialize session state if it does not exist
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'] = ""
if 'upload_completed' not in st.session_state:
    st.session_state['upload_completed'] = False

# Define a function to request and update 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()

# Check if all API keys have been entered
if all_keys_entered:
    # Section to upload the image
    uploaded_file = st.file_uploader("Upload Image to analyze", type=['jpg', 'jpeg', 'png'], on_change=lambda: setattr(st.session_state, 'upload_completed', True))
    if uploaded_file is not None:
        st.session_state['upload_completed'] = True
        left_co, cent_co, _ = st.columns([1, 2, 1])
        with cent_co:
            st.image(uploaded_file, caption="Uploaded Image")
        _, right_co = st.columns([5, 1])
        if right_co.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.session_state['editable_prompt'] = prompt
                st.session_state['upload_completed'] = False

    if st.session_state['upload_completed']:
        # Maintain the content of the editable prompt after generating images
        editable_prompt = st.text_area("Edit the prompt as needed:", value=st.session_state['editable_prompt'], placeholder="Enter your prompt here...", height=150, key='editable_prompt', on_change=lambda: st.session_state.update(editable_prompt=editable_prompt))
        if st.button("Generate New Image", key='generate_image_btn'):
            col1, col2, col3, col4 = st.columns(4)
            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.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.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.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.image(result_path_4, caption="Midjourney Image")
            # Update the prompt in session state to keep the text
            st.session_state['editable_prompt'] = editable_prompt
            st.session_state['upload_completed'] = False  # Disable the text area and button after generating images
else:
    st.warning('Please enter all required API keys to proceed.', icon="⚠️")