import json
from datetime import datetime

import streamlit as st

# Set page config
st.set_page_config(page_title="Conversation Viewer", layout="wide")

@st.cache_data
def load_conversations(uploaded_file):
    """Load conversations from uploaded JSON file"""
    return json.load(uploaded_file)

def format_timestamp(timestamp):
    """Convert timestamp to readable format"""
    try:
        if timestamp > 1e12:  # milliseconds
            timestamp = timestamp / 1000
        return datetime.fromtimestamp(timestamp).strftime('%Y-%m-%d %H:%M:%S')
    except Exception:
        return "N/A"

def search_in_conversation(conv, search_term):
    """Check if search term exists in conversation"""
    if not search_term:
        return True

    search_term = search_term.lower()

    # Search in conversation name
    if search_term in conv['conv']['name'].lower():
        return True

    # Search in messages
    for msg in conv['messages']:
        if 'content' in msg and search_term in msg['content'].lower():
            return True

    return False

def display_message(msg):
    """Display a single message"""
    if msg['role'] == 'user':
        with st.chat_message("user"):
            st.markdown(msg['content'])
    elif msg['role'] == 'assistant':
        with st.chat_message("assistant"):
            # Remove <think> tags if present
            content = msg['content']
            
            # Handle <think> tags if present
            if '<think>' in content and '</think>' in content:
                start = content.find('<think>')
                end = content.find('</think>') + len('</think>')

                # Extract thinking content and response content
                thinking_content = content[start + len('<think>'):content.find('</think>')]
                response_content = content[end:].strip()

                # Display thinking content in a collapsed expander
                with st.expander("🧠 Internal Thinking", expanded=False):
                    st.code(thinking_content.strip(), language="text", wrap_lines=True)

                # Display the actual response
                if response_content:
                    st.markdown(response_content)
            else:
                # No thinking tags, display content normally
                st.markdown(content)

def main():
    st.title("💬 Conversation Viewer")
    st.write("XSRF:", st.get_option("server.enableXsrfProtection"))

    # File uploader
    json_file = st.file_uploader("Upload JSON file", type="json")

    if json_file is None:
        st.info("👆 Please upload a JSON file containing your conversations to get started.")
        return

    try:
        conversations = load_conversations(json_file)

        # Sidebar for search and filtering
        with st.sidebar:
            st.header("Search & Filter")
            search_term = st.text_input("🔍 Search conversations", "")

            st.divider()

            # Filter conversations
            filtered_convs = [
                conv for conv in conversations
                if search_in_conversation(conv, search_term)
            ]

            st.write(f"**{len(filtered_convs)}** conversations found")

            # Sort options
            sort_by = st.selectbox(
                "Sort by",
                ["Most Recent", "Oldest First", "Alphabetical"]
            )

            if sort_by == "Most Recent":
                filtered_convs.sort(key=lambda x: x['conv'].get('lastModified', 0), reverse=True)
            elif sort_by == "Oldest First":
                filtered_convs.sort(key=lambda x: x['conv'].get('lastModified', 0))
            else:
                filtered_convs.sort(key=lambda x: x['conv']['name'])

            st.divider()

            # Conversation list
            st.subheader("Conversations")
            selected_conv_id = st.radio(
                "Select a conversation",
                options=[conv['conv']['id'] for conv in filtered_convs],
                format_func=lambda x: next(
                    (conv['conv']['name'][:50] + "..." if len(conv['conv']['name']) > 50
                     else conv['conv']['name'])
                    for conv in filtered_convs if conv['conv']['id'] == x
                ),
                label_visibility="collapsed"
            )

        # Main content area
        if filtered_convs:
            # Find selected conversation
            selected_conv = next(
                (conv for conv in filtered_convs if conv['conv']['id'] == selected_conv_id),
                filtered_convs[0]
            )

            # Display conversation details
            # have a border around this box
            st.write(selected_conv['conv']['name'])

            col1, col2, col3 = st.columns(3)
            with col1:
                st.metric("Messages", len(selected_conv['messages']) - 1)  # Exclude root
            with col2:
                st.metric("Last Modified", format_timestamp(selected_conv['conv']['lastModified']))
            with col3:
                st.metric("Conv ID", selected_conv['conv']['id'][:8] + "...")

            st.divider()

            # Display messages
            messages = [msg for msg in selected_conv['messages'] if msg['role'] != 'root']

            for msg in messages:
                if msg['role'] in ['user', 'assistant']:
                    display_message(msg)
        else:
            st.info("No conversations found matching your search.")

    except json.JSONDecodeError:
        st.error("❌ Invalid JSON file format. Please upload a valid JSON file.")
    except Exception as e:
        st.error(f"❌ Error loading conversations: {str(e)}")

if __name__ == "__main__":
    main()