"""
Subagent System

Spawns focused subagents for decomposed tasks.

Inspired by Claude Code's Task tool:
- Agents can spawn subagents with specific roles
- Subagents have filtered tool access based on type
- Results flow back to the parent agent
"""

import json
import time
from pathlib import Path
from typing import Optional

from rich.console import Console

from ai_coder.llm.interface import LLMProvider, Message
from ai_coder.tools.base import ToolRegistry

console = Console()


# Agent type configurations
AGENT_TYPES = {
    "explore": {
        "description": "Read-only agent for exploring code, finding files, searching",
        "tools": ["read_file", "list_files", "search_files", "run_command"],
        "prompt": "You are an exploration agent. Search and analyze, but never modify files. Return a concise summary.",
    },
    "code": {
        "description": "Full agent for implementing features and fixing bugs",
        "tools": "*",  # All tools
        "prompt": "You are a coding agent. Implement the requested changes efficiently. Write clean, well-tested code.",
    },
    "plan": {
        "description": "Planning agent for designing implementation strategies",
        "tools": ["read_file", "list_files", "search_files", "run_command"],
        "prompt": "You are a planning agent. Analyze the codebase and output a numbered implementation plan. Do NOT make changes.",
    },
    "review": {
        "description": "Code review agent for finding bugs and quality issues",
        "tools": ["read_file", "list_files", "search_files", "run_command"],
        "prompt": "You are a code review agent. Find bugs, security issues, and code quality problems. Only report issues with confidence >= 70%.",
    },
    "architect": {
        "description": "Architecture agent for system design decisions",
        "tools": ["read_file", "list_files", "search_files", "run_command"],
        "prompt": "You are an architecture agent. Design scalable, maintainable solutions following project patterns.",
    },
    "test": {
        "description": "Testing agent for writing and running tests",
        "tools": "*",
        "prompt": "You are a testing agent. Write comprehensive tests following TDD principles. Use existing test patterns.",
    },
    "debug": {
        "description": "Debugging agent for diagnosing and fixing issues",
        "tools": "*",
        "prompt": "You are a debugging agent. Systematically diagnose the issue: reproduce, isolate, inspect, fix.",
    },
    "security": {
        "description": "Security audit agent for finding vulnerabilities",
        "tools": ["read_file", "list_files", "search_files", "run_command"],
        "prompt": "You are a security audit agent. Check for OWASP Top 10, injection, auth issues, and sensitive data exposure.",
    },
    "refactor": {
        "description": "Refactoring agent for improving code quality",
        "tools": "*",
        "prompt": "You are a refactoring agent. Improve code quality without changing behavior. Keep tests passing.",
    },
}


def get_agent_descriptions() -> str:
    """Generate agent type descriptions for system prompt."""
    return "\n".join(
        f"- {name}: {cfg['description']}"
        for name, cfg in AGENT_TYPES.items()
    )


def run_subagent(
    llm: LLMProvider,
    tool_registry: ToolRegistry,
    project_root: Path,
    description: str,
    prompt: str,
    agent_type: str,
    max_iterations: int = 15,
) -> str:
    """
    Execute a subagent task.
    
    Args:
        llm: LLM provider instance
        tool_registry: Tool registry for tool access
        project_root: Project root directory
        description: Short task description
        prompt: Detailed instructions
        agent_type: Type of agent (explore, code, plan, etc.)
        max_iterations: Max tool iterations
    
    Returns:
        Subagent's final text response
    """
    if agent_type not in AGENT_TYPES:
        return f"Error: Unknown agent type '{agent_type}'. Available: {', '.join(AGENT_TYPES.keys())}"
    
    config = AGENT_TYPES[agent_type]
    
    # Build subagent system prompt
    sub_system = f"""You are a {agent_type} subagent working in {project_root}.

{config['prompt']}

Complete the task and return a clear, concise summary of what you found or did."""
    
    # Filter tools based on agent type
    from ai_coder.core.prompts import get_tool_schemas
    all_tools = get_tool_schemas()
    
    if config["tools"] == "*":
        tools = all_tools
    else:
        tools = [t for t in all_tools if t["function"]["name"] in config["tools"]]
    
    # Build tool instructions
    tool_names = [t["function"]["name"] for t in tools]
    tool_instructions = f"""

## Available Tools: {', '.join(tool_names)}

To call a tool, output JSON:
```json
{{"name": "tool_name", "arguments": {{"arg1": "value1"}}}}
```

Output ONLY ONE tool call per response. When done, provide your final summary WITHOUT a tool call."""
    
    sub_system += tool_instructions
    
    messages = [
        Message(role="system", content=sub_system),
        Message(role="user", content=prompt),
    ]
    
    console.print(f"  [{agent_type}] {description}")
    start = time.time()
    tool_count = 0
    
    import sys
    
    for _ in range(max_iterations):
        response = llm.complete(messages, tools=tools)
        
        # Check for tool calls (native or extracted)
        tool_calls = response.tool_calls
        
        if not tool_calls and response.content:
            # Try to extract JSON tool calls from text
            import re
            patterns = [
                r'```json\s*(\{.*?\})\s*```',
                r'```\s*(\{.*?\})\s*```',
                r'(\{"name":\s*"[^"]+",\s*"arguments":\s*\{.*?\}\})',
            ]
            for pattern in patterns:
                matches = re.findall(pattern, response.content, re.DOTALL)
                for match in matches:
                    try:
                        data = json.loads(match)
                        name = data.get("name", "")
                        args = data.get("arguments", {})
                        if name and name in tool_names:
                            from ai_coder.llm.interface import ToolCall
                            tool_calls = [ToolCall(
                                id=f"sub_{name}_{tool_count}",
                                name=name,
                                arguments=args if isinstance(args, dict) else {},
                            )]
                            break
                    except (json.JSONDecodeError, KeyError):
                        continue
                if tool_calls:
                    break
        
        if not tool_calls:
            # No tool calls - subagent is done
            break
        
        # Execute tools
        for tc in tool_calls:
            tool_count += 1
            tool = tool_registry.get(tc.name)
            if tool:
                try:
                    result = tool.execute(**tc.arguments)
                    output = result.output[:3000] if result.output else ""
                    if result.error:
                        output = f"Error: {result.error}"
                except Exception as e:
                    output = f"Error: {e}"
            else:
                output = f"Unknown tool: {tc.name}"
            
            # Add to messages
            messages.append(Message(
                role="assistant",
                content=response.content or "",
                tool_calls=[{
                    "id": tc.id,
                    "type": "function",
                    "function": {
                        "name": tc.name,
                        "arguments": json.dumps(tc.arguments) if isinstance(tc.arguments, dict) else tc.arguments,
                    }
                }]
            ))
            messages.append(Message(
                role="tool",
                content=output,
                tool_call_id=tc.id,
                name=tc.name,
            ))
        
        elapsed = time.time() - start
        sys.stdout.write(
            f"\r  [{agent_type}] {description} ... {tool_count} tools, {elapsed:.1f}s"
        )
        sys.stdout.flush()
    
    elapsed = time.time() - start
    sys.stdout.write(
        f"\r  [{agent_type}] {description} - done ({tool_count} tools, {elapsed:.1f}s)\n"
    )
    
    # Return final text
    return response.content or "(subagent returned no text)"
