import anthropic
import base64
import json
import os
import pandas as pd
import pytz
import re
import streamlit as st
from datetime import datetime
from gradio_client import Client
from azure.cosmos import CosmosClient, exceptions

# App Configuration
title = "🤖 ArXiv and Claude AI Assistant"
st.set_page_config(page_title=title, layout="wide")

# Cosmos DB configuration
ENDPOINT = "https://acae-afd.documents.azure.com:443/"
Key = os.environ.get("Key")
DATABASE_NAME = os.environ.get("COSMOS_DATABASE_NAME")
CONTAINER_NAME = os.environ.get("COSMOS_CONTAINER_NAME")

# Initialize Anthropic client
anthropic_client = anthropic.Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY"))

# Initialize session state
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

def generate_filename(prompt, file_type):
    """Generate a filename with timestamp and sanitized prompt"""
    central = pytz.timezone('US/Central')
    safe_date_time = datetime.now(central).strftime("%m%d_%H%M")
    safe_prompt = re.sub(r'\W+', '', prompt)[:90]
    return f"{safe_date_time}{safe_prompt}.{file_type}"

def create_file(filename, prompt, response, should_save=True):
    """Create and save a file with prompt and response"""
    if not should_save:
        return
    with open(filename, 'w', encoding='utf-8') as file:
        file.write(f"Prompt:\n{prompt}\n\nResponse:\n{response}")

def save_to_cosmos_db(container, query, response1, response2):
    """Save interaction to Cosmos DB"""
    try:
        if container:
            timestamp = datetime.utcnow().strftime('%Y%m%d%H%M%S%f')
            record = {
                "id": timestamp,
                "name": timestamp,
                "query": query,
                "response1": response1,
                "response2": response2,
                "timestamp": datetime.utcnow().isoformat(),
                "type": "ai_response",
                "version": "1.0"
            }
            container.create_item(body=record)
            st.success(f"Record saved to Cosmos DB with ID: {record['id']}")
    except Exception as e:
        st.error(f"Error saving to Cosmos DB: {str(e)}")

def search_arxiv(query):
    """Search ArXiv using Gradio client"""
    try:
        client = Client("awacke1/Arxiv-Paper-Search-And-QA-RAG-Pattern")
        
        # Get response from Mixtral model
        result_mixtral = client.predict(
            query,
            "mistralai/Mixtral-8x7B-Instruct-v0.1",
            True,
            api_name="/ask_llm"
        )
        
        # Get response from Mistral model
        result_mistral = client.predict(
            query,
            "mistralai/Mistral-7B-Instruct-v0.2",
            True,
            api_name="/ask_llm"
        )
        
        # Get RAG-enhanced response
        result_rag = client.predict(
            query,
            10,  # llm_results_use
            "Semantic Search",
            "mistralai/Mistral-7B-Instruct-v0.2",
            api_name="/update_with_rag_md"
        )
        
        return result_mixtral, result_mistral, result_rag
    except Exception as e:
        st.error(f"Error searching ArXiv: {str(e)}")
        return None, None, None

def main():
    st.title(title)

    # Initialize Cosmos DB client if key is available
    if Key:
        cosmos_client = CosmosClient(ENDPOINT, credential=Key)
        try:
            database = cosmos_client.get_database_client(DATABASE_NAME)
            container = database.get_container_client(CONTAINER_NAME)
        except Exception as e:
            st.error(f"Error connecting to Cosmos DB: {str(e)}")
            container = None
    else:
        st.warning("Cosmos DB Key not found in environment variables")
        container = None

    # Create tabs for different functionalities
    arxiv_tab, claude_tab, history_tab = st.tabs(["ArXiv Search", "Chat with Claude", "History"])

    with arxiv_tab:
        st.header("🔍 ArXiv Search")
        arxiv_query = st.text_area("Enter your research query:", height=100)
        if st.button("Search ArXiv"):
            if arxiv_query:
                with st.spinner("Searching ArXiv..."):
                    result_mixtral, result_mistral, result_rag = search_arxiv(arxiv_query)
                    
                    if result_mixtral:
                        st.subheader("Mixtral Model Response")
                        st.markdown(result_mixtral)
                        
                        st.subheader("Mistral Model Response")
                        st.markdown(result_mistral)
                        
                        st.subheader("RAG-Enhanced Response")
                        if isinstance(result_rag, (list, tuple)) and len(result_rag) > 0:
                            st.markdown(result_rag[0])
                            if len(result_rag) > 1:
                                st.markdown(result_rag[1])
                        
                        # Save results
                        filename = generate_filename(arxiv_query, "md")
                        create_file(filename, arxiv_query, f"{result_mixtral}\n\n{result_mistral}")
                        
                        if container:
                            save_to_cosmos_db(container, arxiv_query, result_mixtral, result_mistral)

    with claude_tab:
        st.header("💬 Chat with Claude")
        user_input = st.text_area("Your message:", height=100)
        if st.button("Send"):
            if user_input:
                with st.spinner("Claude is thinking..."):
                    try:
                        response = anthropic_client.messages.create(
                            model="claude-3-sonnet-20240229",
                            max_tokens=1000,
                            messages=[{"role": "user", "content": user_input}]
                        )
                        
                        claude_response = response.content[0].text
                        st.markdown("### Claude's Response:")
                        st.markdown(claude_response)
                        
                        # Save chat history
                        st.session_state.chat_history.append({
                            "user": user_input,
                            "claude": claude_response,
                            "timestamp": datetime.now().isoformat()
                        })
                        
                        # Save to file
                        filename = generate_filename(user_input, "md")
                        create_file(filename, user_input, claude_response)
                        
                        # Save to Cosmos DB
                        if container:
                            save_to_cosmos_db(container, user_input, claude_response, "")
                            
                    except Exception as e:
                        st.error(f"Error communicating with Claude: {str(e)}")

    with history_tab:
        st.header("📜 Chat History")
        for chat in reversed(st.session_state.chat_history):
            with st.expander(f"Conversation from {chat.get('timestamp', 'Unknown time')}"):
                st.markdown("**Your message:**")
                st.markdown(chat["user"])
                st.markdown("**Claude's response:**")
                st.markdown(chat["claude"])

if __name__ == "__main__":
    main()