import streamlit as st
import requests
import base64
import os

st.title("OpenAI GPT-Image-1 Text-to-Image Demo")

prompt = st.text_area("Please enter image description (prompt)")
api_key = os.environ.get("OPENAI_API_KEY")
if api_key:
    st.info("OPENAI_API_KEY environment variable detected, will use automatically.")
else:
    api_key = st.text_input("Please enter OpenAI API Key", type="password")

uploaded_image = st.file_uploader(
    "(Optional) Upload an image to edit", type=["png", "jpg", "jpeg", "webp"]
)

if st.button("Generate Image") and prompt and api_key:
    if uploaded_image:
        # Call image edit API
        headers = {"Authorization": f"Bearer {api_key}"}
        files = {
            "image": (uploaded_image.name, uploaded_image, uploaded_image.type),
        }
        data = {
            "prompt": prompt,
            "model": "gpt-image-1",
            "size": "1024x1024",
        }
        with st.spinner("Image editing in progress..."):
            response = requests.post(
                "https://api.openai.com/v1/images/edits",
                headers=headers,
                files=files,
                data=data,
            )
        if response.status_code == 200:
            img_b64 = response.json()["data"][0]["b64_json"]
            img_bytes = base64.b64decode(img_b64)
            col1, col2 = st.columns(2)
            with col1:
                st.image(uploaded_image, caption="Original Image", use_column_width=True)
            with col2:
                st.image(img_bytes, caption="Edited Image", use_column_width=True)
        else:
            st.error(f"API Error: {response.text}")
    else:
        # Call image generation API
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        }
        payload = {
            "prompt": prompt,
            "model": "gpt-image-1",
            "output_format": "png",
            "size": "1024x1024",
        }
        with st.spinner("Image generation in progress..."):
            response = requests.post(
                "https://api.openai.com/v1/images/generations",
                headers=headers,
                json=payload,
            )
        if response.status_code == 200:
            img_b64 = response.json()["data"][0]["b64_json"]
            img_bytes = base64.b64decode(img_b64)
            st.image(img_bytes, caption="Generated Image", use_column_width=True)
        else:
            st.error(f"API Error: {response.text}")