#!/usr/bin/env python3
"""
SWE Agent Debug Interface - Shows real-time processing and tool usage
"""

import sys
import logging
import time
from rich.console import Console
from rich.panel import Panel
from rich.text import Text
from rich.live import Live
from rich.columns import Columns
from rich.progress import Progress, SpinnerColumn, TextColumn
import threading
import queue
import json
from datetime import datetime

console = Console()

class DebugHandler(logging.Handler):
    def __init__(self, message_queue):
        super().__init__()
        self.message_queue = message_queue
    
    def emit(self, record):
        try:
            msg = self.format(record)
            timestamp = datetime.now().strftime("%H:%M:%S")
            
            # Parse different types of log messages
            if "tool_usage_tracker" in record.name:
                tool_name = "unknown"
                if "completed" in msg:
                    parts = msg.split("completed")
                    if len(parts) > 1:
                        tool_name = parts[1].split("in")[0].strip()
                
                self.message_queue.put({
                    'type': 'tool_usage',
                    'timestamp': timestamp,
                    'tool': tool_name,
                    'message': msg,
                    'level': record.levelname
                })
            
            elif "httpx" in record.name and "POST" in msg:
                self.message_queue.put({
                    'type': 'api_call',
                    'timestamp': timestamp,
                    'message': "🌐 API Call to Anthropic Claude",
                    'level': 'INFO'
                })
            
            elif any(keyword in msg.lower() for keyword in ['processing', 'analyzing', 'creating', 'executing']):
                self.message_queue.put({
                    'type': 'processing',
                    'timestamp': timestamp,
                    'message': msg,
                    'level': record.levelname
                })
            
            else:
                self.message_queue.put({
                    'type': 'general',
                    'timestamp': timestamp,
                    'message': msg,
                    'level': record.levelname
                })
                
        except Exception:
            pass

class SWEAgentDebugger:
    def __init__(self):
        self.message_queue = queue.Queue()
        self.messages = []
        self.current_tools = []
        self.api_calls = 0
        self.start_time = time.time()
        
    def setup_logging(self):
        """Setup logging to capture agent activity."""
        # Clear existing handlers
        root_logger = logging.getLogger()
        root_logger.handlers = []
        
        # Add our debug handler
        debug_handler = DebugHandler(self.message_queue)
        debug_handler.setLevel(logging.INFO)
        
        # Set up formatters
        formatter = logging.Formatter('%(message)s')
        debug_handler.setFormatter(formatter)
        
        root_logger.addHandler(debug_handler)
        root_logger.setLevel(logging.INFO)
        
        # Also add console handler for errors
        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.ERROR)
        console_handler.setFormatter(logging.Formatter('ERROR: %(message)s'))
        root_logger.addHandler(console_handler)
    
    def create_debug_display(self):
        """Create the debug display panels."""
        # Process new messages
        new_messages = []
        while not self.message_queue.empty():
            try:
                msg = self.message_queue.get_nowait()
                new_messages.append(msg)
                self.messages.append(msg)
                
                if msg['type'] == 'tool_usage':
                    self.current_tools.append(msg)
                elif msg['type'] == 'api_call':
                    self.api_calls += 1
                    
            except queue.Empty:
                break
        
        # Keep only last 20 messages
        self.messages = self.messages[-20:]
        self.current_tools = self.current_tools[-10:]
        
        # Create status panel
        runtime = time.time() - self.start_time
        status_text = f"Runtime: {runtime:.1f}s | API Calls: {self.api_calls} | Tools Used: {len(self.current_tools)}"
        
        # Create recent activity panel
        activity_lines = []
        for msg in self.messages[-8:]:  # Show last 8 messages
            level_color = {
                'INFO': 'green',
                'WARNING': 'yellow',
                'ERROR': 'red',
                'DEBUG': 'blue'
            }.get(msg['level'], 'white')
            
            activity_lines.append(f"[{level_color}]{msg['timestamp']}[/{level_color}] {msg['message'][:80]}...")
        
        if not activity_lines:
            activity_lines = ["[dim]Waiting for agent activity...[/dim]"]
        
        activity_panel = Panel(
            "\n".join(activity_lines),
            title="🔄 Recent Activity",
            border_style="blue"
        )
        
        # Create tools panel
        tool_lines = []
        for tool_msg in self.current_tools[-5:]:  # Show last 5 tools
            tool_lines.append(f"[green]{tool_msg['timestamp']}[/green] 🔧 {tool_msg['tool']}")
        
        if not tool_lines:
            tool_lines = ["[dim]No tools used yet[/dim]"]
            
        tools_panel = Panel(
            "\n".join(tool_lines),
            title="🛠️ Tool Usage",
            border_style="green"
        )
        
        # Create main display
        main_display = Columns([activity_panel, tools_panel])
        
        return Panel(
            main_display,
            title=f"🐛 SWE Agent Debug Monitor - {status_text}",
            border_style="cyan"
        )
    
    def run_with_task(self, task):
        """Run SWE Agent with a task and monitor it."""
        import subprocess
        import threading
        
        console.print(f"[bold green]🚀 Starting SWE Agent Debug Session[/bold green]")
        console.print(f"Task: [yellow]{task}[/yellow]")
        console.print()
        
        self.setup_logging()
        
        def run_agent():
            """Run the agent in background."""
            try:
                result = subprocess.run([
                    sys.executable, 'main.py',
                    '--task', task,
                    '--log-level', 'INFO'
                ], 
                cwd='.', 
                capture_output=True, 
                text=True,
                timeout=120  # 2 minute timeout
                )
                
                # Add final result to queue
                if result.returncode == 0:
                    self.message_queue.put({
                        'type': 'processing',
                        'timestamp': datetime.now().strftime("%H:%M:%S"),
                        'message': '✅ Task completed successfully!',
                        'level': 'INFO'
                    })
                else:
                    self.message_queue.put({
                        'type': 'processing', 
                        'timestamp': datetime.now().strftime("%H:%M:%S"),
                        'message': f'❌ Task failed: {result.stderr[:100]}',
                        'level': 'ERROR'
                    })
                    
            except subprocess.TimeoutExpired:
                self.message_queue.put({
                    'type': 'processing',
                    'timestamp': datetime.now().strftime("%H:%M:%S"), 
                    'message': '⏰ Task timed out after 2 minutes',
                    'level': 'WARNING'
                })
            except Exception as e:
                self.message_queue.put({
                    'type': 'processing',
                    'timestamp': datetime.now().strftime("%H:%M:%S"),
                    'message': f'💥 Error: {str(e)}',
                    'level': 'ERROR'
                })
        
        # Start agent in background
        agent_thread = threading.Thread(target=run_agent, daemon=True)
        agent_thread.start()
        
        # Monitor with live display
        with Live(self.create_debug_display(), refresh_per_second=4) as live:
            try:
                while agent_thread.is_alive():
                    live.update(self.create_debug_display())
                    time.sleep(0.25)
                
                # Show final state for a few seconds
                for _ in range(8):
                    live.update(self.create_debug_display()) 
                    time.sleep(0.5)
                    
            except KeyboardInterrupt:
                console.print("\n[yellow]Debug session stopped by user[/yellow]")
        
        # Show summary
        console.print(f"\n[bold]Debug Session Summary:[/bold]")
        console.print(f"Runtime: {time.time() - self.start_time:.1f}s")
        console.print(f"API Calls: {self.api_calls}")
        console.print(f"Tools Used: {len(self.current_tools)}")
        
        if self.current_tools:
            console.print("\n[bold]Tools that were used:[/bold]")
            unique_tools = set(tool['tool'] for tool in self.current_tools)
            for tool in unique_tools:
                console.print(f"  • {tool}")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        console.print("[bold red]Usage:[/bold red]")
        console.print("python debug_agent.py 'your task here'")
        console.print("\n[bold]Example:[/bold]")
        console.print("python debug_agent.py 'write a bubble sort function in Python'")
        sys.exit(1)
    
    task = " ".join(sys.argv[1:])
    debugger = SWEAgentDebugger()
    debugger.run_with_task(task)