#!/usr/bin/env python3
import json
import sys
import subprocess

def extract_assistant_turns(session_file, num_turns=10):
    """Extract the last N assistant turns with their preceding user messages"""
    with open(session_file, 'r') as f:
        lines = f.readlines()
    
    messages = []
    for line in lines:
        try:
            data = json.loads(line.strip())
            if data.get('type') == 'message' and 'message' in data:
                msg = data['message']
                if msg.get('role') in ['user', 'assistant']:
                    messages.append(msg)
        except json.JSONDecodeError:
            continue
    
    # Find assistant messages and their preceding user messages
    assistant_turns = []
    for i in range(len(messages)):
        if messages[i].get('role') == 'assistant':
            # Look for the most recent user message before this assistant message
            user_msg = None
            for j in range(i-1, -1, -1):
                if messages[j].get('role') == 'user':
                    user_msg = messages[j]
                    break
            
            if user_msg:
                # Extract content from messages
                user_content = ""
                if isinstance(user_msg.get('content'), list):
                    for content in user_msg['content']:
                        if content.get('type') == 'text':
                            user_content = content.get('text', '')
                            break
                else:
                    user_content = user_msg.get('content', '')
                
                assistant_content = ""
                if isinstance(messages[i].get('content'), list):
                    for content in messages[i]['content']:
                        if content.get('type') == 'text':
                            assistant_content = content.get('text', '')
                            break
                else:
                    assistant_content = messages[i].get('content', '')
                
                if user_content and assistant_content:
                    assistant_turns.append({
                        'user': user_content,
                        'assistant': assistant_content
                    })
    
    # Return the last N turns
    return assistant_turns[-num_turns:]

def main():
    session_file = sys.argv[1] if len(sys.argv) > 1 else None
    if not session_file:
        print("Usage: python3 extract_session_turns.py <session_file>")
        sys.exit(1)
    
    turns = extract_assistant_turns(session_file)
    
    for turn in turns:
        json_data = {
            "user": turn['user'],
            "assistant": turn['assistant'],
            "agent_id": "case",
            "session": "main"
        }
        
        # Pipe to auto-memory hook
        try:
            subprocess.run([
                'python3', '/home/wdjones/.openclaw/workspace/tools/auto-memory-hook.py'
            ], input=json.dumps(json_data), text=True, check=True)
        except subprocess.CalledProcessError as e:
            print(f"Error running auto-memory-hook: {e}")

if __name__ == "__main__":
    main()