"""
Agent Architecture Visualizer
Based on: "Dive into Claude Code" (arXiv:2604.14228)
Demonstrates the core agent loop, permission gates, and context management.
"""

import gradio as gr
import time
import random

# Lazy loading - no imports at module level for ML models

def get_demo_state():
    """Initialize or reset the agent demo state."""
    return {
        "iteration": 0,
        "context_tokens": 0,
        "tools_executed": [],
        "permissions_checked": [],
        "status": "idle"
    }

# Permission modes from the paper
PERMISSION_MODES = {
    "read_only": {"risk": "low", "desc": "Read files without modification", "color": "green"},
    "write_file": {"risk": "medium", "desc": "Create or modify files", "color": "yellow"},
    "shell_safe": {"risk": "medium", "desc": "Run safe shell commands", "color": "yellow"},
    "shell_dangerous": {"risk": "high", "desc": "Destructive or network commands", "color": "red"},
    "mcp_read": {"risk": "low", "desc": "Read from MCP servers", "color": "green"},
    "mcp_write": {"risk": "high", "desc": "Write via MCP servers", "color": "red"},
    "subagent": {"risk": "high", "desc": "Spawn isolated subagents", "color": "red"}
}

# Simulated tool outputs
TOOL_CATALOG = {
    "read_file": {"mode": "read_only", "duration": 0.5, "output": "File content loaded into context"},
    "write_file": {"mode": "write_file", "duration": 0.7, "output": "File modified successfully"},
    "list_dir": {"mode": "read_only", "duration": 0.3, "output": "Directory listing retrieved"},
    "shell_ls": {"mode": "shell_safe", "duration": 0.4, "output": "Command executed: ls -la"},
    "shell_rm": {"mode": "shell_dangerous", "duration": 0.6, "output": "⚠️ REQUIRES EXPLICIT CONFIRMATION"},
    "mcp_search": {"mode": "mcp_read", "duration": 0.8, "output": "MCP tool result retrieved"},
    "spawn_subagent": {"mode": "subagent", "duration": 1.2, "output": "🔄 Subagent spawned in isolated worktree"}
}

def simulate_agent_loop(num_iterations: int, enable_permissions: bool, enable_compaction: bool):
    """
    Simulate the core agent loop with optional safety features.
    
    Based on the paper's insight: "The core of the system is a simple while-loop 
    that calls the model, runs tools, and repeats. Most of the code lives in the 
    systems around this loop."
    """
    if num_iterations > 10:
        num_iterations = 10  # Safety limit
    
    state = get_demo_state()
    trace = []
    trace.append("🚀 Agent Loop Started")
    trace.append("=" * 50)
    
    for i in range(num_iterations):
        state["iteration"] = i + 1
        
        # Step 1: Model generates tool request
        tool_name = random.choice(list(TOOL_CATALOG.keys()))
        tool_info = TOOL_CATALOG[tool_name]
        trace.append(f"\n📍 Iteration {i+1}: Model requests '{tool_name}'")
        
        # Step 2: Permission classification (if enabled)
        if enable_permissions:
            mode_info = PERMISSION_MODES[tool_info["mode"]]
            trace.append(f"   🔒 Permission Check: {tool_info['mode']} (risk: {mode_info['risk']})")
            
            if mode_info["risk"] == "high":
                trace.append(f"   ⚠️  HIGH RISK: Requires explicit user confirmation")
                trace.append(f"   ✓ User approved execution")
            elif mode_info["risk"] == "medium":
                trace.append(f"   ℹ️  MEDIUM RISK: Logged for review")
            else:
                trace.append(f"   ✓ LOW RISK: Auto-approved")
            
            state["permissions_checked"].append(tool_info["mode"])
        
        # Step 3: Execute tool
        time.sleep(0.1)  # Simulate processing
        trace.append(f"   🔧 Executing: {tool_info['output']}")
        state["tools_executed"].append(tool_name)
        
        # Step 4: Context grows
        state["context_tokens"] += random.randint(50, 200)
        trace.append(f"   📊 Context: ~{state['context_tokens']} tokens")
        
        # Step 5: Context compaction (if enabled and threshold reached)
        if enable_compaction and state["context_tokens"] > 300:
            trace.append(f"   🗜️  CONTEXT COMPACTION TRIGGERED")
            trace.append(f"      Layer 1: Summarize old messages")
            trace.append(f"      Layer 2: Extract key facts")
            trace.append(f"      Layer 3: Compress tool outputs")
            trace.append(f"      Layer 4: Archive file contents")
            trace.append(f"      Layer 5: Summarize with LLM")
            trace.append(f"   ✓ Compaction complete: {state['context_tokens']} → ~150 tokens")
            state["context_tokens"] = 150
    
    trace.append("\n" + "=" * 50)
    trace.append("✅ Agent Loop Complete")
    trace.append(f"📈 Stats: {state['iteration']} iterations, {len(state['tools_executed'])} tools, {len(state['permissions_checked'])} permission checks")
    
    return "\n".join(trace)

def visualize_permission_modes():
    """Display the 7 permission modes from Claude Code."""
    output = ["🔐 Permission Mode Classification", "=" * 50]
    
    for mode, info in PERMISSION_MODES.items():
        emoji = "🟢" if info["color"] == "green" else "🟡" if info["color"] == "yellow" else "🔴"
        output.append(f"\n{emoji} {mode.upper()}")
        output.append(f"   Risk Level: {info['risk'].upper()}")
        output.append(f"   Description: {info['desc']}")
    
    output.append("\n" + "=" * 50)
    output.append("💡 Insight: ML classifier assigns modes based on tool type and arguments")
    
    return "\n".join(output)

def show_context_compaction():
    """Visualize the 5-layer compaction pipeline."""
    return """
🗜️ Context Compaction Pipeline (5 Layers)
==========================================

When context window approaches limit:

Layer 1: Message Summarization
   └─> Compress old conversation turns

Layer 2: Fact Extraction  
   └─> Extract key facts into structured format

Layer 3: Tool Output Compression
   └─> Summarize verbose tool outputs

Layer 4: File Content Archival
   └─> Replace full file contents with references

Layer 5: LLM-Based Summarization
   └─> Use model to generate compact representation

📊 Result: Maintains semantic coverage while reducing tokens
"""

def show_subagent_delegation():
    """Explain subagent delegation with worktree isolation."""
    return """
🔄 Subagent Delegation Architecture
====================================

Parent Agent                    Subagent
────────────                    ────────
     │                              │
     │── Spawn with worktree ───────>│
     │   (isolated git worktree)      │
     │                              │
     │<── Return results ────────────│
     │   (read-only access to         │
     │    parent context)             │
     │                              │
     ✓ Parent maintains full control
     ✓ Subagent cannot modify parent state
     ✓ Clean termination guaranteed

Use Case: Parallel tool execution, sandboxed experiments
"""

def create_interface():
    """Create the Gradio interface."""
    with gr.Blocks(title="Agent Architecture Visualizer") as demo:
        gr.Markdown("""
        # 🔄 Agent Architecture Visualizer
        
        Interactive demo based on **"Dive into Claude Code: The Design Space of Today's and Future AI Agent Systems"** 
        (arXiv:2604.14228)
        
        This Space visualizes the key architectural insights from the paper:
        - The simple while-loop at the core
        - Permission classification system
        - Context compaction pipeline
        - Subagent delegation
        """)
        
        with gr.Tab("🔄 Agent Loop Simulator"):
            with gr.Row():
                with gr.Column():
                    iterations = gr.Slider(1, 10, value=3, step=1, label="Iterations")
                    enable_perms = gr.Checkbox(value=True, label="Enable Permission Gates")
                    enable_compaction = gr.Checkbox(value=True, label="Enable Context Compaction")
                    run_btn = gr.Button("▶️ Run Simulation", variant="primary")
                
                with gr.Column():
                    output = gr.Textbox(label="Execution Trace", lines=25, max_lines=30)
            
            run_btn.click(
                simulate_agent_loop,
                inputs=[iterations, enable_perms, enable_compaction],
                outputs=output
            )
        
        with gr.Tab("🔐 Permission Modes"):
            gr.Textbox(
                value=visualize_permission_modes(),
                label="7-Mode Permission Classification",
                lines=20,
                interactive=False
            )
        
        with gr.Tab("🗜️ Context Compaction"):
            gr.Textbox(
                value=show_context_compaction(),
                label="5-Layer Compaction Pipeline",
                lines=18,
                interactive=False
            )
        
        with gr.Tab("🔄 Subagent Delegation"):
            gr.Textbox(
                value=show_subagent_delegation(),
                label="Worktree Isolation Architecture",
                lines=18,
                interactive=False
            )
        
        gr.Markdown("""
        ---
        📄 **Paper**: [Dive into Claude Code](https://huggingface.co/papers/2604.14228) | 
        🔗 **Code**: [github.com/VILA-Lab/Dive-into-Claude-Code](https://github.com/VILA-Lab/Dive-into-Claude-Code)
        """)
    
    return demo

# Create and launch
app = create_interface()

if __name__ == "__main__":
    app.launch()