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
from urllib.parse import quote

# 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_path):
    with open(file_path, "rb") as file:
        contents = file.read()
    b64 = base64.b64encode(contents).decode()
    file_name = os.path.basename(file_path)
    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)

# Function to create zip of files
def create_zip_of_files(files):
    import zipfile
    zip_name = "all_files.zip"
    with zipfile.ZipFile(zip_name, 'w') as zipf:
        for file in files:
            zipf.write(file)
    return zip_name

# Function to get zip download link
def get_zip_download_link(zip_file):
    with open(zip_file, 'rb') as f:
        data = f.read()
    b64 = base64.b64encode(data).decode()
    href = f'<a href="data:application/zip;base64,{b64}" download="{zip_file}">Download All</a>'
    return href

# 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_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()