"""
Planner Module

Task decomposition and planning:
- Natural language task parsing
- Step-by-step action planning
- Dynamic re-planning on failure
"""

from dataclasses import dataclass, field
from typing import Optional

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

from ai_coder.llm.interface import LLMProvider, Message, CompletionResponse
from ai_coder.core.context import ProjectContext

console = Console()


PLANNING_PROMPT = '''You are a planning assistant for a software engineering agent.

Given a task description and project context, create a detailed plan of steps to accomplish the task.

## Guidelines

1. Break the task into small, specific, actionable steps
2. Each step should map to a single tool action (read file, write file, run command, etc.)
3. Order steps logically - understand before modifying, test after changes
4. Include verification steps (running tests, checking builds)
5. Consider error handling and rollback

## Available Actions

- read_file: Read file contents
- write_file: Create or overwrite a file
- patch_file: Apply targeted edits to a file
- delete_file: Delete a file
- run_command: Execute a shell command
- search_files: Search for patterns in code
- list_files: List directory contents

## Response Format

Provide a numbered list of steps. Each step should be:
- Clear and specific
- Include the target file or command
- Explain the purpose

Example:
1. Read package.json to understand project dependencies
2. Search for existing authentication code patterns
3. Create src/auth/oauth.js with OAuth client implementation
4. Update src/index.js to import and initialize OAuth
5. Run npm test to verify changes
6. Run npm run lint to check code style
'''


@dataclass
class PlanStep:
    """A single step in the execution plan."""

    number: int
    description: str
    action_type: Optional[str] = None
    target: Optional[str] = None
    completed: bool = False
    result: Optional[str] = None


@dataclass
class ExecutionPlan:
    """Complete execution plan for a task."""

    task: str
    steps: list[PlanStep] = field(default_factory=list)
    current_step: int = 0

    def get_current_step(self) -> Optional[PlanStep]:
        """Get the current step to execute."""
        if self.current_step < len(self.steps):
            return self.steps[self.current_step]
        return None

    def advance(self) -> bool:
        """Move to the next step. Returns True if more steps remain."""
        self.current_step += 1
        return self.current_step < len(self.steps)

    def mark_complete(self, result: str = "") -> None:
        """Mark the current step as complete."""
        if self.current_step < len(self.steps):
            self.steps[self.current_step].completed = True
            self.steps[self.current_step].result = result

    def is_complete(self) -> bool:
        """Check if all steps are complete."""
        return self.current_step >= len(self.steps)

    def get_progress(self) -> str:
        """Get progress string."""
        completed = sum(1 for s in self.steps if s.completed)
        return f"Step {self.current_step + 1}/{len(self.steps)} ({completed} completed)"

    def to_text(self) -> str:
        """Convert plan to text format."""
        lines = [f"Task: {self.task}", "", "Steps:"]
        for step in self.steps:
            status = "✅" if step.completed else "⬜"
            lines.append(f"{status} {step.number}. {step.description}")
        return "\n".join(lines)


class Planner:
    """
    Task planner that creates execution plans using LLM.
    
    Features:
    - Task decomposition
    - Context-aware planning
    - Step parsing
    - Re-planning support
    """

    def __init__(self, llm: LLMProvider, context: Optional[ProjectContext] = None):
        self.llm = llm
        self.context = context

    def create_plan(self, task: str) -> ExecutionPlan:
        """
        Create an execution plan for a task.
        
        Args:
            task: Natural language task description
            
        Returns:
            ExecutionPlan with ordered steps
        """
        # Load prompt from definition if available
        planner_def = Path(__file__).parent.parent / "agents" / "definitions" / "planner.md"
        system_prompt = PLANNING_PROMPT
        
        if planner_def.exists():
            try:
                content = planner_def.read_text(encoding="utf-8")
                # Strip frontmatter
                if content.startswith("---"):
                    parts = content.split("---", 2)
                    if len(parts) >= 3:
                        system_prompt = parts[2].strip()
            except Exception as e:
                console.print(f"[dim]Failed to load planner definition: {e}[/dim]")

        # Build planning prompt
        messages = [
            Message(role="system", content=system_prompt),
        ]

        # Add context if available
        user_content = f"## Task\n\n{task}"

        if self.context:
            user_content += f"\n\n## Project Context\n\n{self.context.get_summary()}"

        messages.append(Message(role="user", content=user_content))

        # Get plan from LLM
        console.print("\n[bold blue]📋 Creating execution plan...[/bold blue]")

        response = self.llm.complete(messages)

        if not response.content:
            # Fallback to simple plan
            return ExecutionPlan(
                task=task,
                steps=[
                    PlanStep(number=1, description=f"Analyze the task: {task}"),
                    PlanStep(number=2, description="Implement the required changes"),
                    PlanStep(number=3, description="Verify the changes work correctly"),
                ]
            )

        # Parse the response into steps
        steps = self._parse_plan(response.content)

        plan = ExecutionPlan(task=task, steps=steps)

        # Display the plan
        self._display_plan(plan)

        return plan

    def replan(self, original_plan: ExecutionPlan, error: str) -> ExecutionPlan:
        """
        Create a new plan after a failure.
        
        Args:
            original_plan: The plan that failed
            error: The error that occurred
            
        Returns:
            New ExecutionPlan with adjusted steps
        """
        messages = [
            Message(role="system", content=PLANNING_PROMPT),
        ]

        user_content = f"""## Original Task
{original_plan.task}

## Previous Plan
{original_plan.to_text()}

## Error Encountered
{error}

## Current Progress
{original_plan.get_progress()}

Please create a revised plan that addresses the error and completes the remaining work.
"""

        messages.append(Message(role="user", content=user_content))

        console.print("\n[bold yellow]🔄 Re-planning after error...[/bold yellow]")

        response = self.llm.complete(messages)

        if not response.content:
            # Keep remaining steps
            remaining = [s for s in original_plan.steps if not s.completed]
            return ExecutionPlan(task=original_plan.task, steps=remaining)

        steps = self._parse_plan(response.content)
        plan = ExecutionPlan(task=original_plan.task, steps=steps)

        self._display_plan(plan)

        return plan

    def _parse_plan(self, content: str) -> list[PlanStep]:
        """Parse LLM response into plan steps."""
        steps = []
        lines = content.strip().split("\n")

        for line in lines:
            line = line.strip()

            # Match numbered list items
            if not line:
                continue

            # Try to parse "1. description" or "1) description" format
            import re
            match = re.match(r"^(\d+)[.)]\s*(.+)$", line)

            if match:
                number = int(match.group(1))
                description = match.group(2).strip()

                # Try to identify action type from description
                action_type = self._infer_action_type(description)
                target = self._extract_target(description)

                steps.append(PlanStep(
                    number=number,
                    description=description,
                    action_type=action_type,
                    target=target,
                ))

        # If no steps parsed, create a default
        if not steps:
            steps = [
                PlanStep(number=1, description="Execute the requested task"),
            ]

        return steps

    def _infer_action_type(self, description: str) -> Optional[str]:
        """Infer the action type from step description."""
        lower = description.lower()

        if any(kw in lower for kw in ["read", "examine", "check", "look at", "review"]):
            return "read_file"
        elif any(kw in lower for kw in ["create", "write", "add new"]):
            return "write_file"
        elif any(kw in lower for kw in ["update", "modify", "edit", "change", "patch"]):
            return "patch_file"
        elif any(kw in lower for kw in ["delete", "remove"]):
            return "delete_file"
        elif any(kw in lower for kw in ["run", "execute", "test", "build", "install"]):
            return "run_command"
        elif any(kw in lower for kw in ["search", "find", "grep"]):
            return "search_files"
        elif any(kw in lower for kw in ["list", "ls", "directory"]):
            return "list_files"

        return None

    def _extract_target(self, description: str) -> Optional[str]:
        """Extract target file/command from description."""
        import re

        # Look for file paths
        file_match = re.search(r"[`'\"]?([a-zA-Z0-9_\-./]+\.[a-zA-Z0-9]+)[`'\"]?", description)
        if file_match:
            return file_match.group(1)

        # Look for commands
        cmd_match = re.search(r"`([^`]+)`", description)
        if cmd_match:
            return cmd_match.group(1)

        return None

    def _display_plan(self, plan: ExecutionPlan) -> None:
        """Display the plan in the console."""
        lines = []
        for step in plan.steps:
            icon = "📌" if step.action_type else "📝"
            lines.append(f"{icon} {step.number}. {step.description}")

        panel = Panel(
            "\n".join(lines),
            title=f"[bold]Execution Plan[/bold] ({len(plan.steps)} steps)",
            border_style="blue",
        )
        console.print(panel)
