#!/usr/bin/env python3

import json
import sys
import subprocess
import os
from pathlib import Path

def extract_assistant_turns(session_file):
    """Extract assistant messages and their preceding user messages"""
    turns = []
    messages = []
    
    # Read all messages from the JSONL file
    with open(session_file, 'r') as f:
        for line in f:
            try:
                data = json.loads(line.strip())
                if data.get('type') == 'message':
                    messages.append(data)
            except json.JSONDecodeError:
                continue
    
    # Find assistant messages and pair with preceding user messages
    for i, msg in enumerate(messages):
        if msg.get('message', {}).get('role') == 'assistant':
            # Find the most recent user message before this assistant message
            user_msg = None
            for j in range(i-1, -1, -1):
                if messages[j].get('message', {}).get('role') == 'user':
                    user_msg = messages[j]
                    break
            
            if user_msg:
                # Extract content text from user message
                user_content = ""
                user_msg_content = user_msg.get('message', {}).get('content', [])
                if isinstance(user_msg_content, list):
                    for item in user_msg_content:
                        if item.get('type') == 'text':
                            user_content += item.get('text', '')
                elif isinstance(user_msg_content, str):
                    user_content = user_msg_content
                
                # Extract content text from assistant message  
                assistant_content = ""
                assistant_msg_content = msg.get('message', {}).get('content', [])
                if isinstance(assistant_msg_content, list):
                    for item in assistant_msg_content:
                        if item.get('type') == 'text':
                            assistant_content += item.get('text', '')
                elif isinstance(assistant_msg_content, str):
                    assistant_content = assistant_msg_content
                
                if user_content.strip() and assistant_content.strip():
                    turns.append({
                        "user": user_content.strip(),
                        "assistant": assistant_content.strip(),
                        "agent_id": "case",
                        "session": "main"
                    })
    
    # Return last 10 turns
    return turns[-10:] if len(turns) > 10 else turns

def main():
    session_file = "/home/wdjones/.openclaw/agents/main/sessions/8fb5bcf0-cc48-4e32-b34a-0c7bdf93b6b5.jsonl"
    
    if not os.path.exists(session_file):
        print(f"Session file not found: {session_file}", file=sys.stderr)
        return 1
    
    turns = extract_assistant_turns(session_file)
    
    if not turns:
        print("No assistant turns found", file=sys.stderr)
        return 1
    
    print(f"Extracted {len(turns)} assistant turns", file=sys.stderr)
    
    # Pipe each turn to the auto-memory hook
    hook_script = "/home/wdjones/.openclaw/workspace/tools/auto-memory-hook.py"
    
    for i, turn in enumerate(turns):
        try:
            # Convert to JSON and pipe to the hook script
            json_data = json.dumps(turn)
            result = subprocess.run(
                ["python3", hook_script],
                input=json_data,
                text=True,
                capture_output=True
            )
            
            if result.returncode == 0:
                print(f"Turn {i+1}: Processed successfully", file=sys.stderr)
            else:
                print(f"Turn {i+1}: Error - {result.stderr}", file=sys.stderr)
                
        except Exception as e:
            print(f"Turn {i+1}: Exception - {e}", file=sys.stderr)
    
    print("Auto-memory indexing complete", file=sys.stderr)
    return 0

if __name__ == "__main__":
    sys.exit(main())