"""
Base Agent

Abstract base class for specialized agents with tool execution.

Integrates:
- Tool execution with agentic loop
- Skills injection for domain knowledge
- Hooks for lifecycle control (pre/post tool use)
- TodoManager for task tracking
- Subagent spawning for decomposition
"""

import json
import re
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any, Optional

from rich.console import Console
from rich.panel import Panel

from ai_coder.llm.interface import LLMProvider, Message, ToolCall
from ai_coder.tools.base import ToolRegistry, ToolExecutor
from ai_coder.core.prompts import get_tool_schemas

console = Console()


class AgentType(Enum):
    """Types of specialized agents."""
    
    CODE_EXPLORER = "code-explorer"
    CODE_ARCHITECT = "code-architect"
    CODE_REVIEWER = "code-reviewer"
    PLANNER = "planner"
    SECURITY_REVIEWER = "security-reviewer"
    TDD_GUIDE = "tdd-guide"
    BUILD_FIXER = "build-fixer"
    REFACTOR = "refactor"
    DEBUGGER = "debugger"
    DATABASE_REVIEWER = "database-reviewer"

    @property
    def display_name(self) -> str:
        """Human-readable name."""
        return self.value.replace("-", " ").title()

    @property
    def color(self) -> str:
        """Color for terminal display."""
        colors = {
            "code-explorer": "yellow",
            "code-architect": "green",
            "code-reviewer": "red",
            "planner": "cyan",
            "security-reviewer": "magenta",
            "tdd-guide": "blue",
            "build-fixer": "bright_yellow",
            "refactor": "bright_green",
            "debugger": "bright_red",
            "database-reviewer": "bright_cyan",
        }
        return colors.get(self.value, "blue")


@dataclass
class AgentResult:
    """Result from an agent execution."""
    
    agent_type: AgentType
    success: bool
    output: str = ""
    key_files: list[str] = field(default_factory=list)
    findings: dict[str, Any] = field(default_factory=dict)
    error: Optional[str] = None

    @classmethod
    def success_result(
        cls,
        agent_type: AgentType,
        output: str,
        key_files: list[str] = None,
        **findings,
    ) -> "AgentResult":
        """Create a successful result."""
        return cls(
            agent_type=agent_type,
            success=True,
            output=output,
            key_files=key_files or [],
            findings=findings,
        )

    @classmethod
    def error_result(cls, agent_type: AgentType, error: str) -> "AgentResult":
        """Create an error result."""
        return cls(agent_type=agent_type, success=False, error=error)


class Agent(ABC):
    """
    Abstract base class for specialized agents.
    
    Features:
    - Agentic loop with tool execution
    - Skills injection for domain knowledge
    - Hooks for lifecycle control
    - TodoManager for task tracking
    - Subagent spawning for task decomposition
    """
    
    MAX_ITERATIONS = 25

    def __init__(
        self,
        llm: LLMProvider,
        tool_registry: ToolRegistry,
        project_root: Path,
    ):
        self.llm = llm
        self.tool_registry = tool_registry
        self.tool_executor = ToolExecutor(tool_registry)
        self.project_root = project_root
        self.messages: list[Message] = []
        
        # Optional subsystems (initialized lazily)
        self._skill_loader = None
        self._hook_runner = None
        self._todo_manager = None

    @property
    @abstractmethod
    def agent_type(self) -> AgentType:
        """The type of this agent."""
        pass

    @property
    @abstractmethod
    def system_prompt(self) -> str:
        """System prompt defining the agent's role and capabilities."""
        pass

    @property
    def allowed_tools(self) -> list[str]:
        """Tools this agent is allowed to use. Override to restrict."""
        return [
            "read_file", "list_files", "search_files",
            "write_file", "patch_file", "delete_file",
            "run_command",
            "web_search", "read_url",
            "todo_write", "use_skill", "spawn_agent",
            "task_complete",
        ]

    @property
    def skill_loader(self):
        """Lazy-init SkillLoader."""
        if self._skill_loader is None:
            from ai_coder.skills.loader import SkillLoader
            skills_dir = Path(__file__).parent.parent / "skills"
            self._skill_loader = SkillLoader(skills_dir)
        return self._skill_loader

    @property
    def hook_runner(self):
        """Lazy-init HookRunner."""
        if self._hook_runner is None:
            from ai_coder.hooks.runner import create_default_hooks
            self._hook_runner = create_default_hooks(self.project_root)
        return self._hook_runner

    @property
    def todo_manager(self):
        """Lazy-init TodoManager."""
        if self._todo_manager is None:
            from ai_coder.workflow.todo import TodoManager
            self._todo_manager = TodoManager()
        return self._todo_manager

    def execute(self, prompt: str) -> AgentResult:
        """
        Execute the agent with the given prompt using an agentic loop.
        
        The loop:
        1. Send prompt + system prompt to LLM
        2. If LLM returns tool calls → execute them, add results, repeat
        3. If LLM returns text only → done, return result
        """
        try:
            console.print(f"[{self.agent_type.color}]>> {self.agent_type.display_name} starting...[/{self.agent_type.color}]")
            
            # Run session start hooks
            self.hook_runner.run_session_hooks("SessionStart")
            
            # Initialize conversation
            self.messages = [
                Message(role="system", content=self._build_system_prompt()),
                Message(role="user", content=prompt),
            ]
            
            # Get tool schemas - FILTERED to only allowed tools
            all_tools = get_tool_schemas()
            tools = [t for t in all_tools if t["function"]["name"] in self.allowed_tools]
            
            # Agentic loop
            final_output = ""
            for iteration in range(self.MAX_ITERATIONS):
                console.print(f"[dim]  Iteration {iteration + 1}...[/dim]")
                
                # Get LLM response
                response = self.llm.complete(self.messages, tools=tools)
                
                # Check for tool calls (native or extracted from JSON text)
                tool_calls = response.tool_calls
                
                if not tool_calls and response.content:
                    extracted = self._extract_json_tool_calls(response.content)
                    if extracted:
                        tool_calls = extracted
                
                # Execute tools if found
                if tool_calls:
                    tool_results = []
                    for tool_call in tool_calls:
                        # Check for task_complete
                        if tool_call.name == "task_complete":
                            final_output = tool_call.arguments.get("summary", response.content or "")
                            break
                        
                        result = self._execute_tool(tool_call)
                        tool_results.append((tool_call, result))
                    else:
                        # Add assistant message with tool calls
                        self.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,
                                }
                            } for tc in tool_calls]
                        ))
                        
                        # Add tool results
                        for tool_call, result in tool_results:
                            self.messages.append(Message(
                                role="tool",
                                content=result,
                                tool_call_id=tool_call.id,
                                name=tool_call.name,
                            ))
                        
                        continue
                    
                    # task_complete was called
                    break
                
                # No tool calls - agent is done
                final_output = response.content or ""
                break
            
            # Run session end hooks
            self.hook_runner.run_session_hooks("SessionEnd")
            
            # Display todo progress if any
            if self.todo_manager.todos:
                self.todo_manager.display()
            
            # Print output
            if final_output:
                console.print(Panel(
                    final_output[:3000] + ("..." if len(final_output) > 3000 else ""),
                    title=f"[bold]{self.agent_type.display_name} Result[/bold]",
                    border_style=self.agent_type.color,
                ))
            
            key_files = self._extract_key_files(final_output)
            if key_files:
                console.print(f"[{self.agent_type.color}]* Found {len(key_files)} key files[/{self.agent_type.color}]")
            
            return AgentResult.success_result(
                agent_type=self.agent_type,
                output=final_output,
                key_files=key_files,
            )
            
        except Exception as e:
            console.print(f"[red]Error: {e}[/red]")
            return AgentResult.error_result(self.agent_type, str(e))

    def _build_system_prompt(self) -> str:
        """Build the full system prompt — only includes sections for tools the agent has."""
        base = self.system_prompt
        allowed = self.allowed_tools
        
        # Skills section — only if agent has use_skill
        skills_section = ""
        if "use_skill" in allowed:
            available_skills = self.skill_loader.get_descriptions()
            if available_skills and available_skills != "(no skills available)":
                skills_section = f"\n## Available Skills (use the use_skill tool):\n{available_skills}"
        
        # Subagent section — only if agent has spawn_agent
        subagent_section = ""
        if "spawn_agent" in allowed:
            from ai_coder.subagents.spawner import get_agent_descriptions
            subagent_section = f"\n## Subagent Types (for spawn_agent tool):\n{get_agent_descriptions()}"
        
        # Build tool examples — ONLY for tools this agent can use
        tool_examples = {
            "read_file": '- **read_file**: `{{"name": "read_file", "arguments": {{"path": "path/to/file"}}}}`',
            "write_file": '- **write_file**: `{{"name": "write_file", "arguments": {{"path": "file.py", "content": "..."}}}}`',
            "list_files": '- **list_files**: `{{"name": "list_files", "arguments": {{"path": "."}}}}`',
            "search_files": '- **search_files**: `{{"name": "search_files", "arguments": {{"pattern": "search_term"}}}}`',
            "run_command": '- **run_command**: `{{"name": "run_command", "arguments": {{"command": "ls"}}}}`',
            "patch_file": '- **patch_file**: `{{"name": "patch_file", "arguments": {{"path": "file", "patches": [{{"search": "old", "replace": "new"}}]}}}}`',
            "use_skill": '- **use_skill**: `{{"name": "use_skill", "arguments": {{"skill": "python"}}}}`',
            "todo_write": '- **todo_write**: `{{"name": "todo_write", "arguments": {{"action": "add", "title": "Do something"}}}}`',
            "spawn_agent": '- **spawn_agent**: `{{"name": "spawn_agent", "arguments": {{"type": "explore", "task": "Describe what to do"}}}}`',
            "task_complete": '- **task_complete**: `{{"name": "task_complete", "arguments": {{"summary": "What was done"}}}}`',
        }
        
        tool_lines = []
        for tool_name in allowed:
            if tool_name in tool_examples:
                tool_lines.append(tool_examples[tool_name])
        tools_text = "\n".join(tool_lines)
        
        context = f"""

## Project Context
- Working directory: {self.project_root}
{skills_section}
{subagent_section}

## How to Use Tools

You MUST use tools to complete your task. To call a tool, output ONLY this JSON:

```json
{{"name": "tool_name", "arguments": {{"key": "value"}}}}
```

### Your Available Tools:
{tools_text}

## Rules:
1. Start by using list_files to see the project structure
2. Read files before modifying them
3. Output EXACTLY ONE tool call per response — nothing else
4. When done, provide your final answer as plain text (no tool call)
"""
        return base + context

    def _extract_json_tool_calls(self, content: str) -> list[ToolCall]:
        """Extract tool calls from JSON in text content."""
        tool_calls = []
        
        patterns = [
            r'```json\s*(\{.*?\})\s*```',
            r'```\s*(\{.*?\})\s*```',
            r'(\{"name":\s*"[^"]+",\s*"arguments":\s*\{.*?\}\})',
        ]
        
        for pattern in patterns:
            matches = re.findall(pattern, 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 self.allowed_tools:
                        if isinstance(args, str):
                            args = json.loads(args)
                        tool_calls.append(ToolCall(
                            id=f"extracted_{name}_{len(tool_calls)}",
                            name=name,
                            arguments=args if isinstance(args, dict) else {},
                        ))
                except (json.JSONDecodeError, KeyError):
                    continue
        
        return tool_calls

    def _execute_tool(self, tool_call: ToolCall) -> str:
        """Execute a tool with hooks and return the result string."""
        
        # Run PreToolUse hooks
        hook_result = self.hook_runner.run_pre_tool_use(tool_call.name, tool_call.arguments)
        if not hook_result.allowed:
            return f"BLOCKED by hook: {hook_result.message}"
        
        # Handle meta tools
        if tool_call.name == "use_skill":
            return self._handle_use_skill(tool_call.arguments)
        
        if tool_call.name == "todo_write":
            return self._handle_todo_write(tool_call.arguments)
        
        if tool_call.name == "spawn_agent":
            return self._handle_spawn_agent(tool_call.arguments)
        
        # Execute regular tool
        tool = self.tool_registry.get(tool_call.name)
        if not tool:
            return f"Error: Unknown tool '{tool_call.name}'"
        
        try:
            console.print(f"[dim]    → {tool_call.name}({', '.join(f'{k}={repr(v)[:50]}' for k, v in tool_call.arguments.items())})[/dim]")
            result = tool.execute(**tool_call.arguments)
            output = result.output[:4000] if result.output else ""
            if result.error:
                output = f"Error: {result.error}"
        except Exception as e:
            output = f"Error executing {tool_call.name}: {str(e)}"
        
        # Run PostToolUse hooks
        self.hook_runner.run_post_tool_use(tool_call.name, tool_call.arguments, output)
        
        return output

    def _handle_use_skill(self, arguments: dict) -> str:
        """Handle skill loading."""
        skill_name = arguments.get("skill", "")
        if not skill_name:
            return f"Available skills: {', '.join(self.skill_loader.list_skills())}"
        
        content = self.skill_loader.get_skill_content(skill_name)
        if content:
            console.print(f"[dim]    → Loaded skill: {skill_name}[/dim]")
            return content
        
        return f"Skill '{skill_name}' not found. Available: {', '.join(self.skill_loader.list_skills())}"

    def _handle_todo_write(self, arguments: dict) -> str:
        """Handle todo operations."""
        action = arguments.get("action", "list")
        
        if action == "add":
            title = arguments.get("title", "Untitled task")
            priority = arguments.get("priority", 0)
            item_id = self.todo_manager.add(title, priority)
            return f"Added: {item_id} - {title}"
        
        elif action == "update":
            item_id = arguments.get("id", "")
            status = arguments.get("status", "done")
            if self.todo_manager.update_status(item_id, status):
                return f"Updated {item_id} to {status}"
            return f"Todo item '{item_id}' not found"
        
        elif action == "remove":
            item_id = arguments.get("id", "")
            if self.todo_manager.remove(item_id):
                return f"Removed {item_id}"
            return f"Todo item '{item_id}' not found"
        
        elif action == "list":
            return self.todo_manager.get_summary()
        
        return f"Unknown action: {action}. Use: add, update, remove, list"

    def _handle_spawn_agent(self, arguments: dict) -> str:
        """Handle spawning a subagent."""
        agent_type = arguments.get("type", "explore")
        task = arguments.get("task", "")
        
        if not task:
            return "Error: 'task' argument required for spawn_agent"
        
        console.print(f"[dim]    → Spawning {agent_type} subagent...[/dim]")
        
        from ai_coder.subagents.spawner import run_subagent
        result = run_subagent(
            llm=self.llm,
            tool_registry=self.tool_registry,
            project_root=self.project_root,
            description=task[:80],
            prompt=task,
            agent_type=agent_type,
        )
        
        return f"[Subagent {agent_type} result]:\n{result[:4000]}"

    def _extract_key_files(self, output: str) -> list[str]:
        """Extract file paths mentioned in the output."""
        patterns = [
            r'`([^`]+\.[a-zA-Z]+)`',
            r'[\'"]([\\w/\\\\.-]+\.[a-zA-Z]+)[\'"]',
            r'(?:^|\s)([\w/\\\\.-]+\.[a-zA-Z]{1,4})(?:\s|$|:)',
        ]
        
        files = set()
        for pattern in patterns:
            matches = re.findall(pattern, output)
            for match in matches:
                if not match.startswith(('http://', 'https://')):
                    if '.' in match and len(match) < 200:
                        files.add(match)
        
        return sorted(files)
