#!/usr/bin/env python3
import sqlite3
import json
from pathlib import Path
from datetime import datetime
import os
import base64

def extract_conversations_from_db(db_path):
    """Extract all conversations from the database with detailed error handling."""
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        
        
        # Get all bubbleId entries which contain the actual chat messages
        cursor.execute("SELECT key, value FROM cursorDiskKV WHERE key LIKE 'bubbleId:%'")
        bubble_rows = cursor.fetchall()
        print(f"Found {len(bubble_rows)} bubbleId entries")
        
        # Group messages by composer ID
        conversations_by_composer = {}
        
        for key, value in bubble_rows:
            if value is None:
                continue
                
            try:
                # Parse the key to extract composer ID and bubble ID
                # Format: bubbleId:composerId:bubbleId
                parts = key.split(':')
                if len(parts) >= 3:
                    composer_id = parts[1]
                    bubble_id = parts[2]
                    
                    # Decode the message data
                    try:
                        data = json.loads(value)
                    except json.JSONDecodeError:
                        try:
                            decoded = base64.b64decode(value)
                            data = json.loads(decoded)
                        except:
                            continue
                    
                    # Only include bubbles that have text content
                    if isinstance(data, dict) and 'text' in data and data['text'].strip():
                        # Initialize composer conversation if not exists
                        if composer_id not in conversations_by_composer:
                            conversations_by_composer[composer_id] = []
                        
                        # Add message to conversation
                        conversations_by_composer[composer_id].append({
                            'id': bubble_id,
                            'data': data,
                            'key': key
                        })
                    
            except Exception as e:
                print(f"Error processing bubble {key}: {e}")
        
        # Convert to conversations list
        conversations = []
        for composer_id, messages in conversations_by_composer.items():
            if messages:  # Only include composers with messages
                conversations.append({
                    'id': composer_id,
                    'messages': messages
                })
        
        print(f"Found {len(conversations)} conversations with messages")
        
        conn.close()
        return conversations
    except Exception as e:
        print(f"Error accessing database {db_path}: {e}")
        return []

def format_message(msg_data):
    """Format a single message with all its components."""
    if not isinstance(msg_data, dict) or 'data' not in msg_data:
        return "Invalid message format"
    
    msg = msg_data['data']
    formatted = []
    
    # Add timestamp if available
    timestamp_str = ""
    if 'timingInfo' in msg and 'clientStartTime' in msg['timingInfo']:
        timestamp = datetime.fromtimestamp(
            msg['timingInfo']['clientStartTime'] / 1000
        ).strftime("%Y-%m-%d %H:%M:%S")
        timestamp_str = f" ({timestamp})"
    
    # Determine message role based on type
    msg_type = msg.get('type', 0)
    if msg_type == 1:  # User message
        role = "User"
    elif msg_type == 2:  # Assistant message
        role = "Assistant"
    else:
        role = "Message"
    
    formatted.append(f"**{role}**{timestamp_str}:")
    
    # Add the main text content
    text = msg.get('text', '').strip()
    if text:
        formatted.append(text)
    
    # Add code blocks if present
    if 'codeBlocks' in msg and msg['codeBlocks']:
        for block in msg['codeBlocks']:
            if isinstance(block, dict):
                lang = block.get('language', '')
                code = block.get('code', '')
                if code:
                    formatted.append(f"```{lang}\n{code}\n```")
    
    # Add tool results if present
    if 'toolResults' in msg and msg['toolResults']:
        formatted.append("\n**Tool Results:**")
        for result in msg['toolResults']:
            if isinstance(result, dict):
                tool_name = result.get('toolName', 'Unknown Tool')
                formatted.append(f"- {tool_name}")
                if 'result' in result:
                    formatted.append(f"  Result: {result['result']}")
    
    return "\n\n".join(formatted) if formatted else "Empty message"

def main():
    """Main function to export all chat history."""
    cursor_dir = Path(os.path.expanduser("~")) / "Library/Application Support/Cursor"
    global_db = cursor_dir / "User/globalStorage/state.vscdb"

    print("Extracting conversations from database...")

    print(f"Extracting conversations from {global_db}")
    conversations = extract_conversations_from_db(global_db)
    
    # Sort conversations by first message timestamp (if available)
    def get_conversation_timestamp(conv):
        if conv['messages']:
            # Try to find a timestamp in the first message
            first_msg = conv['messages'][0]['data']
            if isinstance(first_msg, dict) and 'timingInfo' in first_msg:
                timing = first_msg['timingInfo']
                if 'clientStartTime' in timing:
                    return timing['clientStartTime']
        return 0
    
    conversations.sort(key=get_conversation_timestamp)

    print("Generating markdown...")
    
    # Generate markdown
    output = ["# Cursor Chat History\n"]
    for conv in conversations:
        output.append(f"## Conversation {conv['id']}\n")
        output.append(f"*Messages: {len(conv['messages'])}*\n")
        
        # Sort messages within conversation by timestamp
        def get_message_timestamp(msg):
            if 'data' in msg and 'timingInfo' in msg['data']:
                timing = msg['data']['timingInfo']
                if 'clientStartTime' in timing:
                    return timing['clientStartTime']
            return 0
        
        sorted_messages = sorted(conv['messages'], key=get_message_timestamp)
        
        for msg in sorted_messages:
            output.append(format_message(msg))
            output.append("\n---\n")
        
        output.append("\n")
    
    # Write to file
    output_dir = Path("chat_history_exports")
    output_dir.mkdir(exist_ok=True)
    output_file = output_dir / f"cursor_chat_history_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md"
    
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write("\n".join(output))
    
    print(f"Exported {len(conversations)} conversations to {output_file}")

if __name__ == "__main__":
    main()