"""Unit tests for navi.tools.plan — PlanTool + PlanRunner.

Covers:
- PlanTool schema (both params optional) and the no-runner guard.
- Fresh planning (no reason): planner runs without [RE-PLAN] framing.
- Revision mode (reason given): runner packs reason + goal + todo + scratchpad
  into replan_context, re-runs PlanningEngine with is_replan=True, captures
  PlanReady.plan, logs PlanningDebugData, and replaces the todo.
- PlanningStatus / PlanReady are forwarded to the event sink (UI sees them
  mid-turn); the execute prompt is NOT injected on the tool path.
- The tool result's follow-up instruction follows COMPLEXITY: complex →
  present and wait for confirmation; otherwise proceed.
- The re-plan prompt framing ([RE-PLAN] block) is present only with context.
"""

import asyncio

import pytest

from navi.core.events import PlanReady, PlanningStatus
from navi.core.planning import PlanningEngine
from navi.llm.base import LLMResponse, Message
from navi.tools._internal.base import (
    ToolContext,
    current_event_sink,
    current_plan_runner,
    current_session_id,
    current_user_id,
)
from navi.tools.plan import PlanRunner, PlanTool
from navi.tools import scratchpad as scratchpad_mod
from navi.tools import todo as todo_mod
from tests.conftest_factory import make_profile


class RecordingLLM:
    def __init__(self, responses):
        self.responses = list(responses)
        self.calls = []

    async def complete(self, messages, **kwargs):
        self.calls.append(messages)
        return LLMResponse(
            content=self.responses.pop(0),
            tool_calls=None,
            finish_reason="stop",
        )


class FakeContextBuilder:
    def build_system_prompt(self, profile):
        return "base system prompt"

    def _mcp_context_msg(self, profile=None):
        return None


class FakeKvStore:
    """In-memory KV store for todo + scratchpad tests."""

    def __init__(self):
        self._data: dict[tuple, str] = {}

    async def get(self, user_id, session_id, scope, key):
        return self._data.get((user_id or "", session_id, scope, key))

    async def set(self, user_id, session_id, scope, key, value):
        self._data[(user_id or "", session_id, scope, key)] = value

    async def get_all(self, user_id, session_id, scope):
        return {
            k[3]: v
            for k, v in self._data.items()
            if k[:3] == (user_id or "", session_id, scope)
        }

    async def delete(self, user_id, session_id, scope, key):
        self._data.pop((user_id or "", session_id, scope, key), None)

    async def clear_scope(self, user_id, session_id, scope):
        keys = [k for k in self._data if k[:3] == (user_id or "", session_id, scope)]
        for k in keys:
            del self._data[k]


class FakeSession:
    def __init__(self, sid="sess_plan", user_id="u1"):
        self.id = sid
        self.user_id = user_id
        self.context: list[Message] = []
        self.messages: list[Message] = []
        self.planning_logs: list[dict] = []


@pytest.fixture(autouse=True)
def _fake_kv():
    store = FakeKvStore()
    todo_mod._kv_store = store
    scratchpad_mod._kv_store = store
    yield store
    todo_mod._kv_store = None
    scratchpad_mod._kv_store = None


_FRESH_ANALYSIS = (
    "TASK: build feature\n"
    "GOAL: feature works\n"
    "UNKNOWNS: NONE\nRESOURCES: NONE\nKNOWLEDGE SOURCE ASSESSMENT: NONE\n"
    "KNOWLEDGE CAPTURE: NONE\nCOMPLEXITY: simple\nSUBTASKS:\n1. New step\n"
    "COMMITMENTS: none"
)
_REPLAN_ANALYSIS = (
    "TASK: revised task\n"
    "GOAL: revised goal\n"
    "UNKNOWNS: NONE\nRESOURCES: NONE\nKNOWLEDGE SOURCE ASSESSMENT: NONE\n"
    "KNOWLEDGE CAPTURE: NONE\nCOMPLEXITY: simple\nSUBTASKS:\n1. New step\n"
    "COMMITMENTS: none"
)
_COMPLEX_ANALYSIS = _REPLAN_ANALYSIS.replace("COMPLEXITY: simple", "COMPLEXITY: complex")
_PLAN = (
    "## Plan\n\n"
    "**Task:** revised task\n**Goal:** revised goal\n\n"
    "**Steps:**\n"
    "1. New step one → TOOL: filesystem\n"
    "2. New step two → SELF\n\n"
    "**Parallel:** NONE\n**Risks:** NONE"
)


def _make_runner(responses, session=None, profile=None):
    session = session or FakeSession()
    profile = profile or make_profile("developer")
    llm = RecordingLLM(responses)
    engine = PlanningEngine(FakeContextBuilder())
    runner = PlanRunner(engine, session, profile, llm, mem=None, tool_schemas=[])
    return runner, session, llm, engine, profile


def _bind_runner(runner):
    token = current_plan_runner.set(runner)
    return token


# ── PlanTool schema ───────────────────────────────────────────────────────────


class TestPlanToolSchema:
    def test_name_and_params(self):
        t = PlanTool()
        assert t.name == "plan"
        assert t.parameters["required"] == []
        assert "reason" in t.parameters["properties"]
        assert "updated_goal" in t.parameters["properties"]

    async def test_no_runner_returns_error(self):
        # No current_plan_runner set in this context — both fresh and revision
        # calls must report unavailability.
        t = PlanTool()
        for params in ({}, {"reason": "plan is stale"}):
            result = await t.execute(params, ToolContext())
            assert result.success is False
            assert "not available" in (result.error or "")

    async def test_events_reach_sink_via_contextvar_when_ctx_has_none(self):
        """The agent loop passes ToolContext(event_sink=None) — the ContextVar is
        the real channel. The tool must fall back to current_event_sink, not
        swallow PlanningStatus/PlanReady when ctx.event_sink is None."""
        runner, session, llm, engine, profile = _make_runner([_REPLAN_ANALYSIS, _PLAN])
        sink: asyncio.Queue = asyncio.Queue()
        sink_tok = current_event_sink.set(sink)
        runner_tok = _bind_runner(runner)
        try:
            result = await PlanTool().execute(
                {"reason": "stale"}, ToolContext(event_sink=None)
            )
        finally:
            current_plan_runner.reset(runner_tok)
            current_event_sink.reset(sink_tok)

        assert result.success is True
        events = []
        while not sink.empty():
            events.append(sink.get_nowait())
        assert any(isinstance(e, PlanningStatus) for e in events)
        assert any(isinstance(e, PlanReady) and e.plan == _PLAN for e in events)


# ── PlanRunner + PlanningEngine integration ──────────────────────────────────


class TestPlanRunner:
    async def test_fresh_plan_runs_planner_without_replan_framing(self):
        runner, session, llm, engine, profile = _make_runner([_FRESH_ANALYSIS, _PLAN])

        plan = await runner.plan(None, None, event_sink=None)

        assert plan == _PLAN
        # Phase 1 + Phase 3 ran.
        assert len(llm.calls) == 2
        phase1_prompt = llm.calls[0][0].content
        assert "[RE-PLAN]" not in phase1_prompt
        assert "Reason for re-plan" not in phase1_prompt
        assert "Read the user's latest request." in phase1_prompt

    async def test_revision_runs_planner_with_is_replan_and_returns_plan(self):
        runner, session, llm, engine, profile = _make_runner([_REPLAN_ANALYSIS, _PLAN])

        plan = await runner.plan("config is TOML not JSON", None, event_sink=None)

        assert plan == _PLAN
        assert len(llm.calls) == 2
        phase1_prompt = llm.calls[0][0].content
        assert "[RE-PLAN]" in phase1_prompt
        assert "config is TOML not JSON" in phase1_prompt
        # Top-level runs never see the DIRECT shortcut.
        assert "CRITICAL DIRECT shortcut" not in phase1_prompt
        # Re-plan opening line used instead of the fresh-request one.
        assert "Re-plan based on the [RE-PLAN] context" in phase1_prompt
        assert "Read the user's latest request." not in phase1_prompt

    async def test_replan_context_includes_todo_findings_errors(self):
        # Seed a todo + scratchpad sections the runner must pack into context.
        sid_tok = current_session_id.set("sess_plan")
        uid_tok = current_user_id.set("u1")
        try:
            from navi.tools.todo import set_tasks
            await set_tasks("sess_plan", ["Old step one", "Old step two"])
            await scratchpad_mod._kv_store.set("u1", "sess_plan", "scratchpad", "findings", "config parser expects TOML")
            await scratchpad_mod._kv_store.set("u1", "sess_plan", "scratchpad", "errors", "JSON parse failed at line 3")

            runner, session, llm, engine, profile = _make_runner(
                [_REPLAN_ANALYSIS, _PLAN], session=FakeSession(sid="sess_plan", user_id="u1")
            )

            plan = await runner.plan("config is TOML not JSON", "parse config correctly", event_sink=None)

            assert plan == _PLAN
            phase1_prompt = llm.calls[0][0].content
            assert "Reason for re-plan: config is TOML not JSON" in phase1_prompt
            assert "Updated goal: parse config correctly" in phase1_prompt
            assert "Old step one" in phase1_prompt
            assert "config parser expects TOML" in phase1_prompt
            assert "JSON parse failed at line 3" in phase1_prompt
        finally:
            current_user_id.reset(uid_tok)
            current_session_id.reset(sid_tok)

    async def test_revision_replaces_todo_with_new_steps(self):
        from navi.tools.todo import _load_tasks, set_tasks

        sid_tok = current_session_id.set("sess_plan")
        uid_tok = current_user_id.set("u1")
        try:
            await set_tasks("sess_plan", ["Old step one", "Old step two"])
            assert [t.text for t in await _load_tasks("sess_plan")] == ["Old step one", "Old step two"]

            runner, session, llm, engine, profile = _make_runner(
                [_REPLAN_ANALYSIS, _PLAN], session=FakeSession(sid="sess_plan")
            )

            plan = await runner.plan("stale plan", None, event_sink=None)

            assert plan == _PLAN
            tasks = await _load_tasks("sess_plan")
            assert [t.text for t in tasks] == [
                "New step one → TOOL: filesystem",
                "New step two → SELF",
            ]
        finally:
            current_user_id.reset(uid_tok)
            current_session_id.reset(sid_tok)

    async def test_logs_planning_debug_data_to_session(self):
        runner, session, llm, engine, profile = _make_runner([_REPLAN_ANALYSIS, _PLAN])

        await runner.plan("stale plan", None, event_sink=None)

        assert len(session.planning_logs) == 1
        assert session.planning_logs[0]["result"] == "plan"

    async def test_appends_plan_to_session_context_without_execute_prompt(self):
        """The plan message lands in the session, but the top-level tool path
        must NOT inject the 'Plan is ready. Execute it now' prompt — the tool
        result carries the follow-up instruction instead."""
        runner, session, llm, engine, profile = _make_runner([_REPLAN_ANALYSIS, _PLAN])

        await runner.plan("stale plan", None, event_sink=None)

        plan_msgs = [m for m in session.context if m.role == "assistant"]
        assert any(_PLAN in (m.content or "") for m in plan_msgs)
        prompt_msgs = [
            m for m in session.context
            if m.role == "user" and "Plan is ready" in (m.content or "")
        ]
        assert prompt_msgs == []

    async def test_planning_events_forwarded_to_event_sink(self):
        """PlanningStatus + PlanReady must reach the sink so the UI shows
        planning progress and the plan card mid-turn."""
        runner, session, llm, engine, profile = _make_runner([_REPLAN_ANALYSIS, _PLAN])
        sink: asyncio.Queue = asyncio.Queue()

        await runner.plan("stale plan", None, event_sink=sink)

        events = []
        while not sink.empty():
            events.append(sink.get_nowait())
        assert any(isinstance(e, PlanningStatus) for e in events)
        assert any(isinstance(e, PlanReady) and e.plan == _PLAN for e in events)

    async def test_returns_none_on_planner_failure(self):
        class _BoomLLM:
            async def complete(self, messages, **kwargs):
                raise RuntimeError("llm down")

        session = FakeSession()
        profile = make_profile("developer")
        engine = PlanningEngine(FakeContextBuilder())
        runner = PlanRunner(engine, session, profile, _BoomLLM(), mem=None, tool_schemas=[])

        assert await runner.plan("stale plan", None, event_sink=None) is None


# ── PlanTool follow-up instruction (confirmation by COMPLEXITY) ────────────────


class TestPlanToolConfirmationInstruction:
    async def test_complex_task_asks_to_wait_for_confirmation(self):
        runner, session, llm, engine, profile = _make_runner([_COMPLEX_ANALYSIS, _PLAN])
        token = _bind_runner(runner)
        try:
            result = await PlanTool().execute({"reason": "stale plan"}, ToolContext())
        finally:
            current_plan_runner.reset(token)

        assert result.success is True
        assert "# Revised plan" in result.output
        assert _PLAN in result.output
        assert "COMPLEX" in result.output
        assert "confirmation" in result.output
        assert "Do not start executing yet" in result.output

    async def test_simple_task_proceeds(self):
        runner, session, llm, engine, profile = _make_runner([_FRESH_ANALYSIS, _PLAN])
        token = _bind_runner(runner)
        try:
            result = await PlanTool().execute({}, ToolContext())
        finally:
            current_plan_runner.reset(token)

        assert result.success is True
        assert "# Plan" in result.output
        assert "Proceed with execution now" in result.output
        assert "confirmation" not in result.output

    async def test_fresh_failure_returns_error_without_reason(self):
        class _BoomLLM:
            async def complete(self, messages, **kwargs):
                raise RuntimeError("llm down")

        session = FakeSession()
        profile = make_profile("developer")
        engine = PlanningEngine(FakeContextBuilder())
        runner = PlanRunner(engine, session, profile, _BoomLLM(), mem=None, tool_schemas=[])
        token = _bind_runner(runner)
        try:
            result = await PlanTool().execute({}, ToolContext())
        finally:
            current_plan_runner.reset(token)

        assert result.success is False
        assert "produced no plan" in (result.error or "")

    async def test_revision_failure_returns_error_with_reason(self):
        class _BoomLLM:
            async def complete(self, messages, **kwargs):
                raise RuntimeError("llm down")

        session = FakeSession()
        profile = make_profile("developer")
        engine = PlanningEngine(FakeContextBuilder())
        runner = PlanRunner(engine, session, profile, _BoomLLM(), mem=None, tool_schemas=[])
        token = _bind_runner(runner)
        try:
            result = await PlanTool().execute({"reason": "stale"}, ToolContext())
        finally:
            current_plan_runner.reset(token)

        assert result.success is False
        assert "re-planning produced no plan" in (result.error or "")


# ── PlanningEngine is_replan prompt behavior ──────────────────────────────────


class TestPlanningIsReplan:
    async def test_replan_context_omitted_does_not_inject_block(self):
        """is_replan without replan_context must not inject an empty [RE-PLAN] block."""
        profile = make_profile("developer")
        llm = RecordingLLM([_REPLAN_ANALYSIS])
        engine = PlanningEngine(FakeContextBuilder())
        context = [Message(role="user", content="hi")]

        async for _ev in engine.run(
            context, profile, llm, mem=None, tool_schemas=[],
            is_replan=True, replan_context=None,
        ):
            pass

        phase1_prompt = llm.calls[0][0].content
        assert "[RE-PLAN]" not in phase1_prompt