import streamlit as st
import anthropic
import os
import base64
import glob
import json
import pytz
from datetime import datetime
from streamlit.components.v1 import html
from PIL import Image
import re

# Set up the Anthropic client
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 = []

# Function to get file download link
def get_download_link(file_contents, file_name):
    b64 = base64.b64encode(file_contents.encode()).decode()
    return f'<a href="data:file/txt;base64,{b64}" download="{file_name}">Download {file_name}</a>'

# Function to generate filename
def generate_filename(prompt, file_type):
    central = pytz.timezone('US/Central')
    safe_date_time = datetime.now(central).strftime("%m%d_%H%M")
    replaced_prompt = prompt.replace(" ", "_").replace("\n", "_")
    safe_prompt = "".join(x for x in replaced_prompt if x.isalnum() or x == "_")[:90]
    return f"{safe_date_time}_{safe_prompt}.{file_type}"

# Function to create file
def create_file(filename, prompt, response, should_save=True):
    if not should_save:
        return
    with open(filename, 'w', encoding='utf-8') as file:
        file.write(prompt + "\n\n" + response)

# Function to load file content
def load_file(file_name):
    with open(file_name, "r", encoding='utf-8') as file:
        content = file.read()
    return content

# Function to display glossary entity
def display_glossary_entity(k):
    search_urls = {
        "🚀🌌ArXiv": lambda k: f"/?q={quote(k)}",
        "📖": lambda k: f"https://en.wikipedia.org/wiki/{quote(k)}",
        "🔍": lambda k: f"https://www.google.com/search?q={quote(k)}",
        "🎥": lambda k: f"https://www.youtube.com/results?search_query={quote(k)}",
    }
    links_md = ' '.join([f"[{emoji}]({url(k)})" for emoji, url in search_urls.items()])
    st.markdown(f"**{k}** <small>{links_md}</small>", unsafe_allow_html=True)

# Streamlit app
def main():
    st.title("Claude 3.5 Sonnet API Demo")

    # Sidebar
    st.sidebar.title("File Operations")
    
    # File management
    all_files = glob.glob("*.md")
    all_files = [file for file in all_files if len(os.path.splitext(file)[0]) >= 10]
    all_files.sort(key=lambda x: (os.path.splitext(x)[1], x), reverse=True)

    if st.sidebar.button("🗑 Delete All"):
        for file in all_files:
            os.remove(file)
        st.rerun()

    if st.sidebar.button("⬇️ Download All"):
        zip_file = create_zip_of_files(all_files)
        st.sidebar.markdown(get_zip_download_link(zip_file), unsafe_allow_html=True)

    for file in all_files:
        col1, col2, col3, col4 = st.sidebar.columns([1,3,1,1])
        with col1:
            if st.button("🌐", key="view_"+file):
                st.session_state.current_file = file
                st.session_state.file_content = load_file(file)
        with col2:
            st.markdown(get_table_download_link(file), unsafe_allow_html=True)
        with col3:
            if st.button("📂", key="edit_"+file):
                st.session_state.current_file = file
                st.session_state.file_content = load_file(file)
        with col4:
            if st.button("🗑", key="delete_"+file):
                os.remove(file)
                st.rerun()

    # Main content area
    user_input = st.text_area("Enter your message:", height=100)

    if st.button("Send"):
        if user_input:
            response = client.messages.create(
                model="claude-3-sonnet-20240229",
                max_tokens=1000,
                messages=[
                    {"role": "user", "content": user_input}
                ]
            )
            
            st.write("Claude's response:")
            st.write(response.content[0].text)

            # Save response to file
            filename = generate_filename(user_input, "md")
            create_file(filename, user_input, response.content[0].text)

            # Add to chat history
            st.session_state.chat_history.append({"user": user_input, "claude": response.content[0].text})

    # Display chat history
    st.subheader("Chat History")
    for chat in st.session_state.chat_history:
        st.text_area("You:", chat["user"], height=100, disabled=True)
        st.text_area("Claude:", chat["claude"], height=200, disabled=True)
        st.markdown("---")

    # File content viewer/editor
    if hasattr(st.session_state, 'current_file'):
        st.subheader(f"Editing: {st.session_state.current_file}")
        new_content = st.text_area("File Content:", st.session_state.file_content, height=300)
        if st.button("Save Changes"):
            with open(st.session_state.current_file, 'w', encoding='utf-8') as file:
                file.write(new_content)
            st.success("File updated successfully!")

    # Image Gallery
    st.subheader("Image Gallery")
    image_files = glob.glob("*.png") + glob.glob("*.jpg") + glob.glob("*.jpeg")
    cols = st.columns(3)
    for idx, image_file in enumerate(image_files):
        with cols[idx % 3]:
            img = Image.open(image_file)
            st.image(img, caption=image_file, use_column_width=True)
            display_glossary_entity(os.path.splitext(image_file)[0])

    # Video Gallery
    st.subheader("Video Gallery")
    video_files = glob.glob("*.mp4")
    for video_file in video_files:
        st.video(video_file)
        display_glossary_entity(os.path.splitext(video_file)[0])

if __name__ == "__main__":
    main()