"""
replan — re-run the planner mid-task when the current plan is stale.

Use when discoveries during execution invalidate the original plan (a step turned
out to be unnecessary, the real problem is different from the assumed one, new
constraints appeared). This is NOT for failed steps — those are handled by the
[Adaptive re-plan] signal, which nudges you to revise the todo inline without
re-running the planner.

Integrated approach: replan reuses the same PlanningEngine pipeline (analysis →
review → execution plan) with `is_replan=True`, which frames Phase 1 as a revision
of the existing plan and suppresses the DIRECT shortcut and observe-skip so a
new plan is always produced. The current todo + scratchpad findings/errors are
packed into the re-plan context so the new plan accounts for real progress. The
new plan replaces the todo (set_tasks) — completed work is preserved in the
scratchpad/conversation, but the remaining steps are overwritten.

Costs 1–3 LLM calls (Phase 1 always; Phase 2 only if REFLECT: yes; Phase 3
always). Use selectively, like `reflect`.
"""

from __future__ import annotations

import structlog

from navi.tools._internal.base import Tool, ToolContext, ToolResult, current_replan_runner
from navi.tools.scratchpad import get_section as scratchpad_get_section
from navi.tools.todo import render_todo_lines

log = structlog.get_logger()

# Mirror agent.py's cap so re-plan debug logs don't grow session.planning_logs unbounded.
_MAX_PLANNING_LOGS = 20


class ReplanRunner:
    """Per-run helper bound to the live session/profile/llm.

    Constructed inside the agent run_stream loop (so it sees the current
    iteration's profile/llm/tool_schemas — correct after switch_profile) and
    exposed to the `replan` tool via the `current_replan_runner` ContextVar.
    """

    def __init__(self, planning_engine, session, profile, llm, mem, tool_schemas) -> None:
        self._planning = planning_engine
        self._session = session
        self._profile = profile
        self._llm = llm
        self._mem = mem
        self._tool_schemas = tool_schemas

    async def replan(self, reason: str, updated_goal: str | None) -> str | None:
        """Re-run the planner over the live session context + todo + scratchpad.

        Returns the new plan text, or None if planning produced no plan. The
        planner mutates session.context/messages (appends the new plan + execute
        prompt) and replaces the todo via set_tasks.
        """
        session = self._session
        todo_lines = await render_todo_lines(session.id)
        findings = await scratchpad_get_section(session.id, "findings", user_id=session.user_id)
        errors = await scratchpad_get_section(session.id, "errors", user_id=session.user_id)

        todo_block = "\n".join(todo_lines) if todo_lines else "(empty — no todo steps)"
        findings_block = findings.strip() if findings and findings.strip() else "(none)"
        errors_block = errors.strip() if errors and errors.strip() else "(none)"
        goal_line = updated_goal.strip() if updated_goal and updated_goal.strip() else "(not specified — keep the original goal)"

        replan_context = (
            f"Reason for re-plan: {reason}\n"
            f"Updated goal: {goal_line}\n\n"
            f"Current todo:\n{todo_block}\n\n"
            f"Scratchpad — findings:\n{findings_block}\n\n"
            f"Scratchpad — errors:\n{errors_block}"
        )

        plan_text: str | None = None
        try:
            # Lazy import: navi.core.events pulls in navi.core.__init__ → registry →
            # navi.tools, which would cycle if done at module load. At runtime every
            # module is already imported, so this is safe.
            from navi.core.events import (
                AIHelperTokensUsed,
                PlanReady,
                PlanningDebugData,
                PlanningStatus,
            )

            async for ev in self._planning.run(
                session.context,
                self._profile,
                self._llm,
                self._mem,
                self._tool_schemas,
                messages=session.messages,
                force_plan=True,
                is_replan=True,
                replan_context=replan_context,
            ):
                if isinstance(ev, PlanReady):
                    plan_text = ev.plan
                elif isinstance(ev, PlanningDebugData):
                    session.planning_logs.append(ev.log)
                    if len(session.planning_logs) > _MAX_PLANNING_LOGS:
                        session.planning_logs = session.planning_logs[-_MAX_PLANNING_LOGS:]
                # PlanningStatus / AIHelperTokensUsed are swallowed — replan is a
                # tool call, not a top-level turn, so we don't re-yield UI events.
        except Exception:
            log.warning("replan.runner_failed", exc_info=True)
            return None

        return plan_text


class ReplanTool(Tool):
    name = "replan"
    description = (
        "Re-plan mid-task when the current plan is stale because of what you discovered during execution — "
        "NOT because a step failed (that is handled by the [Adaptive re-plan] signal).\n\n"
        "Call this when:\n"
        "- A step turned out to be unnecessary or the real problem differs from the assumed one.\n"
        "- New constraints, files, or facts surfaced that change the remaining work.\n"
        "- The remaining steps no longer lead to the goal.\n\n"
        "It re-runs the planner over your current context + todo + scratchpad findings/errors and "
        "REPLACES the plan and todo with the new steps. Costs 1–3 LLM calls — use selectively, like `reflect`.\n\n"
        "Do NOT call replan for a single failed step: revise the todo inline instead. "
        "For a small adjustment (drop/merge/reorder 1-2 steps) edit the `todo` directly — replan is for when the remaining plan's overall structure is wrong. "
        "Use `reflect` when you're unsure what's wrong (it surfaces assumptions, no plan change); use `replan` when you know what changed and need a new plan."
    )
    parameters = {
        "type": "object",
        "properties": {
            "reason": {
                "type": "string",
                "description": (
                    "What changed — the discovery that makes the current plan stale. "
                    "One or two sentences, concrete (e.g. 'the config is TOML not JSON, so the parser step is wrong')."
                ),
            },
            "updated_goal": {
                "type": "string",
                "description": (
                    "Optional. The new success criterion if the goal itself shifted. "
                    "Omit to keep the original goal."
                ),
            },
        },
        "required": ["reason"],
    }

    async def execute(self, params: dict, ctx: ToolContext | None = None) -> ToolResult:
        reason = (params.get("reason") or "").strip()
        updated_goal = (params.get("updated_goal") or "").strip() or None

        if not reason:
            return ToolResult(success=False, output="", error="reason is required")

        runner = current_replan_runner.get(None)
        if runner is None:
            return ToolResult(
                success=False,
                output="",
                error="replan is not available in this context (no active agent run)",
            )

        plan_text = await runner.replan(reason, updated_goal)
        if not plan_text:
            return ToolResult(
                success=False,
                output="",
                error="re-planning produced no plan; keep the current plan and continue, or revise the todo inline",
            )

        output = (
            f"# Revised plan\n\n{plan_text}\n\n---\n"
            "Plan revised. The todo has been replaced with the new steps. "
            "Continue from step 1 of the new plan; do not resume the old steps."
        )
        return ToolResult(success=True, output=output)