import streamlit as st
import requests
from datetime import datetime
import time

# Set page configuration as the first Streamlit command
st.set_page_config(page_title="Evo", page_icon="🤖", layout="wide")

# Define the API endpoint
API_URL = "https://startrz-proagents.hf.space/api/v1/prediction/cce3da9a-f1b9-4bf2-aab8-5462f52ac058"

# Custom CSS for advanced styling
st.markdown("""
    <style>
    .chat-container {
        max-height: 500px;
        overflow-y: auto;
        border: 1px solid #e0e0e0;
        border-radius: 10px;
        padding: 10px;
        background-color: #fafafa;
    }
    .user-message {
        background-color: #d9eaff;
        padding: 10px;
        border-radius: 8px;
        margin: 5px 0;
    }
    .assistant-message {
        background-color: #f0f0f0;
        padding: 10px;
        border-radius: 8px;
        margin: 5px 0;
    }
    .sidebar-btn {
        width: 100%;
        margin: 5px 0;
        background-color: #4CAF50;
        color: white;
        border: none;
        padding: 8px;
        border-radius: 5px;
        cursor: pointer;
    }
    .sidebar-btn:hover {
        background-color: #45a049;
    }
    </style>
""", unsafe_allow_html=True)

# Function to query the API with enhanced error handling
def query(payload):
    try:
        response = requests.post(API_URL, json=payload, timeout=15)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.Timeout:
        return {"error": "Request timed out. Please try again later."}
    except requests.exceptions.RequestException as e:
        return {"error": f"API error: {str(e)}"}

# Simulate typing animation for assistant responses
def type_response(text, delay=0.02):
    placeholder = st.empty()
    for i in range(len(text) + 1):
        placeholder.markdown(f"**Assistant**: {text[:i]}")
        time.sleep(delay)
    return placeholder

# Initialize session state
if 'messages' not in st.session_state:
    st.session_state.messages = [{"role": "assistant", "content": "Hello! I'm Evo, your AI Networking Assistant. How can I assist you today?", "timestamp": datetime.now().strftime("%H:%M:%S")}]
if 'cache' not in st.session_state:
    st.session_state.cache = {}
if 'selected_suggestion' not in st.session_state:
    st.session_state.selected_suggestion = None
if 'feedback' not in st.session_state:
    st.session_state.feedback = {}

# Application title and description
st.title("Evo: Your AI Networking Assistant")
st.write("Explore networking topics with me! Type your question or click a suggestion below.")

# Sidebar with advanced suggested questions
with st.sidebar:
    st.header("Suggested Questions")
    categories = {
        "Basics": ["What is TCP/IP?", "How does DNS work?", "What is an IP address?"],
        "Advanced": ["Explain VLANs", "What is SD-WAN?", "How does BGP work?"],
        "Security": ["What is a firewall?", "How does SSL/TLS work?", "What is zero-trust?"]
    }
    for category, questions in categories.items():
        with st.expander(category, expanded=True):
            for q in questions:
                if st.button(q, key=q, help=f"Click to ask: {q}", on_click=lambda q=q: setattr(st.session_state, 'selected_suggestion', q)):
                    pass
    if st.button("Clear Conversation", key="clear", help="Reset the chat"):
        st.session_state.messages = [{"role": "assistant", "content": "Hello! I'm Evo, your AI Networking Assistant. How can I assist you today?", "timestamp": datetime.now().strftime("%H:%M:%S")}]
        st.session_state.cache = {}
        st.session_state.feedback = {}
        st.rerun()

# Display chat history with feedback
st.markdown('<div class="chat-container">', unsafe_allow_html=True)
for i, message in enumerate(st.session_state.messages):
    with st.chat_message(message["role"]):
        st.markdown(f"**{message['role'].capitalize()} ({message['timestamp']})**: {message['content']}")
        if message["role"] == "assistant" and i > 0:  # Skip initial message
            col1, col2 = st.columns([1, 1])
            with col1:
                if st.button("👍", key=f"up_{i}", help="Like this response"):
                    st.session_state.feedback[i] = "positive"
            with col2:
                if st.button("👎", key=f"down_{i}", help="Dislike this response"):
                    st.session_state.feedback[i] = "negative"
            if i in st.session_state.feedback:
                st.write(f"Feedback: {'Liked' if st.session_state.feedback[i] == 'positive' else 'Disliked'}")
st.markdown('</div>', unsafe_allow_html=True)

# Function to handle queries (user input or suggestions)
def handle_query(question):
    timestamp = datetime.now().strftime("%H:%M:%S")
    st.session_state.messages.append({"role": "user", "content": question, "timestamp": timestamp})
    with st.chat_message("user"):
        st.markdown(f"**User ({timestamp})**: {question}")
    
    if question in st.session_state.cache:
        answer = st.session_state.cache[question]
    else:
        with st.spinner("Evo is thinking..."):
            output = query({"question": question})
            answer = output.get("answer", output.get("error", "Sorry, I couldn’t fetch a response."))
            st.session_state.cache[question] = answer
    
    timestamp = datetime.now().strftime("%H:%M:%S")
    st.session_state.messages.append({"role": "assistant", "content": answer, "timestamp": timestamp})
    with st.chat_message("assistant"):
        type_response(f"**Assistant ({timestamp})**: {answer}")

# Process user input or selected suggestion
if user_input := st.chat_input("Ask Evo anything about networking..."):
    handle_query(user_input)
elif st.session_state.selected_suggestion:
    handle_query(st.session_state.selected_suggestion)
    st.session_state.selected_suggestion = None  # Reset after processing