#!/usr/bin/env python3
"""
Real-time SWE Agent Tool Monitor
Watch tools being used live as agents work on tasks.
"""

import sys
import json
import time
import os
from datetime import datetime
from rich.console import Console
from rich.live import Live
from rich.table import Table
from rich.panel import Panel
from rich.text import Text
import threading

console = Console()

class ToolMonitor:
    def __init__(self):
        self.tracker_file = '.tool_usage_tracker.json'
        self.initial_data = {}
        self.active_tools = {}
        self.total_calls = 0
        self.start_time = time.time()
        
        # Load initial state
        if os.path.exists(self.tracker_file):
            try:
                with open(self.tracker_file, 'r') as f:
                    self.initial_data = json.load(f)
                    self.total_calls = sum(len(calls) for calls in self.initial_data.values())
            except:
                pass
    
    def create_display_table(self):
        """Create the main display table."""
        table = Table(title="🔧 SWE Agent Tool Usage - Live Monitor", show_header=True)
        table.add_column("Tool Name", style="cyan", width=25)
        table.add_column("Status", justify="center", width=8)
        table.add_column("Calls", justify="right", width=8)
        table.add_column("Last Used", style="dim", width=12)
        table.add_column("Success Rate", justify="right", width=12)
        
        if os.path.exists(self.tracker_file):
            try:
                with open(self.tracker_file, 'r') as f:
                    current_data = json.load(f)
                
                # Calculate new activity
                new_activity = {}
                for tool_name, calls in current_data.items():
                    initial_count = len(self.initial_data.get(tool_name, []))
                    current_count = len(calls)
                    
                    if current_count > initial_count:
                        new_activity[tool_name] = calls[initial_count:]
                        self.active_tools[tool_name] = time.time()
                    
                    # Add to table if tool has been used
                    if calls:
                        latest_call = calls[-1]
                        success_count = sum(1 for call in calls if call.get('success', True))
                        success_rate = f"{success_count/len(calls)*100:.1f}%"
                        
                        # Status indicator
                        if tool_name in self.active_tools and time.time() - self.active_tools[tool_name] < 2:
                            status = "🟢 ACTIVE"
                        else:
                            status = "⚪ IDLE"
                        
                        # Last used time
                        try:
                            last_used = datetime.fromisoformat(latest_call.get('timestamp', '')).strftime('%H:%M:%S')
                        except:
                            last_used = "unknown"
                        
                        table.add_row(
                            tool_name,
                            status, 
                            str(len(calls)),
                            last_used,
                            success_rate
                        )
                
                # Update initial data for next iteration
                self.initial_data = current_data
                new_total = sum(len(calls) for calls in current_data.values())
                
                # Show new activity in console
                if new_total > self.total_calls:
                    for tool_name, new_calls in new_activity.items():
                        for call in new_calls:
                            status = '✅' if call.get('success', True) else '❌'
                            duration = call.get('duration', 0)
                            console.print(f"[green]🔧 {tool_name}[/green]: {status} {duration:.2f}s", end="")
                            if 'parameters' in call:
                                params = str(call['parameters'])[:50] + "..." if len(str(call['parameters'])) > 50 else str(call['parameters'])
                                console.print(f" | {params}")
                            else:
                                console.print()
                
                self.total_calls = new_total
                
            except Exception as e:
                table.add_row("Error", "❌", "0", "N/A", "0%")
        else:
            table.add_row("No data yet", "⏳", "0", "N/A", "0%")
        
        # Create stats panel
        runtime = time.time() - self.start_time
        stats = f"Runtime: {runtime:.1f}s | Total Calls: {self.total_calls}"
        
        return Panel.fit(table, subtitle=stats)
    
    def monitor(self, duration=60):
        """Run the live monitor."""
        console.print("[bold green]🚀 Starting Real-time Tool Monitor[/bold green]")
        console.print(f"Monitoring for {duration} seconds... Press Ctrl+C to stop early")
        console.print()
        
        end_time = time.time() + duration
        
        with Live(self.create_display_table(), refresh_per_second=2) as live:
            try:
                while time.time() < end_time:
                    live.update(self.create_display_table())
                    time.sleep(0.5)
            except KeyboardInterrupt:
                console.print("\n[yellow]Monitoring stopped by user[/yellow]")
        
        console.print("[bold green]✅ Monitoring complete[/bold green]")

def run_task_and_monitor():
    """Run a sample task and monitor it."""
    import subprocess
    import threading
    
    console.print("[bold blue]🎯 Starting SWE Agent task in background...[/bold blue]")
    
    # Start the task in background
    def run_task():
        subprocess.run([
            'python', 'main.py', 
            '--task', 'analyze the project structure and list all Python files',
            '--log-level', 'WARNING'
        ], cwd='.')
    
    task_thread = threading.Thread(target=run_task, daemon=True)
    task_thread.start()
    
    # Give task a moment to start
    time.sleep(2)
    
    # Start monitoring
    monitor = ToolMonitor()
    monitor.monitor(30)

if __name__ == "__main__":
    if len(sys.argv) > 1:
        if sys.argv[1] == "--with-task":
            run_task_and_monitor()
        elif sys.argv[1] == "--monitor":
            duration = int(sys.argv[2]) if len(sys.argv) > 2 else 30
            monitor = ToolMonitor()
            monitor.monitor(duration)
    else:
        console.print("[bold yellow]SWE Agent Tool Monitor[/bold yellow]")
        console.print("Usage:")
        console.print("  python live_monitor.py --monitor [seconds]     # Just monitor")
        console.print("  python live_monitor.py --with-task             # Run task + monitor")