import streamlit as st
from gradio_client import Client, handle_file
import threading
import time
import os
import uuid
import shutil

# ---------------------------
# CONFIG
# ---------------------------
IMAGE_DIR = "data/images"
os.makedirs(IMAGE_DIR, exist_ok=True)

client = Client("Galaxydude2/a")
client_process = Client("SLSLK/n8")

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

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

if "page" not in st.session_state:
    st.session_state.page = "🖼️ Generieren"

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

# ---------------------------
# UI / THEME
# ---------------------------
st.set_page_config(layout="wide")

st.markdown("""
<style>
.stApp { background-color: #0f0f1a; color: #e0e0ff; }
h1 { color: #ff2e82; text-align:center; }
.stButton > button {
    background-color: #ff2e82;
    color: white;
    border-radius: 10px;
}
.stButton > button:hover { background-color: #ff5ca8; }
</style>
""", unsafe_allow_html=True)

st.markdown("<h1>🎨 Omni Image Studio</h1>", unsafe_allow_html=True)

# ---------------------------
# HELPERS
# ---------------------------
def save_local_copy(src_path):
    if os.path.exists(src_path):
        new_path = os.path.join(IMAGE_DIR, f"{uuid.uuid4()}.png")
        shutil.copy(src_path, new_path)
        return new_path
    return None


def prepare_image(input_image):
    try:
        if hasattr(input_image, "getvalue"):
            path = os.path.join(IMAGE_DIR, f"{uuid.uuid4()}.png")
            with open(path, "wb") as f:
                f.write(input_image.getvalue())
            return path

        elif isinstance(input_image, str) and os.path.exists(input_image):
            return input_image

        return None

    except Exception as e:
        st.error(f"❌ Prepare Error: {e}")
        return None


# ---------------------------
# 🔥 ROBUST API CALL
# ---------------------------
def robust_predict(func, *args, retries=3, delay=2, **kwargs):
    for attempt in range(retries):
        try:
            return func(*args, **kwargs)

        except Exception as e:
            if attempt < retries - 1:
                st.warning(f"⚠️ Retry {attempt+1}/{retries}...")
                time.sleep(delay)
            else:
                st.error("❌ API dauerhaft fehlgeschlagen")

                with st.expander("🔍 Debug Details"):
                    st.code(str(e))

                return None

# ---------------------------
# API FUNCTIONS
# ---------------------------
def generate_image(prompt):
    return robust_predict(
        client.predict,
        prompt=prompt,
        api_name="/generate_image"
    )


def edit_image(input_image, prompt):
    if not input_image or not prompt:
        st.warning("❗ Bild & Prompt erforderlich")
        return None

    image_path = prepare_image(input_image)
    if not image_path:
        return None

    return robust_predict(
        client.predict,
        input_image=handle_file(image_path),
        edit_prompt=prompt,
        api_name="/edit_image"
    )


def process_image(input_image):
    if not input_image:
        st.warning("❗ Bild erforderlich")
        return None

    image_path = prepare_image(input_image)
    if not image_path:
        return None

    return robust_predict(
        client_process.predict,
        image=handle_file(image_path),
        api_name="/process_image"
    )

# ---------------------------
# ASYNC + PROGRESS
# ---------------------------
def run_async(func, *args):
    result = {"done": False, "data": None}

    def task():
        result["data"] = func(*args)
        result["done"] = True

    threading.Thread(target=task).start()
    return result


def show_progress(job):
    bar = st.progress(0)
    status = st.empty()
    percent = 0

    msgs = [
        "🔄 Processing...",
        "🎨 Generating...",
        "🧠 AI thinking...",
        "✨ Almost done..."
    ]

    i = 0
    while not job["done"]:
        percent = min(percent + 5, 95)
        bar.progress(percent)
        status.text(msgs[i % len(msgs)])
        i += 1
        time.sleep(0.3)

    bar.progress(100)
    status.text("✅ Done!")

# ---------------------------
# NAVIGATION
# ---------------------------
page = st.radio(
    "Navigation",
    ["🖼️ Generieren", "✨ Bearbeiten", "🧪 Analyse", "📚 Galerie"],
    index=["🖼️ Generieren", "✨ Bearbeiten", "🧪 Analyse", "📚 Galerie"].index(st.session_state.page),
    horizontal=True
)

st.session_state.page = page

# ---------------------------
# GENERATE
# ---------------------------
if page == "🖼️ Generieren":
    prompt = st.text_area("Prompt")

    if st.session_state.prompt_history:
        st.subheader("🧠 History")
        for i, p in enumerate(reversed(st.session_state.prompt_history[-5:])):
            if st.button(p, key=f"h{i}"):
                prompt = p

    if st.button("Generieren"):
        if prompt:
            st.session_state.prompt_history.append(prompt)

        job = run_async(generate_image, prompt)
        show_progress(job)

        result = job["data"]

        if result:
            st.image(result)

            with open(result, "rb") as f:
                st.download_button("📥 Download", f, file_name="generated.png")

            path = save_local_copy(result)
            if path:
                st.session_state.gallery.append({"path": path, "prompt": prompt})

# ---------------------------
# EDIT
# ---------------------------
elif page == "✨ Bearbeiten":
    if st.session_state.selected_image:
        st.image(st.session_state.selected_image)
        input_image = st.session_state.selected_image
    else:
        input_image = st.file_uploader("Upload Image")

    prompt = st.text_area("Edit Prompt")

    if st.button("Bearbeiten"):
        job = run_async(edit_image, input_image, prompt)
        show_progress(job)

        result = job["data"]

        if result:
            st.image(result)

            with open(result, "rb") as f:
                st.download_button("📥 Download", f, file_name="edited.png")

            path = save_local_copy(result)
            if path:
                st.session_state.gallery.append({"path": path, "prompt": prompt})

# ---------------------------
# ANALYSE
# ---------------------------
elif page == "🧪 Analyse":
    if st.session_state.selected_image:
        st.image(st.session_state.selected_image)
        input_image = st.session_state.selected_image
    else:
        input_image = st.file_uploader("Upload Image")

    if st.button("Analysieren"):
        job = run_async(process_image, input_image)
        show_progress(job)

        result = job["data"]

        if result:
            st.image(result)

            with open(result, "rb") as f:
                st.download_button("📥 Download", f, file_name="analysis.png")

            if os.path.exists(result):
                path = save_local_copy(result)
                st.session_state.gallery.append({"path": path, "prompt": "Analyse"})

# ---------------------------
# GALLERY
# ---------------------------
elif page == "📚 Galerie":
    if not st.session_state.gallery:
        st.info("Keine Bilder")
    else:
        cols = st.columns(3)

        for i, item in enumerate(st.session_state.gallery):
            with cols[i % 3]:
                st.image(item["path"], caption=item["prompt"])

                with open(item["path"], "rb") as f:
                    st.download_button("📥", f, file_name=f"img_{i}.png", key=f"d{i}")

                c1, c2 = st.columns(2)

                if c1.button("✏️", key=f"e{i}"):
                    st.session_state.selected_image = item["path"]
                    st.session_state.page = "✨ Bearbeiten"
                    st.rerun()

                if c2.button("🧪", key=f"a{i}"):
                    st.session_state.selected_image = item["path"]
                    st.session_state.page = "🧪 Analyse"
                    st.rerun()

        if st.button("🗑️ Clear"):
            st.session_state.gallery = []