import streamlit as st
from graph import run_graph
from PIL import Image
import io
import tempfile
import os

st.set_page_config(
    page_title="AI Image Studio",
    layout="wide"
)

st.title("🧠 AI Image Studio (Agents + LangGraph)")


# ---------------- SESSION STATE ----------------
if "messages" not in st.session_state:
    st.session_state.messages = []

if "uploaded_paths" not in st.session_state:
    st.session_state.uploaded_paths = []


# ---------------- SIDEBAR ----------------
st.sidebar.header("📁 Upload Files")

uploaded_files = st.sidebar.file_uploader(
    "Upload Images",
    type=["png", "jpg", "jpeg"],
    accept_multiple_files=True
)

# Convert Streamlit file → REAL FILE PATHS
if uploaded_files:

    paths = []

    for file in uploaded_files:

        temp_dir = "uploads"
        os.makedirs(temp_dir, exist_ok=True)

        file_path = os.path.join(temp_dir, file.name)

        with open(file_path, "wb") as f:
            f.write(file.getbuffer())

        paths.append(file_path)

    st.session_state.uploaded_paths = paths

    st.sidebar.success(f"{len(paths)} files saved")


# Preview
if st.session_state.uploaded_paths:

    st.sidebar.subheader("Preview")

    for path in st.session_state.uploaded_paths:

        st.sidebar.image(path, caption=os.path.basename(path))


# ---------------- CHAT UI ----------------
st.subheader("💬 Chat with AI Agent")

for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.write(msg["content"])


user_input = st.chat_input("Generate / Edit / Analyze image...")


# ---------------- RUN AGENT ----------------
if user_input:

    st.session_state.messages.append({
        "role": "user",
        "content": user_input
    })

    with st.chat_message("user"):
        st.write(user_input)

    with st.chat_message("assistant"):
        with st.spinner("AI Agent thinking..."):

            result = run_graph(
                user_input,
                uploaded_files=st.session_state.uploaded_paths
            )

            output = result.get("result")


            # ---------------- IMAGE OUTPUT ----------------
            if isinstance(output, Image.Image):

                st.image(output, caption="Result", use_container_width=True)

                buf = io.BytesIO()
                output.save(buf, format="PNG")

                st.download_button(
                    "⬇ Download",
                    buf.getvalue(),
                    "result.png",
                    "image/png"
                )


            # ---------------- DICT OUTPUT ----------------
            elif isinstance(output, dict):

                st.json(output)


            # ---------------- TEXT OUTPUT ----------------
            else:

                st.write(output)


    st.session_state.messages.append({
        "role": "assistant",
        "content": "Done"
    })