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

st.set_page_config(initial_sidebar_state="collapsed", layout="centered")

def generate_image(api_key, prompt, aspect_ratio, mode, model, seed, output_format, strength=None, negative_prompt=None, image_file=None):
    headers = {
        "authorization": f"Bearer {api_key}",
        "accept": "image/*"
    }
    data = {
        "prompt": prompt,
        "mode": mode,
        "model": model,
        "output_format": output_format
    }
    
    files = {"none": ''}

    if mode == "text-to-image":
        files["aspect_ratio"] = (None, aspect_ratio)
    elif mode == "image-to-image":
        if seed != 0:
            data['seed'] = seed
    if negative_prompt:
        data['negative_prompt'] = negative_prompt
    if image_file is not None:
        image_bytes = image_file.getvalue()
        files = {'image': ('image.jpg', image_bytes, 'image/jpeg')}
    if strength is not None:
        files["strength"] = (None, str(strength))
    if seed is not None:
        files["seed"] = (None, str(seed))
    response = requests.post(
        "https://api.stability.ai/v2beta/stable-image/generate/sd3",
        headers=headers,
        files=files,
        data=data
    )
    
    return response

def main():
    st.image("static/SD3_webui_logo_image.png", use_column_width=True)
    
    if 'api_key' not in st.session_state:
        st.session_state.api_key = ""
    if 'prompt' not in st.session_state:
        st.session_state.prompt = ""
    if 'negative_prompt' not in st.session_state:
        st.session_state.negative_prompt = ""
    if 'image_data' not in st.session_state:
        st.session_state.image_data = None
    if 'image_path' not in st.session_state:
        st.session_state.image_path = None
    
    api_key = st.text_input("Enter your API key:", type="password", value=st.session_state.api_key)
    models = st.selectbox("Select model:", ["sd3", "sd3-turbo"], index=0)
    mode = st.selectbox("Select generation mode:", ["text-to-image", "image-to-image"], index=0)
    prompt = st.text_area("Enter positive prompt:", value=st.session_state.prompt)
    negative_prompt = st.text_input("Enter negative prompt: (optional)", value=st.session_state.negative_prompt)
    aspect_ratios = st.selectbox("Select the aspect ratio:", ["1:1", "2:3", "3:2", "4:5", "5:4", "16:9", "21:9", "9:16", "9:21"], index=0) if mode == "text-to-image" else None
    seed = st.number_input("Set seed: (randomize = 0)", min_value=0, max_value=4294967294, value=0, step=1)
    strength = st.slider("Denoising strength:", min_value=0.0, max_value=1.0, value=0.5, step=0.01) if mode == "image-to-image" else None
    output_formats = st.selectbox("Select the output format:", ["jpeg", "png"], index=1)
    
    st.session_state.api_key = api_key
    st.session_state.prompt = prompt
    st.session_state.negative_prompt = negative_prompt
    
    uploaded_file = st.file_uploader("Upload for image-to-image generation:", type=["png", "jpg", "jpeg"]) if mode == "image-to-image" else None

    if st.button(':blue[Generate Image]', use_container_width=True):
        with st.spinner('Generating Image...'):
            result = generate_image(api_key, prompt, aspect_ratios, mode, models, seed, output_formats, strength, negative_prompt, uploaded_file)
            if result.status_code == 200:
                st.session_state.image_data = Image.open(io.BytesIO(result.content))
                st.session_state.image_path = os.path.join(f"gen_{models}_{seed}_{time.strftime('%H.%M.%S').lower()}.{output_formats}")
                st.session_state.image_data.save(st.session_state.image_path)
                st.success(f'Image saved to {st.session_state.image_path}')
            else:
                st.error('Failed to generate image: ' + str(result.json()))

    if st.session_state.image_data is not None:
        st.image(st.session_state.image_data, use_column_width=True)
        with open(st.session_state.image_path, "rb") as file:
            st.download_button(label=":green[Download Image]", data=file, file_name=st.session_state.image_path, mime=f"image/{output_formats}", use_container_width=True)

    with st.sidebar:
        st.link_button("Stability.Ai _ API Documentation", "https://platform.stability.ai/docs/api-reference", use_container_width=True)
        with st.expander("Image Generation Costs", expanded=True):
            st.markdown("""
            - SD3 - 6.5 credits per image or $0.065
            - SD3 Turbo - 4.0 credits per image or $0.04
            
            Additional credits can be purchased via the account_page button below.
            
            Credits cost $10 per 1,000 credits, which is enough credits for roughly 154 SD3 images or 250 SD3 Turbo images.
            """)
        st.link_button("Account_Page", "https://platform.stability.ai/account/credits", use_container_width=True)

if __name__ == "__main__":
    main()