"""Planner agent — explores the task and produces a structured milestone plan."""

from __future__ import annotations

from typing import Any

from .config import config
from .knowledge import KnowledgeBank
from .llm_client import LLMClient
from .prompts import PLANNER_SYSTEM_PROMPT
from .task_logger import get_agent_logger, log_agent_api_call
from .tools import (
    SUBMIT_PLAN,
    web_extract,
    web_search,
)
from .types import Milestone, Plan
from .utils import ToolResult

_task_log = get_agent_logger("planner")

PLANNER_TOOLS = [
    SUBMIT_PLAN,
]


def _handle_web_search(planner: Planner, inp: dict) -> tuple[Plan | None, ToolResult]:
    result = web_search(inp["query"], inp.get("max_results", 10))
    return None, result


def _handle_web_fetch(planner: Planner, inp: dict) -> tuple[Plan | None, ToolResult]:
    result = web_extract(inp["urls"], max_chars=inp.get("max_chars", 8000))
    return None, result


def _handle_submit_plan(planner: Planner, inp: dict) -> tuple[Plan | None, ToolResult]:
    milestones = [Milestone(description=desc) for desc in inp.get("milestones", [])]
    infeasible_reason = (inp.get("infeasible_reason") or "").strip() or None
    if len(milestones) > 25:
        return None, ToolResult(
            output=f"Plan rejected. Too many milestones: {len(milestones)} milestones found. Please check your formatted structure correctly."
        )
    plan = Plan(
        task_instruction=planner.instruction,
        milestones=milestones,
        revision_count=planner.revision_count,
        infeasible_reason=infeasible_reason,
    )
    return plan, ToolResult(output="Plan accepted.")


_TOOL_DISPATCH: dict[str, Any] = {
    # "web_search": _handle_web_search,
    # "web_fetch": _handle_web_fetch,
    "submit_plan": _handle_submit_plan,
}


class Planner:
    """Planner: breaks the task into structured milestones via ``submit_plan``."""

    def __init__(
        self,
        env: Any,
        knowledge: KnowledgeBank,
        *args: Any,
        **kwargs: Any,
    ) -> None:
        self.env = env
        self.llm = LLMClient(
            name="planner",
            system_prompt=PLANNER_SYSTEM_PROMPT,
            tools=[t.schema for t in PLANNER_TOOLS],
            model=config.planner_model,
            provider=config.provider,
            only_n_most_recent_images=config.planner_n_recent_images,
            effort=config.planner_effort,
        )
        self.knowledge = knowledge
        self.instruction = ""
        self.revision_count = 0

        global _task_log
        _task_log = get_agent_logger("planner")

    def reset(self) -> None:
        self.llm.reset()

    def log_usage(self) -> None:
        _task_log.info(self.llm.print_usage())

    def run(
        self,
        instruction: str,
        screenshot_b64: str,
        *,
        replan: bool = False,
        current_plan: Plan | None = None,
        verifier_diagnosis: str = "",
        verifier_suggestion: str = "",
    ) -> Plan:
        """Run the planner (or replanner) and return a ``Plan``."""
        revision_count = (current_plan.revision_count + 1) if replan and current_plan else 0
        self.instruction = instruction
        self.revision_count = revision_count

        if replan:
            prompt = self._build_replan_prompt(instruction, current_plan, verifier_diagnosis, verifier_suggestion)
        else:
            prompt = self._build_plan_prompt(instruction)

        images = None
        if screenshot_b64:
            images = [screenshot_b64]
            prompt += "\nThe current screenshot is attached."
        self.llm.add_user_message(text=prompt, images=images)

        budget = config.planner_max_steps
        label = "REPLANNER" if replan else "PLANNER"

        _task_log.info(
            "\n=== %s START | revision=%d ===\nTask: %s",
            label,
            revision_count,
            instruction,
        )

        for _ in range(budget):
            plan = self._execute_step()
            if plan is not None:
                _task_log.info(
                    "\n=== %s END | milestones=%d ===\n%s",
                    label,
                    len(plan.milestones),
                    plan.format(),
                )
                return plan

        _task_log.info("\n=== %s END | budget exhausted ===", label)

        if replan and current_plan:
            return current_plan

        fallback = Plan(
            task_instruction=instruction,
            milestones=[Milestone(description=f"Complete the task: {instruction}")],
        )
        return fallback

    def _execute_step(self) -> Plan | None:
        """Dispatch a single LLM turn's tool calls. Returns Plan if submit_plan fired."""
        tool_calls, text_blocks, thinking_blocks, _ = self.llm.generate_and_process()

        if not tool_calls:
            return None

        tool_call_results: list[tuple[str, ToolResult]] = []
        final_plan: Plan | None = None

        for i, tc in enumerate(tool_calls):
            handler = _TOOL_DISPATCH.get(tc["name"])
            if handler is None:
                tool_call_results.append((tc["id"], ToolResult(error=f"Unknown tool: {tc['name']}")))
                continue

            plan, result = handler(self, tc["input"])
            tool_call_results.append((tc["id"], result))

            if plan is not None:
                final_plan = plan
                for remaining in tool_calls[i + 1 :]:
                    tool_call_results.append(
                        (
                            remaining["id"],
                            ToolResult(error="Not executed: a prior tool returned control."),
                        )
                    )
                break

        log_agent_api_call(
            agent_name="planner",
            tool_calls=tool_calls,
            text_blocks=text_blocks,
            thinking_blocks=thinking_blocks,
            tool_call_results=tool_call_results,
            agent_action=final_plan,
        )

        self.llm.add_tool_results(tool_call_results)
        return final_plan

    def _build_plan_prompt(self, instruction: str) -> str:
        kb_text = self.knowledge.format_for_prompt()
        prompt = f"Task: {instruction}"
        if kb_text:
            prompt += f"\n\n{kb_text}"
        return prompt

    def _build_replan_prompt(
        self,
        instruction: str,
        plan: Plan,
        verifier_diagnosis: str,
        verifier_suggestion: str,
    ) -> str:
        kb_text = self.knowledge.format_for_prompt()

        parts = [
            f"TASK: {instruction}",
            (f"\nPREVIOUS PLAN (revision {plan.revision_count}):\n" + f"\n{plan.format()}"),
            (
                "\nVERIFIER DIAGNOSIS (what the VM state actually is vs. required):\n"
                + (verifier_diagnosis or "(none provided)")
            ),
            (
                "\nVERIFIER SUGGESTION (a different approach to try this time):\n"
                + (verifier_suggestion or "(none provided)")
            ),
        ]

        if kb_text:
            parts.append(f"\n{kb_text}")

        parts.append("\nThe current screenshot is attached.")
        parts.append(
            "\nINSTRUCTIONS:"
            "\n- The previous attempt failed. Use the diagnosis + suggestion to pick"
            " a different approach — don't just retry the same steps."
            "\n- Call submit_plan with the revised milestone list."
        )

        return "\n".join(parts)
