import os
import json
import requests
from datetime import datetime
import streamlit as st
import pandas as pd

# API configuration
API_URL = "http://localhost:8000/api"

# Set page config
st.set_page_config(
    page_title="BuffaloRAG - UB International Student Assistant",
    page_icon="🦬",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Styling
st.markdown("""
<style>
    .main {
        background-color: #f0f2f6;
    }
    .sidebar .sidebar-content {
        background-color: #2d51a3;
        color: white;
    }
    .stButton>button {
        background-color: #2d51a3;
        color: white;
    }
    .source-card {
        background-color: white;
        padding: 10px;
        border-radius: 5px;
        margin-bottom: 10px;
        border-left: 4px solid #2d51a3;
    }
</style>
""", unsafe_allow_html=True)

# App title
st.title("🦬 BuffaloRAG - UB International Student Assistant")
st.markdown("Ask questions about resources, policies, and procedures for international students at the University at Buffalo.")

# Initialize session state
if 'messages' not in st.session_state:
    st.session_state.messages = []

if 'sources' not in st.session_state:
    st.session_state.sources = []

# Sidebar content
with st.sidebar:
    st.image("https://www.buffalo.edu/content/www/brand/identity/university-logo/_jcr_content/bottompar/image.img.1280.high.jpg/1629127701437.jpg", width=200)
    st.header("Useful Links")
    
    # Common categories
    categories = {
        "Immigration": [
            ("I-20 Information", "https://www.buffalo.edu/international-student-services/immigration/maintainingVisaStatus/i-20.html"),
            ("SEVIS Transfer", "https://www.buffalo.edu/international-student-services/immigration/sevis-transfer.html"),
            ("Visa Information", "https://www.buffalo.edu/international-student-services/immigration/visa.html")
        ],
        "Employment": [
            ("OPT Information", "https://www.buffalo.edu/international-student-services/immigration/f1-student/employment/opt.html"),
            ("CPT Information", "https://www.buffalo.edu/international-student-services/immigration/f1-student/employment/cpt.html"),
            ("On-Campus Employment", "https://www.buffalo.edu/international-student-services/immigration/f1-student/employment/on-campus.html")
        ],
        "Student Life": [
            ("Housing", "https://www.buffalo.edu/campusliving.html"),
            ("Dining", "https://www.buffalo.edu/campusdining.html"),
            ("Student Clubs", "https://www.buffalo.edu/studentlife/life-on-campus/clubs.html")
        ],
        "Academics": [
            ("Academic Calendar", "https://registrar.buffalo.edu/calendars/academic/"),
            ("Course Registration", "https://registrar.buffalo.edu/registration/"),
            ("Library", "https://library.buffalo.edu/")
        ]
    }
    
    # Display links by category
    for category, links in categories.items():
        st.subheader(category)
        for name, url in links:
            st.markdown(f"[{name}]({url})")
    
    # Admin section (hidden in collapsed section)
    with st.expander("Admin Tools"):
        if st.button("Scrape New Content"):
            with st.spinner("Starting scraper..."):
                response = requests.post(f"{API_URL}/scrape", json={
                    "seed_url": "https://www.buffalo.edu/international-student-services.html",
                    "max_pages": 100
                })
                if response.status_code == 200:
                    st.success("Scraping started successfully!")
                else:
                    st.error(f"Error starting scraper: {response.text}")
        
        if st.button("Refresh Index"):
            with st.spinner("Refreshing index..."):
                response = requests.post(f"{API_URL}/refresh-index")
                if response.status_code == 200:
                    st.success("Index refresh started successfully!")
                else:
                    st.error(f"Error refreshing index: {response.text}")

# Main content area - Chat interface
col1, col2 = st.columns([2, 1])

with col1:
    # Display chat messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])
    
    # Input box for new questions
    if prompt := st.chat_input("Ask a question about UB international student services..."):
        # Add user message to chat history
        st.session_state.messages.append({"role": "user", "content": prompt})
        
        # Display user message
        with st.chat_message("user"):
            st.markdown(prompt)
        
        # Call API to get response
        with st.spinner("Thinking..."):
            try:
                response = requests.post(f"{API_URL}/ask", json={"query": prompt})
                
                if response.status_code == 200:
                    result = response.json()
                    answer = result["response"]
                    sources = result["sources"]
                    
                    # Store sources in session state
                    st.session_state.sources = sources
                    
                    # Add assistant message to chat history
                    st.session_state.messages.append({"role": "assistant", "content": answer})
                    
                    # Display assistant message
                    with st.chat_message("assistant"):
                        st.markdown(answer)
                else:
                    st.error(f"Error: {response.text}")
            except Exception as e:
                st.error(f"Error: {str(e)}")

with col2:
    # Display sources for the latest response
    st.subheader("Sources")
    
    if st.session_state.sources:
        for i, source in enumerate(st.session_state.sources):
            st.markdown(f"""
            <div class="source-card">
                <h4>{source['title']}</h4>
                <p><a href="{source['url']}" target="_blank">View source</a></p>
                <small>Relevance: {source['score']:.2f}</small>
            </div>
            """, unsafe_allow_html=True)
    else:
        st.info("Ask a question to see relevant sources.")

# Footer
st.markdown("---")
st.markdown("Built with BuffaloRAG - AI Assistant for International Students at UB")