"""Persistent todo plans — one row per (user, session) in PostgreSQL.

Todo used to live as an opaque JSON string inside the generic KV store
(``session_store``, ``scope='todo'``). It survived restarts there too, but the
row was untyped, every read parsed JSON, and nothing could tell a live plan
from one abandoned by a crashed run. This module gives the plan its own table
with a revision counter and an ``updated_at`` the recovery path can reason
about.

The KV row is *migrated*, then removed: the first time the table is created,
existing plans are copied over so a session that was mid-task before this
change resumes with its plan and statuses intact, and the old rows are deleted
afterwards (a stale copy would otherwise resurrect a plan the user cleared).

The tool API is unchanged — :class:`navi.tools.todo.TodoTool` still speaks
strings; structured per-step fields (``source``, ``step_id``,
``protocol_run_id``) ride along in the JSON and default to empty.
"""

from __future__ import annotations

import asyncio
import json
from datetime import datetime
from typing import Any, Protocol

import asyncpg
import structlog

log = structlog.get_logger()

_DDL = """
CREATE TABLE IF NOT EXISTS session_todos (
    user_id    TEXT NOT NULL DEFAULT '',
    session_id TEXT NOT NULL,
    tasks      JSONB NOT NULL DEFAULT '[]'::jsonb,
    revision   INTEGER NOT NULL DEFAULT 0,
    updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
    PRIMARY KEY (user_id, session_id)
);
"""

# One-time copy of the plans that predate this table. ``ON CONFLICT DO NOTHING``
# makes it idempotent: a plan that already exists here is newer than (or equal
# to) its KV copy, and must not be overwritten by a stale row.
_MIGRATE_INSERT = """
INSERT INTO session_todos (user_id, session_id, tasks, revision, updated_at)
SELECT COALESCE(user_id, ''), session_id, value::jsonb, 0, updated_at
  FROM session_store
 WHERE scope = 'todo' AND key = 'tasks' AND value <> ''
ON CONFLICT (user_id, session_id) DO NOTHING
"""

# The KV row must not survive the migration. Otherwise a plan the user clears
# here would be resurrected from the stale copy on the next process start.
_MIGRATE_DROP_SOURCE = """
DELETE FROM session_store WHERE scope = 'todo' AND key = 'tasks'
"""


def _norm_uid(user_id: str | None) -> str:
    return user_id or ""


class TodoStore(Protocol):
    """Storage contract the todo tool depends on."""

    async def load(self, user_id: str | None, session_id: str) -> list[dict[str, Any]] | None:
        ...

    async def load_meta(self, user_id: str | None, session_id: str) -> dict | None:
        """The row's revision/updated_at, or None when no plan exists."""
        ...

    async def save(self, user_id: str | None, session_id: str, tasks: list[dict[str, Any]]) -> None:
        ...

    async def clear(self, user_id: str | None, session_id: str) -> None:
        ...


class PgTodoStore:
    """Todo plans in their own table; DDL and migration run once, lazily."""

    def __init__(self, pool: asyncpg.Pool) -> None:
        self._pool = pool
        self._ready = False
        self._lock = asyncio.Lock()

    async def _ensure_ready(self) -> None:
        if self._ready:
            return
        async with self._lock:
            if self._ready:
                return
            async with self._pool.acquire() as conn:
                await conn.execute(_DDL)
                try:
                    await conn.execute(_MIGRATE_INSERT)
                    await conn.execute(_MIGRATE_DROP_SOURCE)
                except asyncpg.PostgresError as e:
                    # A missing session_store (fresh install, table created by
                    # another subsystem later) is not a reason to lose todos for
                    # good — log it and carry on with an empty table.
                    log.warning("todo.migration_skipped", error=str(e))
            self._ready = True

    async def load(self, user_id: str | None, session_id: str) -> list[dict[str, Any]] | None:
        await self._ensure_ready()
        async with self._pool.acquire() as conn:
            row = await conn.fetchrow(
                "SELECT tasks FROM session_todos WHERE user_id = $1 AND session_id = $2",
                _norm_uid(user_id), session_id,
            )
        if row is None:
            return None
        raw = row["tasks"]
        if isinstance(raw, str):
            # asyncpg returns JSONB as str unless a codec is registered.
            try:
                raw = json.loads(raw)
            except ValueError:
                return None
        return list(raw) if isinstance(raw, list) else None

    async def load_meta(self, user_id: str | None, session_id: str) -> dict | None:
        await self._ensure_ready()
        async with self._pool.acquire() as conn:
            row = await conn.fetchrow(
                "SELECT revision, updated_at FROM session_todos "
                "WHERE user_id = $1 AND session_id = $2",
                _norm_uid(user_id), session_id,
            )
        return (
            {"revision": row["revision"], "updated_at": row["updated_at"]}
            if row is not None else None
        )

    async def save(self, user_id: str | None, session_id: str, tasks: list[dict[str, Any]]) -> None:
        await self._ensure_ready()
        async with self._pool.acquire() as conn:
            await conn.execute(
                """
                INSERT INTO session_todos (user_id, session_id, tasks, revision, updated_at)
                VALUES ($1, $2, $3::jsonb, 1, now())
                ON CONFLICT (user_id, session_id) DO UPDATE SET
                    tasks      = EXCLUDED.tasks,
                    revision   = session_todos.revision + 1,
                    updated_at = now()
                """,
                _norm_uid(user_id), session_id, json.dumps(tasks, ensure_ascii=False),
            )

    async def clear(self, user_id: str | None, session_id: str) -> None:
        await self._ensure_ready()
        async with self._pool.acquire() as conn:
            await conn.execute(
                "DELETE FROM session_todos WHERE user_id = $1 AND session_id = $2",
                _norm_uid(user_id), session_id,
            )


class KvTodoStore:
    """Adapter over the legacy KV row.

    Kept for two reasons: tests inject a fake KvStore, and an instance that has
    not run the migration yet still answers from the old location.
    """

    def __init__(self, kv) -> None:
        self._kv = kv

    async def load(self, user_id: str | None, session_id: str) -> list[dict[str, Any]] | None:
        raw = await self._kv.get(_norm_uid(user_id), session_id, "todo", "tasks")
        if not raw:
            return None
        try:
            data = json.loads(raw)
        except ValueError:
            return None
        return list(data) if isinstance(data, list) else None

    async def load_meta(self, user_id: str | None, session_id: str) -> dict | None:
        # The KV row carries no revision and no timestamp the store exposes.
        return None

    async def save(self, user_id: str | None, session_id: str, tasks: list[dict[str, Any]]) -> None:
        await self._kv.set(
            _norm_uid(user_id), session_id, "todo", "tasks",
            json.dumps(tasks, ensure_ascii=False),
        )

    async def clear(self, user_id: str | None, session_id: str) -> None:
        await self._kv.clear_scope(_norm_uid(user_id), session_id, "todo")


def plan_timestamp(meta: dict | None) -> datetime | None:
    """``updated_at`` from a :meth:`TodoStore.load_meta` result, if it has one."""
    if not meta:
        return None
    value = meta.get("updated_at")
    return value if isinstance(value, datetime) else None