"""
Agent Orchestrator Module

Main agent execution loop:
- Task intake and initialization
- Context loading
- Planning invocation
- Tool execution
- Error handling and iteration
- Session management
"""

import json
import uuid
from pathlib import Path
from typing import Optional

from rich.console import Console
from rich.panel import Panel
from rich.progress import Progress, SpinnerColumn, TextColumn

from ai_coder.agent.planner import Planner, ExecutionPlan
from ai_coder.core.context import ContextLoader, ProjectContext
from ai_coder.core.prompts import PromptBuilder, get_tool_schemas, get_anthropic_tools
from ai_coder.llm.interface import LLMProvider, Message, CompletionResponse
from ai_coder.memory.store import MemoryStore, TaskStatus
from ai_coder.safety.guards import get_safety_guard, configure_safety, SafetyConfig
from ai_coder.tools.base import ToolRegistry, ToolExecutor, create_tool_registry
from ai_coder.hooks.system import HookSystem

console = Console()


class AgentConfig:
    """Configuration for the agent."""

    def __init__(
        self,
        max_iterations: int = 10,
        auto_confirm: bool = False,
        verbose: bool = False,
        provider: str = "openai",
        model: str = "gpt-4o",
    ):
        self.max_iterations = max_iterations
        self.auto_confirm = auto_confirm
        self.verbose = verbose
        self.provider = provider
        self.model = model


class AgentOrchestrator:
    """
    Main orchestrator for the AI coding agent.
    
    Coordinates:
    - Task initialization
    - Context loading
    - Planning
    - Tool execution
    - Error handling
    - Session management
    """

    def __init__(
        self,
        llm: LLMProvider,
        project_root: Path,
        config: Optional[AgentConfig] = None,
    ):
        self.llm = llm
        self.project_root = project_root.resolve()
        self.config = config or AgentConfig()

        # Initialize components
        self.context_loader = ContextLoader()
        self.context: Optional[ProjectContext] = None

        self.tool_registry = create_tool_registry(self.project_root)
        self.tool_executor = ToolExecutor(self.tool_registry)
        
        # Initialize hooks
        self.hook_system = HookSystem(self.project_root)
        # Load default hooks or user hooks? For now default config can be loaded here
        # self.hook_system.load_hooks(load_config().get("hooks", {}))
        # But load_config isn't imported. Let's lazily load or just init.
        # Actually, let's look for hooks.yaml or config.yaml
        try:
             import yaml
             config_path = self.project_root / ".ai-coder" / "hooks.yaml"
             if config_path.exists():
                 with open(config_path) as f:
                     self.hook_system.load_hooks(yaml.safe_load(f) or {})
             # Also load from main config
             main_config_path = self.project_root / "config.yaml"
             if main_config_path.exists():
                 with open(main_config_path) as f:
                     cfg = yaml.safe_load(f) or {}
                     if "hooks" in cfg:
                         self.hook_system.load_hooks(cfg["hooks"])
        except Exception as e:
            console.print(f"[dim]Failed to load hooks: {e}[/dim]")

        self.memory = MemoryStore(
            persist_path=self.project_root / ".ai-coder" / "session.json"
        )

        self.prompt_builder: Optional[PromptBuilder] = None
        self.planner: Optional[Planner] = None

        # Configure safety
        safety_config = SafetyConfig(
            confirm_delete=not self.config.auto_confirm,
            restrict_to_project=True,
            project_root=self.project_root,
        )
        configure_safety(safety_config)
        get_safety_guard().set_project_root(self.project_root)

    def run(self, task: str) -> bool:
        """
        Execute a task from start to finish.
        
        Args:
            task: Natural language task description
            
        Returns:
            True if task completed successfully, False otherwise
        """
        task_id = str(uuid.uuid4())[:8]

        console.print(Panel(
            f"[bold]{task}[/bold]",
            title="🚀 Starting Task",
            border_style="green",
        ))

        # Initialize session
        self.memory.start_task(task_id, task)

        try:
            # Load project context
            self._load_context()

            # Initialize prompt builder and planner
            self.prompt_builder = PromptBuilder(self.context)
            self.planner = Planner(self.llm, self.context)

            # Main execution loop
            success = self._execution_loop(task)

            if success:
                self.memory.complete_task("Task completed successfully")
                console.print("\n[bold green]✅ Task completed successfully![/bold green]")
            else:
                self.memory.fail_task("Task failed after max iterations")
                console.print("\n[bold red]❌ Task failed[/bold red]")

            return success

        except KeyboardInterrupt:
            console.print("\n[yellow]Task cancelled by user[/yellow]")
            self.memory.update_task_status(TaskStatus.CANCELLED)
            return False

        except Exception as e:
            console.print(f"\n[bold red]Fatal error: {e}[/bold red]")
            self.memory.fail_task(str(e))
            if self.config.verbose:
                console.print_exception()
            return False

    def _load_context(self) -> None:
        """Load project context."""
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console,
        ) as progress:
            progress.add_task("Loading project context...", total=None)
            self.context = self.context_loader.load(self.project_root)

        console.print(f"[dim]Loaded {self.context.total_files} files from {self.project_root.name}[/dim]")

    def _execution_loop(self, task: str) -> bool:
        """Main execution loop."""
        iteration = 0

        # Add initial task message
        self.prompt_builder.add_message("user", self.prompt_builder.build_task_prompt(task))

        while iteration < self.config.max_iterations:
            iteration += 1
            self.memory.increment_iteration()

            if self.config.verbose or iteration > 5:
                 console.print(f"\n[bold cyan]─── Iteration {iteration}/{self.config.max_iterations} ───[/bold cyan]")

            # Get LLM response
            response = self._get_llm_response()

            if response is None:
                continue

            # Handle response
            if response.has_tool_calls:
                # Execute tool calls
                all_success = True
                for tool_call in response.tool_calls:
                    success = self._execute_tool(tool_call.id, tool_call.name, tool_call.arguments)

                    if tool_call.name == "task_complete":
                        return True

                    if tool_call.name == "request_help":
                        # Wait for user input and continue
                        continue

                    if not success:
                        all_success = False

                # Add continuation prompt
                if all_success:
                    self.prompt_builder.add_message(
                        "user",
                        "Continue with the next step, or call task_complete if finished."
                    )
                else:
                    self.prompt_builder.add_message(
                        "user",
                        "The previous step encountered an issue. Please diagnose and fix it."
                    )

            elif response.content:
                # Text response without tool calls
                console.print(f"\n[dim]Agent: {response.content[:500]}[/dim]")
                self.memory.record_llm_response(response.content)

                # Prompt for tool use
                self.prompt_builder.add_message(
                    "user",
                    "Please use the available tools to make progress on the task."
                )

        console.print(f"\n[yellow]Max iterations ({self.config.max_iterations}) reached[/yellow]")
        return False

    def _get_llm_response(self) -> Optional[CompletionResponse]:
        """Get a response from the LLM."""
        try:
            messages = self.prompt_builder.get_messages()

            # Convert to Message objects
            msg_objects = [
                Message(
                    role=m["role"],
                    content=m["content"],
                    tool_call_id=m.get("tool_call_id"),
                    tool_calls=m.get("tool_calls"),
                    name=m.get("name"),
                )
                for m in messages
            ]

            # Get tool schemas for the provider
            if self.config.provider == "anthropic":
                tools = get_anthropic_tools()
            else:
                tools = get_tool_schemas()

            if self.config.verbose:
                console.print(f"[dim]Sending {len(msg_objects)} messages to LLM...[/dim]")

            response = self.llm.complete(msg_objects, tools=tools)

            # Add assistant message to history
            if response.content:
                self.prompt_builder.add_message("assistant", response.content)

            if response.has_tool_calls:
                # Record tool calls in message format (OpenAI-compatible format)
                tool_calls_data = [
                    {
                        "id": tc.id,
                        "type": "function",  # Required by HuggingFace and other OpenAI-compatible APIs
                        "function": {"name": tc.name, "arguments": json.dumps(tc.arguments)},
                    }
                    for tc in response.tool_calls
                ]
                self.prompt_builder.conversation_history[-1]["tool_calls"] = tool_calls_data

            return response

        except Exception as e:
            console.print(f"[red]LLM error: {e}[/red]")
            self.memory.record_error(str(e))
            return None

    def _execute_tool(self, tool_id: str, tool_name: str, arguments: dict) -> bool:
        """Execute a tool and record the result."""
        console.print(f"\n[bold yellow]>> {tool_name}[/bold yellow]")

        if self.config.verbose:
            console.print(f"[dim]Arguments: {json.dumps(arguments, indent=2)}[/dim]")

        # Handle special tools
        if tool_name == "task_complete":
            summary = arguments.get("summary", "Task completed")
            console.print(f"\n[green]>> {summary}[/green]")
            return True

        if tool_name == "request_help":
            question = arguments.get("question", "Need clarification")
            console.print(f"\n[bold cyan]? Agent needs help:[/bold cyan]")
            console.print(f"   {question}")

            # Get user input
            try:
                user_response = console.input("\n[bold]Your response: [/bold]")
                self.prompt_builder.add_message("tool", user_response)
                self.prompt_builder.conversation_history[-1]["tool_call_id"] = tool_id
                self.prompt_builder.conversation_history[-1]["name"] = tool_name
                return True
            except (KeyboardInterrupt, EOFError):
                return False

        # Run PreToolUse hooks
        if not self.hook_system.run_pre_tool_hooks(tool_name, arguments):
             return False # Blocked
             
        # Execute regular tool
        result = self.tool_executor.execute(tool_name, arguments)
        
        # Run PostToolUse hooks
        hook_feedback = self.hook_system.run_post_tool_hooks(tool_name, result.output or "")
        if hook_feedback:
             if result.output:
                 result.output += f"\n\n[Hook Feedback]:\n{hook_feedback}"
             else:
                 result.output = f"[Hook Feedback]:\n{hook_feedback}"
             # If hook had output, user/agent should see it.
             console.print(f"[dim]Hook Feedback: {len(hook_feedback)} chars[/dim]")

        # Record in memory
        self.memory.record_tool_call(
            tool_name,
            arguments,
            result.output or result.error or "",
            is_error=result.is_error,
        )

        # Add tool result to conversation
        result_content = result.to_message()
        self.prompt_builder.add_message("tool", result_content)
        self.prompt_builder.conversation_history[-1]["tool_call_id"] = tool_id
        self.prompt_builder.conversation_history[-1]["name"] = tool_name

        if result.is_error:
            console.print(f"[red]Error: {result.error}[/red]")
        elif result.output:
            # Truncate long output
            output = result.output
            if len(output) > 500:
                output = output[:500] + f"\n... ({len(result.output) - 500} more chars)"
            console.print(f"[dim]{output}[/dim]")

        return result.is_success

    def get_session_summary(self) -> str:
        """Get a summary of the current session."""
        return self.memory.get_session_summary()
