#------------------------------------------------------------------------
# Import
#------------------------------------------------------------------------

import streamlit as st
import requests
from PIL import Image
import io
import os

#------------------------------------------------------------------------
# HF API
#------------------------------------------------------------------------

# Retrieve the HF API key from environment variables 
hf_api_key = os.getenv('HF_API_KEY')
if not hf_api_key:
    raise ValueError("HF_API_KEY not set in environment variables")

API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
headers = {"Authorization": f"Bearer {hf_api_key}"}

#------------------------------------------------------------------------
# Configurations
#------------------------------------------------------------------------
# Streamlit page setup
st.set_page_config(
    page_title="Stxtement | Image Generation", 
    page_icon=":art:", 
    layout="centered", 
    initial_sidebar_state="auto",
    menu_items={
        'Get Help': 'mailto:info@mtss.ai',
        'About': "This app is built to support spreadsheet analysis"
    }
)

#------------------------------------------------------------------------
# Sidebar
#------------------------------------------------------------------------
with st.sidebar:
    # Password input field
    # password = st.text_input("Enter Password:", type="password")
    
    # Set the desired width in pixels
    image_width = 300  
    # Define the path to the image
    # image_path = "mimtss.png"
    # # Display the image
    # st.image(image_path, width=image_width)
    
    # Set the title
    st.title("MTSS.ai")
    
    # Toggle for Help and Report a Bug
    with st.expander("Need help and report a bug"):
        st.write("""
        **Contact**: Cheyne LeVesseur, PhD  
        **Email**: info@mtss.ai
        """)
    st.divider()
    st.subheader('User Instructions')
    
    # Principles text with Markdown formatting
    User_Instructions = """
    Enter a detailed description of the image you want to generate, and the app will create it based on your prompt.
    """
    st.markdown(User_Instructions)

#------------------------------------------------------------------------
# Define functions
#------------------------------------------------------------------------

# SIMPLE CODE
# def query(payload):
#     response = requests.post(API_URL, headers=headers, json=payload)
#     if response.status_code != 200:
#         st.error(f"Error: {response.status_code} - {response.text}")
#         return None
#     return response.content

# def generate_image(prompt):
#     image_bytes = query({"inputs": prompt})
#     if image_bytes:
#         return Image.open(io.BytesIO(image_bytes))
#     return None

# def main():
#     st.title("Stxtement | Image Generation")

#     prompt = st.text_input("Enter a prompt for image generation:")
    
#     if st.button("Generate Image"):
#         if prompt:
#             image = generate_image(prompt)
#             if image:
#                 st.image(image, caption="Generated Image")
#         else:
#             st.warning("Please enter a prompt.")

# COMPREHENSIVE CODE
def query(payload):
    response = requests.post(API_URL, headers=headers, json=payload)
    if response.status_code != 200:
        st.error(f"Error: {response.status_code} - {response.text}")
        return None
    return response.content

def generate_image(prompt):
    image_bytes = query({"inputs": prompt})
    if image_bytes:
        return Image.open(io.BytesIO(image_bytes))
    return None

def generate_image_callback():
    prompt = st.session_state.get("prompt_input", "")
    if prompt:
        st.session_state["prompt"] = prompt
        # The spinner is handled outside this function
        image = generate_image(prompt)
        if image:
            st.session_state["image"] = image
    else:
        st.warning("Please enter a prompt.")

def reset_callback():
    # Clear session state variables
    st.session_state["prompt"] = ""
    st.session_state["prompt_input"] = ""
    st.session_state["image"] = None

def main():
    st.title("Stxtement | Image Generation")

    # Input field for the prompt, tied to session state
    st.text_input(
        "Enter a prompt for image generation:",
        value=st.session_state.get("prompt_input", ""),
        key="prompt_input"
    )

    # Generate Image button
    generate_button_clicked = st.button("Generate Image")

    # Create a placeholder for the spinner after the button
    spinner_placeholder = st.empty()

    if generate_button_clicked:
        if st.session_state.get("prompt_input", ""):
            with spinner_placeholder:
                with st.spinner('Generating image...'):
                    generate_image_callback()
        else:
            st.warning("Please enter a prompt.")

    # Show the image if it exists in the session state
    if st.session_state.get("image"):
        st.image(st.session_state["image"], caption="Generated Image")

        # Download button
        image_bytes = io.BytesIO()
        st.session_state["image"].save(image_bytes, format='PNG')

        st.download_button(
            label="Download Image",
            data=image_bytes.getvalue(),
            file_name="generated_image.png",
            mime="image/png"
        )

        # Reset button
        st.button("Reset", on_click=reset_callback)

    # If no image but prompt exists, show reset button to clear prompt
    elif st.session_state.get("prompt_input"):
        st.button("Reset", on_click=reset_callback)

#------------------------------------------------------------------------
# Main Guard
#------------------------------------------------------------------------

if __name__ == "__main__":
    main()