/application/storage/db/repositories/todos.py

"""Repository for the ``todos`` table.

Covers the operations in ``application/agents/tools/todo_list.py``.

The Mongo schema uses ``todo_id`` (a per-tool monotonic integer that the
LLM uses as its handle) and ``status`` ("open"/"completed"). The Postgres
schema mirrors that with a dedicated ``todo_id INTEGER`` column (unique
per ``tool_id`` via a partial index) for the LLM-facing handle, while the
primary key remains a UUID, and ``status`` is collapsed to a ``completed``
boolean. ``legacy_mongo_id`` lets the backfill stay idempotent across
reruns.
"""

from __future__ import annotations

from typing import Any, Optional

from sqlalchemy import Connection, text

from application.storage.db.base_repository import looks_like_uuid, row_to_dict


class TodosRepository:
    def __init__(self, conn: Connection) -> None:
        self._conn = conn

    def create(
        self,
        user_id: str,
        tool_id: str,
        title: str,
        *,
        todo_id: Optional[int] = None,
        legacy_mongo_id: Optional[str] = None,
    ) -> dict:
        """Insert a todo row.

        Allocates the per-tool monotonic ``todo_id`` inside the same
        transaction when the caller does not supply one. The allocation
        is ``COALESCE(MAX(todo_id), 0) + 1`` scoped to ``tool_id``; the
        partial unique index ``todos_tool_todo_id_uidx`` enforces
        correctness if two callers race.
        """
        if todo_id is None:
            todo_id = self._conn.execute(
                text(
                    "SELECT COALESCE(MAX(todo_id), 0) + 1 FROM todos "
                    "WHERE tool_id = CAST(:tool_id AS uuid)"
                ),
                {"tool_id": tool_id},
            ).scalar_one()

        result = self._conn.execute(
            text(
                """
                INSERT INTO todos (user_id, tool_id, todo_id, title, legacy_mongo_id)
                VALUES (:user_id, CAST(:tool_id AS uuid), :todo_id, :title, :legacy_mongo_id)
                RETURNING *
                """
            ),
            {
                "user_id": user_id,
                "tool_id": tool_id,
                "todo_id": todo_id,
                "title": title,
                "legacy_mongo_id": legacy_mongo_id,
            },
        )
        return row_to_dict(result.fetchone())

    def get(self, todo_uuid: str, user_id: str) -> Optional[dict]:
        """Look up a todo by its UUID primary key."""
        result = self._conn.execute(
            text("SELECT * FROM todos WHERE id = CAST(:id AS uuid) AND user_id = :user_id"),
            {"id": todo_uuid, "user_id": user_id},
        )
        row = result.fetchone()
        return row_to_dict(row) if row is not None else None

    def list_for_user_tool(self, user_id: str, tool_id: str) -> list[dict]:
        result = self._conn.execute(
            text(
                "SELECT * FROM todos WHERE user_id = :user_id "
                "AND tool_id = CAST(:tool_id AS uuid) ORDER BY created_at"
            ),
            {"user_id": user_id, "tool_id": tool_id},
        )
        return [row_to_dict(r) for r in result.fetchall()]

    def list_for_tool(self, user_id: str, tool_id: str) -> list[dict]:
        """Return all todos for a (user, tool) ordered by ``todo_id``."""
        result = self._conn.execute(
            text(
                "SELECT * FROM todos WHERE user_id = :user_id "
                "AND tool_id = CAST(:tool_id AS uuid) "
                "ORDER BY todo_id NULLS LAST, created_at"
            ),
            {"user_id": user_id, "tool_id": tool_id},
        )
        return [row_to_dict(r) for r in result.fetchall()]

    def get_by_tool_and_todo_id(
        self, user_id: str, tool_id: str, todo_id: int
    ) -> Optional[dict]:
        result = self._conn.execute(
            text(
                "SELECT * FROM todos WHERE user_id = :user_id "
                "AND tool_id = CAST(:tool_id AS uuid) AND todo_id = :todo_id"
            ),
            {"user_id": user_id, "tool_id": tool_id, "todo_id": todo_id},
        )
        row = result.fetchone()
        return row_to_dict(row) if row is not None else None

    def update_title(self, todo_uuid: str, user_id: str, title: str) -> bool:
        result = self._conn.execute(
            text(
                "UPDATE todos SET title = :title, updated_at = now() "
                "WHERE id = CAST(:id AS uuid) AND user_id = :user_id"
            ),
            {"id": todo_uuid, "user_id": user_id, "title": title},
        )
        return result.rowcount > 0

    def update_title_by_tool_and_todo_id(
        self, user_id: str, tool_id: str, todo_id: int, title: str
    ) -> bool:
        result = self._conn.execute(
            text(
                "UPDATE todos SET title = :title, updated_at = now() "
                "WHERE user_id = :user_id AND tool_id = CAST(:tool_id AS uuid) "
                "AND todo_id = :todo_id"
            ),
            {
                "user_id": user_id,
                "tool_id": tool_id,
                "todo_id": todo_id,
                "title": title,
            },
        )
        return result.rowcount > 0

    def set_completed(
        self,
        user_id_or_uuid: str,
        tool_id_or_user_id: str,
        todo_id_or_completed: Any,
        completed: Optional[bool] = None,
    ) -> bool:
        """Mark a todo's ``completed`` flag.

        Two call shapes are supported during the migration window:

        * Legacy UUID form (kept for existing tests):
          ``set_completed(todo_uuid, user_id, completed: bool)``.
        * Per-tool integer-handle form (used by the tool's dual-write):
          ``set_completed(user_id, tool_id, todo_id: int, completed: bool)``.
        """
        if completed is None:
            # Legacy three-arg form: (todo_uuid, user_id, completed)
            todo_uuid = user_id_or_uuid
            user_id = tool_id_or_user_id
            completed_value = bool(todo_id_or_completed)
            result = self._conn.execute(
                text(
                    "UPDATE todos SET completed = :completed, updated_at = now() "
                    "WHERE id = CAST(:id AS uuid) AND user_id = :user_id"
                ),
                {"id": todo_uuid, "user_id": user_id, "completed": completed_value},
            )
            return result.rowcount > 0

        # New form: (user_id, tool_id, todo_id, completed)
        user_id = user_id_or_uuid
        tool_id = tool_id_or_user_id
        todo_id = int(todo_id_or_completed)
        result = self._conn.execute(
            text(
                "UPDATE todos SET completed = :completed, updated_at = now() "
                "WHERE user_id = :user_id AND tool_id = CAST(:tool_id AS uuid) "
                "AND todo_id = :todo_id"
            ),
            {
                "user_id": user_id,
                "tool_id": tool_id,
                "todo_id": todo_id,
                "completed": bool(completed),
            },
        )
        return result.rowcount > 0

    def delete(self, todo_uuid: str, user_id: str) -> bool:
        result = self._conn.execute(
            text("DELETE FROM todos WHERE id = CAST(:id AS uuid) AND user_id = :user_id"),
            {"id": todo_uuid, "user_id": user_id},
        )
        return result.rowcount > 0

    def delete_by_tool_and_todo_id(
        self, user_id: str, tool_id: str, todo_id: int
    ) -> bool:
        result = self._conn.execute(
            text(
                "DELETE FROM todos WHERE user_id = :user_id "
                "AND tool_id = CAST(:tool_id AS uuid) AND todo_id = :todo_id"
            ),
            {"user_id": user_id, "tool_id": tool_id, "todo_id": todo_id},
        )
        return result.rowcount > 0

    def get_by_legacy_id(self, legacy_mongo_id: str) -> Optional[dict]:
        legacy_mongo_id = str(legacy_mongo_id) if legacy_mongo_id is not None else None
        result = self._conn.execute(
            text("SELECT * FROM todos WHERE legacy_mongo_id = :legacy"),
            {"legacy": legacy_mongo_id},
        )
        row = result.fetchone()
        return row_to_dict(row) if row is not None else None

    def get_any(self, identifier: str, user_id: str) -> Optional[dict]:
        """Resolve a todo by PG UUID or legacy Mongo ObjectId.

        Picks the lookup path from the id shape so non-UUID input never
        reaches ``CAST(:id AS uuid)`` — that cast raises on the server
        and poisons the enclosing transaction, making any subsequent
        query on the same connection fail.
        """
        if looks_like_uuid(identifier):
            doc = self.get(identifier, user_id)
            if doc is not None:
                return doc
        legacy = self.get_by_legacy_id(identifier)
        if legacy and legacy.get("user_id") == user_id:
            return legacy
        return None

    def update_by_legacy_id(
        self,
        legacy_mongo_id: str,
        *,
        title: Optional[str] = None,
        completed: Optional[bool] = None,
    ) -> bool:
        legacy_mongo_id = str(legacy_mongo_id) if legacy_mongo_id is not None else None
        sets = []
        params: dict[str, Any] = {"legacy": legacy_mongo_id}
        if title is not None:
            sets.append("title = :title")
            params["title"] = title
        if completed is not None:
            sets.append("completed = :completed")
            params["completed"] = bool(completed)
        if not sets:
            return False
        sets.append("updated_at = now()")
        sql = "UPDATE todos SET " + ", ".join(sets) + " WHERE legacy_mongo_id = :legacy"
        result = self._conn.execute(text(sql), params)
        return result.rowcount > 0