import streamlit as st
import os
from dotenv import load_dotenv
from agent.graph import app
from agent.state import AgentState

load_dotenv()

st.set_page_config(page_title="AutoStream AI Sales Assistant", page_icon="🎬", layout="centered")

# Custom CSS for minimalist, cooler UI
st.markdown("""
<style>
    .stChatFloatingInputContainer {
        bottom: 20px;
    }
    .main {
        background-color: #0E1117;
    }
    h1 {
        color: #E2E8F0;
        font-family: 'Helvetica Neue', Helvetica, Arial, sans-serif;
        font-weight: 700;
        text-align: center;
        margin-bottom: 2rem;
    }
    .subtitle {
        color: #94A3B8;
        text-align: center;
        margin-bottom: 2rem;
        font-size: 1.1rem;
    }
    .stAlert {
        border-radius: 8px;
    }
</style>
""", unsafe_allow_html=True)

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

    st.session_state.agent_state = AgentState(
        conversation_history=[],
        current_message="",
        detected_intent=None,
        retrieved_documents=[],
        user_name=None,
        user_email=None,
        creator_platform=None,
        lead_ready=False,
        response=""
    )

    st.session_state.messages.append({"role": "assistant", "content": "Hello! I'm the AutoStream assistant. I can answer questions about our features and pricing. How can I help you today?"})

st.markdown("<h1>🎬 AutoStream Assistant</h1>", unsafe_allow_html=True)
st.markdown("<div class='subtitle'>Ask about features and pricing, or sign up for a plan instantly!</div>", unsafe_allow_html=True)

if not os.environ.get("OPENAI_API_KEY"):
    st.info("ℹ️ OPENAI_API_KEY is not set. The system will fall back to a local Qwen model and HuggingFace embeddings.")

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

if prompt := st.chat_input("What would you like to know?"):

    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    st.session_state.agent_state["current_message"] = prompt

    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            try:
                result_state = app.invoke(st.session_state.agent_state)
                st.session_state.agent_state = result_state

                response = result_state["response"]

                st.session_state.agent_state["conversation_history"].append({"role": "user", "content": prompt})
                st.session_state.agent_state["conversation_history"].append({"role": "assistant", "content": response})

                if len(st.session_state.agent_state["conversation_history"]) > 12:
                    st.session_state.agent_state["conversation_history"] = st.session_state.agent_state["conversation_history"][-12:]

                st.markdown(response)

                with st.expander("Agent Reasoning & State", expanded=False):
                    st.write(f"**Detected Intent:** `{result_state.get('detected_intent', 'UNKNOWN')}`")
                    if result_state.get("retrieved_documents") and result_state.get("detected_intent") in ["PRODUCT_QUERY", "PRICING_QUERY"]:
                        st.write(f"**RAG Retrieval:** Found {len(result_state['retrieved_documents'])} relevant knowledge chunks.")

                    st.write("**Lead Data:**")
                    st.json({
                        "user_name": result_state.get("user_name"),
                        "user_email": result_state.get("user_email"),
                        "creator_platform": result_state.get("creator_platform"),
                        "lead_ready": result_state.get("lead_ready")
                    })

            except Exception as e:
                response = f"An error occurred: {str(e)}"
                st.error(response)

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