import streamlit as st
import time
import random
import json
from datetime import datetime
import pytz
import platform
import uuid
import extra_streamlit_components as stx
import requests
from urllib.parse import quote
from openai import OpenAI

# Set page config
st.set_page_config(page_title="Personalized Real-Time Chat with ArXiv Search and AI", page_icon="💬", layout="wide")

# Initialize cookie manager
cookie_manager = stx.CookieManager()

# File to store chat history and user data
CHAT_FILE = "chat_history.txt"

# Function to save chat history and user data to file
def save_data():
    with open(CHAT_FILE, 'w') as f:
        json.dump({
            'messages': st.session_state.messages,
            'users': st.session_state.users
        }, f)

# Function to load chat history and user data from file
def load_data():
    try:
        with open(CHAT_FILE, 'r') as f:
            data = json.load(f)
            st.session_state.messages = data['messages']
            st.session_state.users = data['users']
    except FileNotFoundError:
        st.session_state.messages = []
        st.session_state.users = []

# Load data at the start
if 'data_loaded' not in st.session_state:
    load_data()
    st.session_state.data_loaded = True

# Function to get or create user
def get_or_create_user():
    user_id = cookie_manager.get(cookie='user_id')
    if not user_id:
        user_id = str(uuid.uuid4())
        cookie_manager.set('user_id', user_id)
    
    user = next((u for u in st.session_state.users if u['id'] == user_id), None)
    if not user:
        user = {
            'id': user_id,
            'name': random.choice(['Alice', 'Bob', 'Charlie', 'David', 'Eve', 'Frank', 'Grace', 'Henry']),
            'browser': f"{platform.system()} - {st.session_state.get('browser_info', 'Unknown')}"
        }
        st.session_state.users.append(user)
        save_data()
    
    return user

# Initialize session state
if 'messages' not in st.session_state:
    st.session_state.messages = []
if 'users' not in st.session_state:
    st.session_state.users = []
if 'current_user' not in st.session_state:
    st.session_state.current_user = get_or_create_user()

# ArXiv search function
def search_arxiv(query):
    base_url = "http://export.arxiv.org/api/query?"
    search_query = f"search_query=all:{quote(query)}&start=0&max_results=5"
    response = requests.get(base_url + search_query)
    
    if response.status_code == 200:
        results = []
        for entry in response.text.split('<entry>')[1:]:
            title = entry.split('<title>')[1].split('</title>')[0]
            summary = entry.split('<summary>')[1].split('</summary>')[0]
            link = entry.split('<id>')[1].split('</id>')[0]
            results.append(f"Title: {title}\nSummary: {summary}\nLink: {link}\n")
        return "\n".join(results)
    else:
        return "Error fetching results from ArXiv."

# Initialize OpenAI client
client = OpenAI(api_key=st.secrets['OPENAI_API_KEY'])
MODEL = "gpt-4-0125-preview"  # Use the appropriate model

# Function to get AI response
def get_ai_response(prompt, context=""):
    try:
        messages = [
            {"role": "system", "content": "You are a helpful assistant in a chat room that can also search ArXiv."},
            {"role": "user", "content": f"Context: {context}\n\nUser Query: {prompt}"}
        ]
        response = client.chat.completions.create(
            model=MODEL,
            messages=messages,
            max_tokens=150,
            temperature=0.7
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"An error occurred: {str(e)}"

# Sidebar for user information and settings
with st.sidebar:
    st.title("User Info")
    st.write(f"Current User: {st.session_state.current_user['name']}")
    st.write(f"Browser: {st.session_state.current_user['browser']}")
    
    new_name = st.text_input("Change your name:")
    if st.button("Update Name"):
        if new_name:
            for user in st.session_state.users:
                if user['id'] == st.session_state.current_user['id']:
                    user['name'] = new_name
                    st.session_state.current_user['name'] = new_name
                    save_data()
                    st.success(f"Name updated to {new_name}")
                    st.rerun()
    
    st.title("Active Users")
    for user in st.session_state.users:
        st.write(f"{user['name']} ({user['browser']})")

# Main chat area
st.title("Personalized Real-Time Chat with ArXiv Search and AI")

# Display chat messages
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# Input for new message
if prompt := st.chat_input("Type your message or ArXiv search query:"):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # Check if it's an ArXiv search query
    if prompt.lower().startswith("arxiv:"):
        query = prompt[6:].strip()
        with st.chat_message("assistant"):
            with st.spinner("Searching ArXiv..."):
                search_results = search_arxiv(query)
                st.markdown(f"Search results for '{query}':\n\n{search_results}")
                
                # Get AI commentary on the search results
                ai_commentary = get_ai_response(f"Provide a brief analysis of these ArXiv search results: {search_results}")
                st.markdown(f"\nAI Analysis:\n{ai_commentary}")
        
        st.session_state.messages.append({"role": "assistant", "content": f"Search results for '{query}':\n\n{search_results}\n\nAI Analysis:\n{ai_commentary}"})
    else:
        # Get AI response for regular chat
        with st.chat_message("assistant"):
            with st.spinner("AI is thinking..."):
                ai_response = get_ai_response(prompt)
                st.markdown(ai_response)
        st.session_state.messages.append({"role": "assistant", "content": ai_response})

    save_data()
    st.rerun()

# Polling for updates
if st.button("Refresh Chat"):
    st.rerun()

# Auto-refresh
if 'last_refresh' not in st.session_state:
    st.session_state.last_refresh = time.time()

if time.time() - st.session_state.last_refresh > 5:  # Refresh every 5 seconds
    st.session_state.last_refresh = time.time()
    st.rerun()