import streamlit as st
import asyncio
import time
from google.adk.sessions import InMemorySessionService
from google.adk.runners import Runner
from google.genai import types
from agent import root_agent  # ← your agent file

# ────────────────────────────────────────────────
# CONFIG
# ────────────────────────────────────────────────
APP_NAME = "leave_policy_app"
session_service = InMemorySessionService()
runner = Runner(
    agent=root_agent,
    app_name=APP_NAME,
    session_service=session_service,
)

# ────────────────────────────────────────────────
# PAGE CONFIG + BASIC STYLING
# ────────────────────────────────────────────────
st.set_page_config(
    page_title="Leave Policy Assistant",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded",
)

# Minimal modern styling (safe for 2025+ Streamlit)
st.markdown("""
    <style>
        .stChatMessage {
            margin-bottom: 1.1rem !important;
            border-radius: 16px !important;
        }
        .stChatInput > div > div {
            border-radius: 24px !important;
            padding: 0.5rem 1rem;
        }
    </style>
""", unsafe_allow_html=True)

# ────────────────────────────────────────────────
# SIDEBAR
# ────────────────────────────────────────────────
with st.sidebar:
    st.title("🤖 Leave Assistant")
    st.markdown("Ask about leave policy, entitlement, approval, holidays, etc.")
    if "user_id" in st.session_state and st.session_state.user_id:
        st.info(f"**You are:** {st.session_state.user_id}")
    if st.button("Clear conversation", use_container_width=True):
        st.session_state.messages = []
        st.rerun()

# ────────────────────────────────────────────────
# SESSION STATE INIT
# ────────────────────────────────────────────────
defaults = {
    "user_id": "",
    "chat_started": False,
    "messages": [],
}
for key, value in defaults.items():
    if key not in st.session_state:
        st.session_state[key] = value

# ────────────────────────────────────────────────
# LOGIN SCREEN
# ────────────────────────────────────────────────
if not st.session_state.chat_started:
    st.title("Welcome to Leave Policy Assistant")
    st.markdown("Your AI helper for all leave-related questions")

    with st.form("start_form", clear_on_submit=False):
        user_id_input = st.text_input(
            "Enter your name or Employee ID",
            value=st.session_state.user_id,
            placeholder="e.g. Vishal123 or Vishal Patel",
            key="user_id_input"
        )
        if st.form_submit_button("Start Chatting", type="primary", use_container_width=True):
            cleaned = user_id_input.strip()
            if cleaned:
                st.session_state.user_id = cleaned
                st.session_state.chat_started = True
                # Welcome message from bot
                st.session_state.messages.append({
                    "role": "assistant",
                    "content": f"Hi **{cleaned}**! 👋 I'm your leave policy assistant.\nHow may I assist you today?"
                })
                st.rerun()
            else:
                st.error("Please enter a valid name or Employee ID")

else:
    # ────────────────────────────────────────────────
    # MAIN CHAT INTERFACE
    # ────────────────────────────────────────────────
    user_id = st.session_state.user_id
    session_id = f"{user_id}_session_{int(time.time())}"  # optional: prevent collision

    # Show history
    for msg in st.session_state.messages:
        role = "user" if msg["role"] == "user" else "assistant"
        avatar = "👤" if role == "user" else "🤖"

        with st.chat_message(role, avatar=avatar):
            st.markdown(msg["content"])

    # Input
    user_input = st.chat_input("Ask anything about leave policies…")

    if user_input:
        # Append & display user message immediately
        st.session_state.messages.append({"role": "user", "content": user_input})
        with st.chat_message("user", avatar="👤"):
            st.markdown(user_input)

        # Assistant thinking placeholder
        with st.chat_message("assistant", avatar="🤖"):
            placeholder = st.empty()
            placeholder.markdown("▋ Thinking…")

        # ── Async agent logic ───────────────────────────────────
        async def run_agent():
            try:
                # Get or create session
                session = await session_service.get_session(
                    app_name=APP_NAME,
                    user_id=user_id,
                    session_id=session_id,
                )
                if not session:
                    await session_service.create_session(
                        app_name=APP_NAME,
                        user_id=user_id,
                        session_id=session_id,
                    )

                # Prepare message
                user_content = types.Content(
                    role="user",
                    parts=[types.Part.from_text(text=user_input)],
                )

                response_text = ""

                async for event in runner.run_async(
                    user_id=user_id,
                    session_id=session_id,
                    new_message=user_content,
                ):
                    if event.is_final_response():
                        if event.content and event.content.parts:
                            response_text = event.content.parts[0].text
                        break

                return response_text or "(No response generated)"

            except Exception as exc:
                return f"**Error occurred:** {str(exc)}"

        # Execute async function in sync Streamlit context
        try:
            response = asyncio.run(run_agent())
        except RuntimeError as e:
            # Common in some environments when event loop is already running
            loop = asyncio.get_event_loop()
            if loop.is_running():
                response = loop.run_until_complete(run_agent())
            else:
                response = f"Async loop error: {str(e)}"

        # Replace placeholder with real answer
        placeholder.markdown(response)

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

        # Refresh UI
        st.rerun()