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

from graph import run_graph

# -------------------------------------------------------
# PAGE CONFIG
# -------------------------------------------------------

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 = []

if "analysis" not in st.session_state:
    st.session_state.analysis = None

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

# -------------------------------------------------------
# SIDEBAR
# -------------------------------------------------------

st.sidebar.header("📁 Upload Images")

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

# -------------------------------------------------------
# SAVE FILES
# -------------------------------------------------------

if uploaded_files:

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

    paths = []

    for file in uploaded_files:

        file_path = os.path.join("uploads", file.name)

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

        paths.append(file_path)

    # Save paths
    st.session_state.uploaded_paths = paths

    # ---------------------------------------------------
    # AUTO ANALYZE ONLY IF NEW IMAGE
    # ---------------------------------------------------

    if paths != st.session_state.last_uploaded:

        st.session_state.last_uploaded = paths

        with st.spinner("🔍 Analyzing image..."):

            result = run_graph(
                "analyze this image",
                uploaded_files=paths
            )

        analysis = result.get("result")

        st.session_state.analysis = analysis

        if isinstance(analysis, dict):

            description = analysis.get(
                "description",
                "Image analyzed."
            )

            resolution = analysis.get("resolution", "")

            style = analysis.get("style", "")

            lighting = analysis.get("lighting", "")

            objects = analysis.get("objects", [])

            message = f"""📷 **Image Analysis**

**Description**

{description}
"""

            if resolution:
                message += f"\n**Resolution:** {resolution}"

            if style:
                message += f"\n\n**Style:** {style}"

            if lighting:
                message += f"\n\n**Lighting:** {lighting}"

            if objects:
                message += f"\n\n**Objects:** {', '.join(objects)}"

            message += """

---

### What would you like to do?

• Remove background

• Make it Anime

• Make it Pixar style

• Enhance quality

• Replace objects

• Change colors

• Add new objects

• Extract text
"""

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

        else:

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

# -------------------------------------------------------
# SIDEBAR PREVIEW
# -------------------------------------------------------

if st.session_state.uploaded_paths:

    st.sidebar.subheader("🖼 Preview")

    for img in st.session_state.uploaded_paths:

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

# -------------------------------------------------------
# CHAT
# -------------------------------------------------------

st.subheader("💬 Chat with AI Agent")

for msg in st.session_state.messages:

    with st.chat_message(msg["role"]):

        st.markdown(msg["content"])

# -------------------------------------------------------
# USER INPUT
# -------------------------------------------------------

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.markdown(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")

            assistant_message = ""

            # ------------------------------------------------
            # IMAGE OUTPUT
            # ------------------------------------------------

            if isinstance(output, Image.Image):

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

                buffer = io.BytesIO()

                output.save(buffer, format="PNG")

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

                assistant_message = "✅ Image generated successfully."

            # ------------------------------------------------
            # DICTIONARY OUTPUT
            # ------------------------------------------------

            elif isinstance(output, dict):

                st.json(output)

                assistant_message = output.get(
                    "description",
                    str(output)
                )

            # ------------------------------------------------
            # TEXT OUTPUT
            # ------------------------------------------------

            else:

                st.markdown(str(output))

                assistant_message = str(output)

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