frontend/app.py

import streamlit as st
import requests
import json
import logging
import uuid
from pathlib import Path
import sys

# --- Path Setup ---
project_root = str(Path(__file__).parent.parent)
if project_root not in sys.path:
    sys.path.append(project_root)

# --- Imports from our backend logic ---
from insucompass.core.agent_orchestrator import app as unified_orchestrator_app
from insucompass.services.zip_client import get_geo_data_from_zip

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# --- Streamlit Page Configuration ---
st.set_page_config(
    page_title="InsuCompass AI",
    page_icon="🧭",
    layout="wide",
)

# --- Header ---
LOGO_PATH = Path(__file__).parent / "assets" / "InsuCompass_Logo.png"
INSU_AVATAR = Path(__file__).parent / "assets" / "robot-chat.jpg"
# Create two columns
col1, col2 = st.columns([0.4, 5])  # Adjust the width ratio as needed
with col1:
    if LOGO_PATH.exists():
        st.image(str(LOGO_PATH), width=70)
with col2:
    st.title("InsuCompass")
st.caption("Your AI guide to U.S. Health Insurance")

# --- Session State Initialization ---
if "thread_id" not in st.session_state:
    st.session_state.thread_id = str(uuid.uuid4())
    st.session_state.user_profile = {}
    st.session_state.chat_history = []
    st.session_state.is_profile_complete = False
    st.session_state.phase = "zip"

# --- Helper Functions ---
@st.cache_data
def get_geodata(_zip_code: str):
    """Calls the zip_client function directly."""
    return get_geo_data_from_zip(_zip_code)

def run_orchestrator(user_message: str):
    """Calls the LangGraph orchestrator directly."""
    thread_config = {"configurable": {"thread_id": st.session_state.thread_id}}
    
    inputs = {
        "user_profile": st.session_state.user_profile,
        "user_message": user_message,
        "is_profile_complete": st.session_state.is_profile_complete,
        "conversation_history": st.session_state.chat_history,
    }
    logger.info(f"Invoking orchestrator for thread_id: {st.session_state.thread_id}")
    
    try:
        final_state = unified_orchestrator_app.invoke(inputs, config=thread_config)
        st.session_state.user_profile = final_state["user_profile"]
        st.session_state.chat_history = final_state["conversation_history"]
        st.session_state.is_profile_complete = final_state["is_profile_complete"]
    except Exception as e:
        logger.error(f"Error invoking orchestrator: {e}", exc_info=True)
        st.error("An internal error occurred in the AI agent. Please try again.")

# --- UI Rendering Functions ---

def display_zip_form():
    st.header("Let's Start with Your Location")
    zip_code_input = st.text_input("Enter your 5-digit ZIP code:", max_chars=5)
    if st.button("Confirm ZIP Code"):
        if zip_code_input and len(zip_code_input) == 5 and zip_code_input.isdigit():
            with st.spinner("Verifying ZIP code..."):
                geo_data = get_geodata(zip_code_input)
            if geo_data:
                st.session_state.user_profile.update(geo_data.to_dict())
                st.session_state.phase = "basic_profile"
                st.rerun()
            else:
                st.error("Invalid ZIP code or could not retrieve location data.")
        else:
            st.error("Please enter a valid 5-digit ZIP code.")

def display_basic_profile_form():
    st.header("Tell Us More About You")
    st.write(f"**Location:** {st.session_state.user_profile.get('city', '')}, {st.session_state.user_profile.get('state', '')}")
    
    with st.form("basic_profile_form"):
        st.subheader("Please provide some basic information:")
        
        # (CORRECTED) Fields now start empty and are validated.
        age = st.number_input("Your Age", min_value=1, max_value=120, value=None, placeholder="Enter your age")
        gender = st.selectbox("Your Gender", [None, "Male", "Female", "Non-binary", "Prefer not to say"], index=0, format_func=lambda x: "Select an option" if x is None else x)
        household_size = st.number_input("Household Size (including yourself)", min_value=1, value=None, placeholder="e.g., 1")
        income = st.number_input("Estimated Annual Household Income ($)", min_value=0, step=1000, value=None, placeholder="e.g., 55000")
        employment_status = st.selectbox("Your Employment Status", [None, "Employed with employer coverage", "Employed without coverage", "Unemployed", "Retired", "Student", "Self-employed"], index=0, format_func=lambda x: "Select an option" if x is None else x)
        citizenship = st.selectbox("Your Citizenship Status", [None, "US Citizen", "Lawful Permanent Resident", "Other legal resident", "Non-resident"], index=0, format_func=lambda x: "Select an option" if x is None else x)
        
        submitted = st.form_submit_button("Start My Personalized Session")
        
        if submitted:
            # Validation check
            if not all([age, gender, household_size, income is not None, employment_status, citizenship]):
                st.error("Please fill out all the fields to continue.")
            else:
                st.session_state.user_profile.update({
                    "age": age, "gender": gender, "household_size": household_size,
                    "income": income, "employment_status": employment_status,
                    "citizenship": citizenship,
                    "medical_history": None, "medications": None, "special_cases": None
                })
                st.session_state.phase = "chat"
                # Trigger the first profile question by calling the orchestrator
                with st.spinner("Starting your personalized profile conversation..."):
                    run_orchestrator("START_PROFILE_BUILDING")
                st.rerun()

def display_chat_interface():
    st.header("Let's Chat!")
    
    # Display chat history from session state
    for message in st.session_state.chat_history:
        if ":" in message:
            role, content = message.split(":", 1)
            if role.lower() == "user":
                with st.chat_message("user", avatar="👤"):
                    st.markdown(content.strip())
            else: # Agent
                with st.chat_message("assistant", avatar=INSU_AVATAR):
                    st.markdown(content.strip())
    
    # User input
    if prompt := st.chat_input("Your message..."):
        with st.spinner("InsuCompass AI is thinking..."):
            run_orchestrator(prompt)
        st.rerun()

# --- Main Application Flow Control ---
if st.session_state.phase == "zip":
    display_zip_form()
elif st.session_state.phase == "basic_profile":
    display_basic_profile_form()
else: # 'chat' phase
    display_chat_interface()

# --- Sidebar for Debugging ---
with st.sidebar:
    st.header("Session State")
    st.json(st.session_state.to_dict(), expanded=False)
    if st.button("Reset Session"):
        st.session_state.clear()
        st.rerun()