#!/usr/bin/env python3
"""
Extract chat history from OpenClaw session URL and generate JS data file.
Usage: python3 extract_chat.py <session-url-or-key> [output-dir]
"""

import json
import sys
import os
import re
from urllib.parse import urlparse, parse_qs
from datetime import datetime, timezone, timedelta

tz_beijing = timezone(timedelta(hours=8))

def parse_time(ts_str):
    """Parse ISO timestamp string to Beijing time string."""
    try:
        ts_str = ts_str.replace('Z', '')
        dt = datetime.fromisoformat(ts_str)
        dt_utc = dt.replace(tzinfo=timezone.utc)
        dt_beijing = dt_utc.astimezone(tz_beijing)
        return dt_beijing.strftime('%Y-%m-%d %H:%M:%S')
    except:
        return None

def extract_session_key(url_or_key):
    """Extract session key from URL or return as-is if already a key."""
    if 'session=' in url_or_key:
        parsed = urlparse(url_or_key)
        params = parse_qs(parsed.query)
        if 'session' in params:
            return params['session'][0]
    return url_or_key

def find_session_file(session_key):
    """Find the JSONL file path for a given session key."""
    # session_key format: agent:c108-v1-1811:session-name
    parts = session_key.split(':')
    if len(parts) < 2:
        return None
    agent_id = parts[1]
    sessions_json_path = f'/root/.openclaw/agents/{agent_id}/sessions/sessions.json'
    
    if not os.path.exists(sessions_json_path):
        return None
    
    with open(sessions_json_path, 'r') as f:
        sessions = json.load(f)
    
    if session_key in sessions:
        return sessions[session_key].get('sessionFile')
    return None

def parse_jsonl_to_conversations(jsonl_path):
    """Parse JSONL file and build conversation pairs."""
    messages = []
    with open(jsonl_path, 'r') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            obj = json.loads(line)
            if obj.get('type') == 'message':
                role = obj['message']['role']
                content = obj['message'].get('content', [])
                timestamp = obj.get('timestamp', '')
                
                if isinstance(content, list):
                    text = ''.join([c.get('text', '') for c in content if c.get('type') == 'text'])
                else:
                    text = str(content)
                
                messages.append({
                    'role': role,
                    'timestamp': timestamp,
                    'text': text.strip(),
                    'has_text': bool(text.strip())
                })
    
    # Build conversation pairs
    conversations = []
    i = 0
    while i < len(messages):
        msg = messages[i]
        if msg['role'] == 'user':
            user_text = msg['text']
            user_time = msg['timestamp']
            
            # Collect ALL consecutive AI messages until next user
            ai_messages = []
            j = i + 1
            while j < len(messages) and messages[j]['role'] != 'user':
                if messages[j]['role'] == 'assistant':
                    ai_messages.append(messages[j])
                j += 1
            
            # Merge all AI messages with text
            ai_text = '\n\n'.join([m['text'] for m in ai_messages if m['has_text']])
            ai_time = ai_messages[-1]['timestamp'] if ai_messages else ''
            
            conversations.append({
                'user': user_text,
                'user_time': parse_time(user_time) if user_time else '',
                'ai': ai_text,
                'ai_time': parse_time(ai_time) if ai_time else ''
            })
            i += 1
        else:
            i += 1
    
    return conversations

def generate_js(conversations, session_key, output_path):
    """Generate JS file from conversations."""
    session_name = session_key.split(':')[-1] if ':' in session_key else session_key
    
    output_data = {
        'session': session_name,
        'session_id': session_key,
        'total_pairs': len(conversations),
        'initiator': 'session initiator',
        'conversations': conversations
    }
    
    js_content = f'''// Chat History - {session_name}
// Session Key: {session_key}
// Generated by chat-history-extractor skill

const chatData = {json.dumps(output_data, ensure_ascii=False, indent=2)};
'''
    
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(js_content)
    
    return len(conversations)

def main():
    if len(sys.argv) < 2:
        print("Usage: python3 extract_chat.py <session-url-or-key> [output-dir]")
        sys.exit(1)
    
    url_or_key = sys.argv[1]
    output_dir = sys.argv[2] if len(sys.argv) > 2 else os.getcwd()
    
    # Extract session key
    session_key = extract_session_key(url_or_key)
    print(f"Session key: {session_key}")
    
    # Find session file
    jsonl_path = find_session_file(session_key)
    if not jsonl_path:
        print(f"ERROR: Could not find session file for {session_key}")
        sys.exit(1)
    print(f"JSONL path: {jsonl_path}")
    
    # Parse conversations
    conversations = parse_jsonl_to_conversations(jsonl_path)
    print(f"Found {len(conversations)} conversation pairs")
    
    # Generate JS file
    os.makedirs(output_dir, exist_ok=True)
    js_path = os.path.join(output_dir, 'chat_history.js')
    count = generate_js(conversations, session_key, js_path)
    print(f"Generated: {js_path}")
    print(f"Total pairs: {count}")

if __name__ == '__main__':
    main()
