"""Session-scoped task plan manager — backed by PostgreSQL KV store."""
from __future__ import annotations

import json
from dataclasses import dataclass

from navi.llm.base import Message
from navi.tools._internal.base import (
    Tool,
    ToolContext,
    ToolResult,
    current_session_id,
    current_todo_session_id,
    current_user_id,
)

_STATUS_ICON: dict[str, str] = {
    "pending":     "○",
    "in_progress": "◎",
    "done":        "✓",
    "failed":      "✗",
    "skipped":     "—",
}

@dataclass
class _Task:
    text: str
    status: str = "pending"
    validation: str = ""  # how the result was verified (required for done, encouraged for failed)
    milestone: str = ""   # grouping label (e.g. "A"/"B") from the plan's Milestones; "" = ungrouped


# Global KV store reference — injected at startup by registry.py
_kv_store = None


def set_kv_store(kv) -> None:
    """Inject the shared KvStore instance (called once at startup)."""
    global _kv_store
    _kv_store = kv


def _sid(explicit: str | None = None) -> str:
    # Sub-agents set current_todo_session_id to their ephemeral run id so their
    # plan/todo stay isolated from the parent session's todo. It must win over
    # ctx.session_id (which is still the parent session id for a sub-agent, since
    # other tools rely on it). The parent agent run leaves it unset, so todo falls
    # back to the explicit / regular session id.
    return (
        current_todo_session_id.get()
        or explicit
        or current_session_id.get()
        or "__default__"
    )


def _uid(explicit: str | None = None) -> str | None:
    return explicit if explicit is not None else current_user_id.get(None)


async def _load_tasks(sid: str) -> list[_Task]:
    if _kv_store is None:
        return []
    raw = await _kv_store.get(_uid(), sid, "todo", "tasks")
    if not raw:
        return []
    try:
        data = json.loads(raw)
        return [_Task(**item) for item in data]
    except Exception:
        return []


async def _save_tasks(sid: str, tasks: list[_Task]) -> None:
    if _kv_store is None:
        return
    data = [{"text": t.text, "status": t.status, "validation": t.validation, "milestone": t.milestone} for t in tasks]
    await _kv_store.set(_uid(), sid, "todo", "tasks", json.dumps(data))


class TodoTool(Tool):
    name = "todo"
    description = (
        "Task plan tracker. Your todo list is automatically populated from the plan at the start of each task — "
        "you do NOT need to call 'set'. "
        "Indexes are 1-based. Call 'update' with status='in_progress' when you start a step. "
        "Call 'update' immediately after completing or failing each step — before moving to the next. "
        "When marking a step 'done', you MUST provide a 'validation' field describing how you verified the result. "
        "When marking a step 'failed', provide 'validation' explaining what went wrong and what you tried. "
        "Before final response, make sure every completed step, including the final step, is marked done with validation. "
        "Call 'view' to re-orient yourself after sub-agent execution or long tool chains. "
        "Use 'set' only when you need to replace the plan mid-task (rare). "
        "Statuses: pending → in_progress → done / failed / skipped."
    )
    parameters = {
        "type": "object",
        "properties": {
            "op": {
                "type": "string",
                "enum": ["set", "view", "update", "clear"],
                "description": (
                    "set — create/replace the Master Plan with a list of task milestones; "
                    "view — show the current state of the plan; "
                    "update — change the status of a specific task; "
                    "clear — reset the plan"
                ),
            },
            "tasks": {
                "type": "array",
                "items": {"type": "string"},
                "description": "Ordered list of task descriptions (required for 'set').",
            },
            "index": {
                "type": "integer",
                "description": "1-based task index (required for 'update').",
            },
            "status": {
                "type": "string",
                "enum": ["pending", "in_progress", "done", "failed", "skipped"],
                "description": "New status for the task (required for 'update').",
            },
            "validation": {
                "type": "string",
                "description": (
                    "Required when status='done': briefly describe how you verified the result "
                    "(e.g. 'ran test_tool — output matched expected', 'read the file and confirmed content'). "
                    "Encouraged when status='failed': describe what went wrong and what you tried."
                ),
            },
        },
        "required": ["op"],
    }

    def __init__(self, kv_store=None) -> None:
        if kv_store is not None:
            set_kv_store(kv_store)

    async def execute(self, params: dict, ctx: ToolContext | None = None) -> ToolResult:
        sid = _sid(ctx.session_id if ctx else None)
        op  = params.get("op")

        if op == "set":
            raw = params.get("tasks") or []
            if not raw:
                return ToolResult(success=False, output="", error="'tasks' list is required for 'set'")
            tasks = [_Task(text=str(t)) for t in raw]
            await _save_tasks(sid, tasks)
            return ToolResult(success=True, output=self._render(sid, tasks))

        if op == "view":
            tasks = await _load_tasks(sid)
            if not tasks:
                return ToolResult(success=True, output="No plan set for this session.")
            return ToolResult(success=True, output=self._render(sid, tasks))

        if op == "update":
            tasks = await _load_tasks(sid)
            if not tasks:
                return ToolResult(success=False, output="", error="No plan set. Use 'set' first.")
            idx        = params.get("index")
            status     = params.get("status")
            validation = (params.get("validation") or "").strip()
            if not idx or not status:
                return ToolResult(success=False, output="", error="'index' and 'status' are required for 'update'")
            if idx < 1 or idx > len(tasks):
                return ToolResult(success=False, output="", error=f"index {idx} is out of range (plan has {len(tasks)} tasks)")

            if status == "done" and not validation:
                return ToolResult(
                    success=False,
                    output=(
                        f"Cannot mark step {idx} as done without validation.\n"
                        "Provide a 'validation' field describing how you verified the result before marking it done.\n"
                        "Example: \"ran test_tool — all assertions passed\" or \"read the output file and confirmed expected content\".\n"
                        "If you haven't verified yet — verify first, then update."
                    ),
                    error="validation_required",
                )

            if status == "failed" and not validation:
                # Soft prompt — don't block, but encourage explanation
                tasks[idx - 1].status = status
                await _save_tasks(sid, tasks)
                return ToolResult(
                    success=True,
                    output=(
                        self._render(sid, tasks) + "\n\n"
                        "[Tip: next time add a 'validation' field when marking a step failed — "
                        "describe what went wrong and what you tried. "
                        "This helps with re-planning.]"
                    ),
                )

            tasks[idx - 1].status = status
            tasks[idx - 1].validation = validation
            await _save_tasks(sid, tasks)
            return ToolResult(success=True, output=self._render(sid, tasks))

        if op == "clear":
            if _kv_store is not None:
                await _kv_store.clear_scope(_uid(ctx.user_id if ctx else None), sid, "todo")
            return ToolResult(success=True, output="Plan cleared.")

        return ToolResult(success=False, output="", error=f"Unknown op: {op!r}")

    def _render(self, sid: str, tasks: list[_Task]) -> str:
        if not tasks:
            return "Plan is empty."
        n = len(tasks)
        done = sum(1 for t in tasks if t.status == "done")
        lines = [f"Plan — {done}/{n} done:"]
        for i, t in enumerate(tasks, 1):
            icon   = _STATUS_ICON.get(t.status, "?")
            suffix = f" ({t.status})" if t.status not in ("pending", "done") else ""
            validation_note = f" [verified: {t.validation}]" if t.validation else ""
            milestone_prefix = f"[{t.milestone}] " if t.milestone else ""
            lines.append(f"  {icon} {i}. {milestone_prefix}{t.text}{suffix}{validation_note}")
        return "\n".join(lines)


# ── Public API for agent.py (avoids tight coupling to internal storage) ──

async def get_task_snapshot(session_id: str) -> frozenset[tuple[str, str]]:
    """Return a frozenset of (task_text, status) for the session's todo list."""
    try:
        tasks = await _load_tasks(session_id)
        return frozenset((t.text, t.status) for t in tasks)
    except Exception:
        return frozenset()


async def get_failed_steps(session_id: str) -> frozenset[tuple[int, str]]:
    """Return a frozenset of (1-based index, task_text) for failed steps."""
    try:
        tasks = await _load_tasks(session_id)
        return frozenset(
            (i + 1, t.text)
            for i, t in enumerate(tasks)
            if t.status == "failed"
        )
    except Exception:
        return frozenset()


async def has_open_steps(session_id: str) -> bool:
    """True if the plan has any step still pending/in_progress (i.e. not closed
    as done/skipped/failed). An empty plan returns False.

    Used by the final-turn intercept in agent.py: a "stopped to narrate" turn
    (text without a tool call) only bypasses the tool-loop nudges (anti-stall,
    adaptive re-plan) because those fire in the *next* iteration's pre_turn —
    which never comes. ``has_open_steps`` is the gate that decides whether the
    turn is genuinely finished or should be intercepted and continued.
    """
    snap = await get_task_snapshot(session_id)
    return any(status in ("pending", "in_progress") for _text, status in snap)


async def get_progress_message(session_id: str, *, first_iteration: bool = False) -> Message | None:
    """Build a compact system reminder with current todo state."""
    try:
        tasks = await _load_tasks(session_id)
        if not tasks:
            return None
        n = len(tasks)
        done = sum(1 for t in tasks if t.status == "done")
        failed = sum(1 for t in tasks if t.status == "failed")
        in_progress = sum(1 for t in tasks if t.status == "in_progress")
        pending = n - done - failed - in_progress

        # Progress line
        icon_done = _STATUS_ICON["done"]
        icon_failed = _STATUS_ICON["failed"]
        lines = [f"TODO progress: {done}/{n} {icon_done}  {failed} {icon_failed}  ({in_progress} in progress, {pending} pending)"]

        # List remaining tasks
        remaining = [t for t in tasks if t.status not in ("done", "skipped")]
        if remaining:
            lines.append("Remaining tasks:")
            for i, t in enumerate(remaining, 1):
                lines.append(f"  {i}. {t.text}")
        else:
            lines.append("All tasks completed.")

        # Adaptive discipline note
        if failed > 0 and failed >= done:
            lines.append(
                "Discipline: you have more failures than completions. "
                "Stop and diagnose the root cause before continuing."
            )
        elif failed > 0 and failed == 1 and done >= 2:
            lines.append(
                "Discipline: one failure so far — acceptable if the root cause is understood. "
                "Document the fix in your response."
            )
        elif failed > 1:
            lines.append(
                "Discipline: multiple failures detected. Consider replanning or delegating."
            )

        if first_iteration:
            lines.append(
                "Discipline: this is the FIRST iteration. Prioritise the highest-impact "
                "remaining task and complete it fully before moving to the next."
            )
        elif pending == 0 and failed == 0:
            lines.append(
                "Discipline: all tasks are done or in progress. Finalise the response now."
            )
        else:
            lines.append(
                "Discipline: focus on completing ONE task per iteration. "
                "Do not start the next task until the current one is verified and marked done."
            )

        return Message(role="system", content="\n".join(lines))
    except Exception:
        return None


async def set_tasks(session_id: str, steps) -> None:
    """Auto-populate the todo list from plan steps.

    ``steps`` accepts either the legacy ``list[str]`` (milestone='') or the
    planner's ``list[tuple[milestone, text]]`` produced by
    ``planning._parse_plan_steps``. The legacy form keeps existing callers
    (and tests) working unchanged.
    """
    tasks: list[_Task] = []
    for s in steps:
        if isinstance(s, tuple):
            milestone, text = s
        else:
            milestone, text = "", str(s)
        tasks.append(_Task(text=text, milestone=milestone))
    await _save_tasks(session_id, tasks)


async def step_text_for_update(index, ctx=None) -> str | None:
    """Return the current text of step ``index`` for an in-flight ``update``
    call, so the client can render *which* step is changing in the call card
    (the LLM's update args carry only the index, not the text).

    Resolves the row via the same ``_sid``/``_uid`` logic the tool itself uses,
    so a sub-agent's isolated todo row is read correctly
    (``current_todo_session_id`` wins over ``ctx.session_id``).
    """
    if index is None:
        return None
    try:
        idx = int(index)
    except (TypeError, ValueError):
        return None
    sid = _sid(ctx.session_id if ctx is not None else None)
    tasks = await _load_tasks(sid)
    if not tasks or idx < 1 or idx > len(tasks):
        return None
    return tasks[idx - 1].text


async def started_metadata_for_call(tc, ctx=None) -> dict:
    """Client-rendering metadata for a ``ToolStarted`` event (duck-typed ``tc``
    with ``.name`` and ``.arguments``). For a ``todo`` ``update`` call, attaches
    the step text at the given index; empty for everything else. Used by both
    the parent agent and the sub-agent runner so the call card shows which step
    is changing before the tool runs.
    """
    if tc.name == "todo" and (tc.arguments or {}).get("op") == "update":
        step_text = await step_text_for_update(tc.arguments.get("index"), ctx)
        if step_text:
            return {"step_text": step_text}
    return {}


async def load_tasks_for(
    session_id: str, user_id: str | None, kv=None
) -> list[dict]:
    """Load the todo list for an explicit (user_id, session_id) pair as
    JSON-serialisable dicts — no ContextVars.

    Used by the REST endpoint (GET /sessions/{id}/todos) and the TodoUpdated
    emitter, both of which run outside tool execution where the per-session
    ContextVars are not reliably set. Pass ``kv`` to bypass the module-global
    store (the REST route resolves its own KvStore via deps).
    """
    store = kv if kv is not None else _kv_store
    if store is None:
        return []
    raw = await store.get(user_id, session_id, "todo", "tasks")
    if not raw:
        return []
    try:
        data = json.loads(raw)
    except Exception:
        return []
    return [
        {
            "index": i + 1,
            "text": item.get("text", ""),
            "status": item.get("status", "pending"),
            "validation": item.get("validation", ""),
            "milestone": item.get("milestone", ""),
        }
        for i, item in enumerate(data)
    ]


async def render_todo_lines(session_id: str) -> list[str]:
    """Return a list of formatted todo lines for goal anchoring."""
    tasks = await _load_tasks(session_id)
    if not tasks:
        return []
    lines = []
    for i, t in enumerate(tasks):
        icon = _STATUS_ICON.get(t.status, "?")
        lines.append(f"  {icon} [{i}] {t.text} ({t.status})")
    return lines
