"""
Blackreach Planner - Break complex goals into subtasks.

For simple goals like "search for cats":
  → Just run the ReAct loop directly

For complex goals like "download 5 papers about transformers from arxiv":
  → Break into subtasks:
    1. Go to arxiv.org
    2. Search for "transformers"
    3. Download paper 1
    4. Download paper 2
    ... etc.

The planner uses a smarter model (or more tokens) to reason about the task,
then the executor uses a faster model for each step.
"""

import json
import re
from typing import List, Optional
from dataclasses import dataclass

from blackreach.llm import LLM, LLMConfig


# P0-PERF: Pre-compiled regex patterns
_RE_NUMBERS_GT_1 = re.compile(r'\b([2-9]|[1-9]\d+)\b')
_RE_JSON_OBJECT = re.compile(r'\{[\s\S]*\}')


@dataclass
class Subtask:
    """A single subtask generated by the planner."""
    description: str
    expected_outcome: str
    optional: bool = False


@dataclass
class Plan:
    """A plan with subtasks."""
    goal: str
    subtasks: List[Subtask]
    estimated_steps: int


class Planner:
    """
    Plans complex goals by breaking them into subtasks.

    Uses the LLM to analyze the goal and generate a sequence
    of concrete, actionable subtasks.
    """

    PLAN_PROMPT = """You are a task planner for an autonomous web browser agent.

GOAL: {goal}

Break this goal into a sequence of concrete, actionable subtasks.
Each subtask should be simple enough to complete in 1-5 browser actions.

RULES:
- Be specific about URLs to visit
- Be specific about what to search for
- Be specific about what to click/download
- Each subtask should have a clear success condition
- Keep the list focused (3-8 subtasks usually)

OUTPUT FORMAT (JSON):
{{
    "subtasks": [
        {{"description": "Go to arxiv.org", "expected_outcome": "On arxiv homepage"}},
        {{"description": "Search for transformers", "expected_outcome": "Search results visible"}},
        ...
    ],
    "estimated_total_steps": 15
}}

JSON:"""

    def __init__(self, llm_config: Optional[LLMConfig] = None):
        self.llm = LLM(llm_config)

    def is_simple_goal(self, goal: str) -> bool:
        """
        Check if goal is simple enough to run directly without planning.

        Simple goals:
        - Single action like "go to wikipedia"
        - Single search like "search for cats"
        - Short, focused tasks

        Complex goals (need planning):
        - Multiple items to download
        - Multi-site navigation
        - Conditional logic ("if X then Y")
        """
        # Simple keyword checks
        simple_indicators = [
            "go to",
            "navigate to",
            "search for",
            "find",
            "click",
            "open",
        ]

        # Complex indicators
        complex_indicators = [
            "download",  # downloads usually need planning
            "all",       # "download all" needs iteration
            "each",      # "for each" needs iteration
            "multiple",
            "several",
            "papers",    # academic papers are complex
            "files",
            "images",
            "then",      # sequential steps
            "and then",
            "after that",
        ]

        goal_lower = goal.lower()

        # Check for complex indicators
        for indicator in complex_indicators:
            if indicator in goal_lower:
                return False

        # If it has number > 1, it's complex (uses pre-compiled pattern)
        numbers = _RE_NUMBERS_GT_1.findall(goal)
        if numbers:
            return False

        # Check if it matches simple patterns
        for indicator in simple_indicators:
            if goal_lower.startswith(indicator):
                return True

        # Default to simple for short goals
        return len(goal.split()) < 10

    def plan(self, goal: str) -> Optional[Plan]:
        """
        Generate a plan for the goal.

        Returns None if the goal is simple enough to run directly.
        """
        if self.is_simple_goal(goal):
            return None

        prompt = self.PLAN_PROMPT.format(goal=goal)

        response = self.llm.generate(
            "You are a task planning assistant. Output ONLY valid JSON.",
            prompt
        )

        # Parse the response (uses pre-compiled regex)

        # Find JSON in response
        json_match = _RE_JSON_OBJECT.search(response)
        if not json_match:
            # Couldn't parse plan, return a default
            return Plan(
                goal=goal,
                subtasks=[Subtask(
                    description=f"Complete: {goal}",
                    expected_outcome="Goal achieved"
                )],
                estimated_steps=20
            )

        try:
            data = json.loads(json_match.group())

            subtasks = []
            for st in data.get("subtasks", []):
                subtasks.append(Subtask(
                    description=st.get("description", ""),
                    expected_outcome=st.get("expected_outcome", ""),
                    optional=st.get("optional", False)
                ))

            return Plan(
                goal=goal,
                subtasks=subtasks,
                estimated_steps=data.get("estimated_total_steps", len(subtasks) * 5)
            )
        except json.JSONDecodeError:
            # Fallback
            return Plan(
                goal=goal,
                subtasks=[Subtask(
                    description=f"Complete: {goal}",
                    expected_outcome="Goal achieved"
                )],
                estimated_steps=20
            )

    def format_plan(self, plan: Plan) -> str:
        """Format plan for display."""
        lines = [f"Plan for: {plan.goal}", ""]

        for i, st in enumerate(plan.subtasks, 1):
            optional = " (optional)" if st.optional else ""
            lines.append(f"  {i}. {st.description}{optional}")
            lines.append(f"     → {st.expected_outcome}")

        lines.append("")
        lines.append(f"Estimated steps: ~{plan.estimated_steps}")

        return "\n".join(lines)


# Convenience function
def maybe_plan(goal: str, llm_config: Optional[LLMConfig] = None) -> Optional[Plan]:
    """Plan goal if complex, return None if simple."""
    planner = Planner(llm_config)
    return planner.plan(goal)
