import streamlit as st
import requests
import base64
from streamlit import switch_page

API_URL = "https://subashdvorak-brainstroming-fast-api.hf.space"

st.set_page_config(page_title="AI Business Ideation Platform", layout="wide")
st.header("Welcome to RT-GENAI")

# Initialize session state
session_defaults = {
    'business_details': {},
    'details_completed': False,
    'final_ideation': [],
    'human_interactions': [],
    'brainstorm_response': {},
    'final_story': "",
    'generated_image': "",
    'selected_topics': [],
    'selected_from_brainstorm': [],
    'story': "",
    'brainstorming_topics': [],
    'topics_to_pass': [],
    'base64_images': [],
    'history': [],
    'history_index': -1,
    'trigger_brainstorm': False
}
for key, val in session_defaults.items():
    if key not in st.session_state:
        st.session_state[key] = val

# Step 1: Context Analysis or Orchestration
st.subheader("How can I help you?")
msg = st.text_input("")

# Image uploader
uploaded_files = st.file_uploader("📂 Upload reference images (optional)", type=['png', 'jpg', 'jpeg'], key="initial_image_uploader")

if uploaded_files:
    base64_images = []
    with st.spinner("Encoding uploaded images..."):
        for file in uploaded_files:
            base64_images.append(base64.b64encode(file.read()).decode('utf-8'))
        st.session_state.base64_images = base64_images
else:
    print('File not uploaded')

# Submit button
if st.button("Submit"):
    if not st.session_state.details_completed:
        with st.spinner("Analyzing your business context..."):
            resp = requests.post(f"{API_URL}/context-analysis", json={"message": msg})
        if resp.ok:
            response = resp.json()
            st.markdown(response.get('response'))
            if response.get('complete'):
                st.session_state.details_completed = True
                st.session_state.business_details = response.get("business_details", {})
                st.success("Context analysis completed. Now you can interact freely.")
        else:
            st.error("Context analysis failed.")
    else:
        with st.spinner("Processing orchestration..."):
            resp = requests.post(f"{API_URL}/orchestration", json={"message": msg, "image_base64": st.session_state.base64_images})

        if resp.ok:
            response = resp.json()
            tool = response.get("tool_response")
            message = response.get("message_response")
            image_caption = response.get("image_caption")
            st.write('Tool response:', tool)

            # ========== Tool Handling ==========
            if "ideation" in tool:
                with st.spinner("Generating ideas..."):
                    resp = requests.post(f"{API_URL}/ideation")
                if resp.ok:
                    st.session_state.final_ideation = resp.json()["response"]["improver_response"][-1]
                    ideas = eval(st.session_state.final_ideation)
                    for i, idea in enumerate(ideas, 1):
                        st.markdown(f"*Idea:{i}* - {idea}")
                    st.success("Ideas generated successfully.")
                else:
                    st.error("Failed to generate ideas.")

            elif "human-idea-refining" in tool:
                with st.spinner("Refining ideas with your input..."):
                    resp = requests.post(f"{API_URL}/human-idea-refining", json={"query": message})
                if resp.ok:
                    refined = resp.json()['response']
                    st.session_state.human_interactions.append(refined)
                    st.markdown(st.session_state.human_interactions[-1])
                    st.success("Ideas refined.")
                else:
                    st.error("Idea refinement failed.")

            elif "generate-story" in tool:
                st.info("Redirecting to Brainstorm page...")
                switch_page("pages/brainstorm.py")

            elif "generate-ultimate-story" in tool:
                with st.spinner("Generating final story..."):
                    resp = requests.post(f"{API_URL}/generate-final-story")
                if resp.ok:
                    st.session_state.final_story = resp.json()["response"]
                    st.markdown(st.session_state.final_story)
                    st.success("Final story generated.")
                else:
                    st.error("Failed to generate story.")

            elif "generate-image" in tool:
                with st.spinner("Generating image from the idea..."):
                    resp = requests.post(f"{API_URL}/generate-image")
                if resp.ok:
                    st.session_state.generated_image = resp.json()["response"]
                    st.image(f"data:image/png;base64,{st.session_state.generated_image}", use_column_width=True)
                    st.success("Image generated.")
                else:
                    st.error("Image generation failed.")

            else:
                st.markdown(f"🤖 {message}")
                if "no any information" not in image_caption.lower():
                    st.markdown(f"🤖 {image_caption}")
        else:
            st.error("Orchestration failed.")