"""Unit tests for navi.tools.replan — ReplanTool + ReplanRunner.

Covers:
- ReplanTool schema/param requirements and the no-runner / missing-reason guards.
- ReplanRunner packs reason + goal + todo + scratchpad into replan_context,
  re-runs PlanningEngine with is_replan=True (DIRECT + observe-skip suppressed),
  captures PlanReady.plan, logs PlanningDebugData, and replaces the todo.
- The re-plan prompt framing ([RE-PLAN] block, re-plan opening line) is present
  and the DIRECT shortcut / observe-skip are suppressed for is_replan.
"""

import pytest

from navi.core.events import PlanReady
from navi.core.planning import PlanningEngine
from navi.llm.base import LLMResponse, Message
from navi.tools._internal.base import (
    ToolContext,
    current_replan_runner,
    current_session_id,
    current_user_id,
)
from navi.tools.replan import ReplanRunner, ReplanTool
from navi.tools import scratchpad as scratchpad_mod
from navi.tools import todo as todo_mod
from tests.conftest_factory import FakePool, 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_replan", 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


_REPLAN_ANALYSIS = (
    "TASK: revised task\n"
    "GOAL: revised goal\n"
    "MODE: act\n"
    "UNKNOWNS: NONE\nRESOURCES: NONE\nKNOWLEDGE SOURCE ASSESSMENT: NONE\n"
    "KNOWLEDGE CAPTURE: NONE\nCOMPLEXITY: simple\nSUBTASKS:\n1. New step\n"
    "REFLECT: no\nCOMMITMENTS: none"
)
_REPLAN_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(store, responses, session=None, profile=None):
    session = session or FakeSession()
    profile = profile or make_profile("developer", planning_phase2_enabled=False)
    llm = RecordingLLM(responses)
    engine = PlanningEngine(FakeContextBuilder())
    runner = ReplanRunner(engine, session, profile, llm, mem=None, tool_schemas=[])
    return runner, session, llm, engine, profile


# ── ReplanTool schema ───────────────────────────────────────────────────────


class TestReplanToolSchema:
    def test_name_and_required_params(self):
        t = ReplanTool()
        assert t.name == "replan"
        assert t.parameters["required"] == ["reason"]
        assert "reason" in t.parameters["properties"]
        assert "updated_goal" in t.parameters["properties"]

    async def test_missing_reason_returns_error(self):
        t = ReplanTool()
        result = await t.execute({}, ToolContext())
        assert result.success is False
        assert "reason is required" in (result.error or "")

    async def test_no_runner_returns_error(self):
        # No current_replan_runner set in this context.
        t = ReplanTool()
        result = await t.execute({"reason": "plan is stale"}, ToolContext())
        assert result.success is False
        assert "not available" in (result.error or "")


# ── ReplanRunner + PlanningEngine is_replan integration ──────────────────────


class TestReplanRunner:
    async def test_replan_runs_planner_with_is_replan_and_returns_plan(self):
        runner, session, llm, engine, profile = _make_runner(
            None, [_REPLAN_ANALYSIS, _REPLAN_PLAN]
        )

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

        assert plan == _REPLAN_PLAN
        # Phase 1 + Phase 3 ran (phase2 disabled).
        assert len(llm.calls) == 2
        # Phase 1 prompt carries the [RE-PLAN] framing and the reason.
        phase1_prompt = llm.calls[0][0].content
        assert "[RE-PLAN]" in phase1_prompt
        assert "config is TOML not JSON" in phase1_prompt
        # DIRECT shortcut suppressed for re-plan.
        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.
        # ContextVars (current_session_id / current_user_id) must be set for the
        # whole flow — render_todo_lines reads user_id from current_user_id, and
        # planning.run's set_tasks scopes by current_session_id. In production the
        # WS handler / run_stream set these; here we set them explicitly.
        sid_tok = current_session_id.set("sess_replan")
        uid_tok = current_user_id.set("u1")
        try:
            from navi.tools.todo import set_tasks
            await set_tasks("sess_replan", ["Old step one", "Old step two"])
            await scratchpad_mod._kv_store.set("u1", "sess_replan", "scratchpad", "findings", "config parser expects TOML")
            await scratchpad_mod._kv_store.set("u1", "sess_replan", "scratchpad", "errors", "JSON parse failed at line 3")

            runner, session, llm, engine, profile = _make_runner(
                None, [_REPLAN_ANALYSIS, _REPLAN_PLAN], session=FakeSession(sid="sess_replan", user_id="u1")
            )

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

            assert plan == _REPLAN_PLAN
            phase1_prompt = llm.calls[0][0].content
            # The packed context surfaces all four pieces.
            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_replan_replaces_todo_with_new_steps(self):
        from navi.tools.todo import _load_tasks, set_tasks

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

            runner, session, llm, engine, profile = _make_runner(
                None, [_REPLAN_ANALYSIS, _REPLAN_PLAN], session=FakeSession(sid="sess_replan")
            )

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

            assert plan == _REPLAN_PLAN
            # planning.run's set_tasks replaced the todo with the new plan steps.
            # _parse_plan_steps keeps the full step line including the executor tag.
            tasks = await _load_tasks("sess_replan")
            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_replan_logs_planning_debug_data_to_session(self):
        runner, session, llm, engine, profile = _make_runner(
            None, [_REPLAN_ANALYSIS, _REPLAN_PLAN]
        )

        await runner.replan("stale plan", None)

        # PlanningDebugData log captured (is_subagent=False for replan).
        assert len(session.planning_logs) == 1
        assert session.planning_logs[0]["result"] == "plan"

    async def test_replan_appends_plan_and_prompt_to_session_context(self):
        runner, session, llm, engine, profile = _make_runner(
            None, [_REPLAN_ANALYSIS, _REPLAN_PLAN]
        )

        await runner.replan("stale plan", None)

        # New plan assistant message + execute prompt appended.
        plan_msgs = [m for m in session.context if m.role == "assistant"]
        assert any(_REPLAN_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 len(prompt_msgs) == 1

    async def test_replan_returns_none_on_planner_failure(self):
        # First LLM call raises → runner catches and returns None.
        class _BoomLLM:
            async def complete(self, messages, **kwargs):
                raise RuntimeError("llm down")

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

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


# ── PlanningEngine is_replan prompt/observe-skip behavior ────────────────────


class TestPlanningIsReplan:
    async def test_is_replan_suppresses_observe_skip(self):
        """An observe-classified analysis must NOT skip Phase 3 when is_replan —
        a stale plan always needs a new plan."""
        profile = make_profile(
            "developer",
            planning_phase2_enabled=False,
            observe_skips_plan_enabled=True,
        )
        observe_analysis = (
            "TASK: look at X\nGOAL: report X\nMODE: observe\n"
            "UNKNOWNS: NONE\nRESOURCES: NONE\nKNOWLEDGE SOURCE ASSESSMENT: NONE\n"
            "KNOWLEDGE CAPTURE: NONE\nCOMPLEXITY: simple\nSUBTASKS:\n1. List\n"
            "REFLECT: no\nCOMMITMENTS: none"
        )
        llm = RecordingLLM([observe_analysis, _REPLAN_PLAN])
        engine = PlanningEngine(FakeContextBuilder())
        context = [Message(role="user", content="look at X")]

        events = []
        async for ev in engine.run(
            context, profile, llm, mem=None, tool_schemas=[],
            force_plan=True, is_replan=True, replan_context="Reason for re-plan: stale",
        ):
            events.append(ev)

        # Observe-skip suppressed → Phase 3 ran, plan produced.
        assert len(llm.calls) == 2
        assert any(isinstance(e, PlanReady) for e in events)

    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", planning_phase2_enabled=False)
        llm = RecordingLLM(["DIRECT"])  # force_plan+is_replan suppress DIRECT → needs full analysis, but one call is enough to inspect prompt
        engine = PlanningEngine(FakeContextBuilder())
        context = [Message(role="user", content="hi")]

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

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