import os
import streamlit as st
import requests
import uuid
import time
from datetime import datetime, timedelta

# Get the API URL from environment variable
API_URL = os.getenv('API_URL')

if not API_URL:
    st.error("API_URL environment variable is not set. Please configure the API URL.")
    st.stop()

# Set session timeout to 15 minutes
SESSION_TIMEOUT = 15 * 60  # 15 minutes in seconds
# Set API request timeout to 2 minutes for Render (cold starts can be slow)
API_TIMEOUT = 120  # seconds

st.title("Longevity Assistant")

# Initialize session state
if "messages" not in st.session_state:
    st.session_state.messages = []
    st.session_state.session_id = str(uuid.uuid4())
    st.session_state.last_activity = datetime.now()
    
    # Get welcome message from API
    try:
        with st.spinner("Connecting to server..."):
            welcome_response = requests.get(f"{API_URL}/welcome", timeout=API_TIMEOUT)
            if welcome_response.status_code == 200:
                welcome_data = welcome_response.json()
                
                # Add assistant welcome message to chat history
                st.session_state.messages.append({
                    "role": "assistant", 
                    "content": welcome_data["welcome_message"],
                    "links": []
                })
                
                # Store suggested questions
                st.session_state.suggested_questions = welcome_data.get("suggested_questions", [])
            else:
                raise Exception(f"Server returned status code: {welcome_response.status_code}")
    except requests.Timeout:
        st.error("Server is taking too long to respond. This might be due to a cold start. Please refresh the page.")
        st.session_state.messages.append({
            "role": "assistant", 
            "content": "The server is warming up. Please wait a moment and refresh the page.",
            "links": []
        })
        st.session_state.suggested_questions = []
    except Exception as e:
        st.error(f"Error connecting to the server: {str(e)}")
        st.session_state.messages.append({
            "role": "assistant", 
            "content": "Welcome to the Longevity Assistant! How can I help you today?",
            "links": []
        })
        st.session_state.suggested_questions = []

elif "last_activity" in st.session_state:
    # Check if session has expired
    time_inactive = (datetime.now() - st.session_state.last_activity).total_seconds()
    if time_inactive > SESSION_TIMEOUT:
        # Reset session
        st.session_state.messages = []
        st.session_state.session_id = str(uuid.uuid4())
    
    # Update last activity time
    st.session_state.last_activity = datetime.now()

# Display chat messages
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.write(message["content"])
        if "links" in message and message["links"]:
            for link in message["links"]:
                st.markdown(f"[{link['name']}]({link['url']})")

# Display suggested questions as buttons (only if no messages from user yet)
if len(st.session_state.messages) == 1 and hasattr(st.session_state, 'suggested_questions'):
    st.write("Try asking about:")
    cols = st.columns(2)
    for i, question in enumerate(st.session_state.suggested_questions):
        with cols[i % 2]:
            if st.button(question, key=f"suggested_{i}"):
                # Add user message to chat history
                st.session_state.messages.append({"role": "user", "content": question})
                
                # Display user message immediately
                with st.chat_message("user"):
                    st.write(question)
                
                # Process the question (reusing the chat input logic)
                with st.chat_message("assistant"):
                    with st.spinner("Thinking..."):
                        try:
                            response = requests.post(
                                f"{API_URL}/chat",
                                json={"session_id": st.session_state.session_id, "message": question},
                                timeout=API_TIMEOUT
                            )
                            
                            if response.status_code == 200:
                                response_data = response.json()
                                # Store response for history
                                st.session_state.messages.append({
                                    "role": "assistant",
                                    "content": response_data["response"],
                                    "links": response_data["links"]
                                })
                                # Display the response
                                st.write(response_data["response"])
                                if response_data["links"]:
                                    for link in response_data["links"]:
                                        st.markdown(f"[{link['name']}]({link['url']})")
                            else:
                                st.error(f"Error: {response.status_code}")
                                st.session_state.messages.append({
                                    "role": "assistant",
                                    "content": "Sorry, I encountered an error while processing your request.",
                                    "links": []
                                })
                        except Exception as e:
                            st.error(f"Error connecting to the server: {str(e)}")
                            st.session_state.messages.append({
                                "role": "assistant",
                                "content": "Sorry, I'm having trouble connecting to the server.",
                                "links": []
                            })
                
                # Force a rerun to update the UI
                st.rerun()

# Chat input
if prompt := st.chat_input():
    # Add user message to chat history
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    # Display user message immediately
    with st.chat_message("user"):
        st.write(prompt)
    
    # Display assistant "thinking" message with spinner
    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            # Get bot response
            try:
                response = requests.post(
                    f"{API_URL}/chat",
                    json={"session_id": st.session_state.session_id, "message": prompt},
                    timeout=API_TIMEOUT
                )
                
                if response.status_code == 200:
                    response_data = response.json()
                    # Store response for history
                    st.session_state.messages.append({
                        "role": "assistant",
                        "content": response_data["response"],
                        "links": response_data["links"]
                    })
                    # Display the response
                    st.write(response_data["response"])
                    if response_data["links"]:
                        for link in response_data["links"]:
                            st.markdown(f"[{link['name']}]({link['url']})")
                else:
                    st.error(f"Error: {response.status_code}")
                    st.session_state.messages.append({
                        "role": "assistant",
                        "content": "Sorry, I encountered an error while processing your request.",
                        "links": []
                    })
            except Exception as e:
                st.error(f"Error connecting to the server: {str(e)}")
                st.session_state.messages.append({
                    "role": "assistant",
                    "content": "Sorry, I'm having trouble connecting to the server.",
                    "links": []
                })