"""Project scoped private message storage, sharing task transaction locks."""
from psycopg import AsyncConnection

from app.infra.db.postgresql.operational_tasks import OperationalTaskRepository


class AgentMessageRepository(OperationalTaskRepository):
    def __init__(self, connection: AsyncConnection, project_id):
        super().__init__(connection, project_id)

    async def get_message(self, message_id: str) -> dict | None:
        cursor = await self.connection.execute(
            "SELECT * FROM workbench.agent_messages WHERE project_id=%s AND message_id=%s",
            (self.project_id, message_id),
        )
        row = await cursor.fetchone()
        return self._message(row) if row else None

    async def get_message_request(self, owner_id: str, session_id: str, request_id: str) -> dict | None:
        cursor = await self.connection.execute(
            """SELECT * FROM workbench.agent_messages WHERE project_id=%s
               AND source_owner_id=%s AND source_session_id=%s AND request_id=%s""",
            (self.project_id, owner_id, session_id, request_id),
        )
        row = await cursor.fetchone()
        return self._message(row) if row else None

    @staticmethod
    def _message(row) -> dict:
        return dict(row)

    async def insert_message(self, message: dict) -> None:
        keys = [key for key in message if key != "project_id"]
        await self.connection.execute(
            f"INSERT INTO workbench.agent_messages (project_id,{','.join(keys)}) "
            f"VALUES ({','.join(['%s'] * (len(keys) + 1))})",
            (self.project_id, *(message[key] for key in keys)),
        )

    async def list_session_messages(
        self, owner_id: str, session_id: str, limit: int, *, cursor_at=None, cursor_id=None, pending_only=False,
    ) -> list[dict]:
        cursor = await self.connection.execute(
            """SELECT * FROM workbench.agent_messages WHERE project_id=%s
               AND conversation_owner_id=%s AND (source_session_id=%s OR
                    (direction='to_chat' AND target_session_id=%s))
               AND (%s::timestamptz IS NULL OR (created_at,message_id)<(%s::timestamptz,%s::uuid))
               AND (NOT %s OR status IN ('pending','enqueued'))
               ORDER BY created_at DESC,message_id DESC LIMIT %s""",
            (self.project_id, owner_id, session_id, session_id, cursor_at, cursor_at, cursor_id, pending_only, limit),
        )
        return [self._message(row) for row in await cursor.fetchall()]

    async def list_run_session_messages(
        self, session_id: str, limit: int, *, cursor_at=None, cursor_id=None, pending_only=False,
    ) -> list[dict]:
        cursor = await self.connection.execute(
            """SELECT m.* FROM workbench.agent_messages m
               JOIN workbench.operational_task_runs r ON r.project_id=m.project_id AND r.run_id=m.run_id
               WHERE m.project_id=%s AND r.agent_session_id=%s
                 AND r.status IN ('claimed','running')
                 AND (%s::timestamptz IS NULL OR (m.created_at,m.message_id)<(%s::timestamptz,%s::uuid))
                 AND (NOT %s OR (m.direction='to_task' AND m.status IN ('pending','enqueued')))
               ORDER BY m.created_at DESC,m.message_id DESC LIMIT %s""",
            (self.project_id, session_id, cursor_at, cursor_at, cursor_id, pending_only, limit),
        )
        return [self._message(row) for row in await cursor.fetchall()]

    async def list_pending_messages(self, limit: int, *, cursor_at=None, cursor_id=None) -> list[dict]:
        cursor = await self.connection.execute(
            """SELECT m.* FROM workbench.agent_messages m
               JOIN workbench.operational_task_runs r
                 ON r.project_id=m.project_id AND r.run_id=m.run_id
               WHERE m.project_id=%s AND m.status IN ('pending','enqueued')
                 AND (m.direction='to_chat' OR r.status IN ('claimed','running'))
                 AND (%s::timestamptz IS NULL OR (m.created_at,m.message_id)>(%s::timestamptz,%s::uuid))
               ORDER BY m.created_at,m.message_id LIMIT %s""",
            (self.project_id, cursor_at, cursor_at, cursor_id, limit),
        )
        return [self._message(row) for row in await cursor.fetchall()]

    async def update_message_delivery(self, message_id: str, update: dict) -> None:
        keys = list(update)
        await self.connection.execute(
            f"UPDATE workbench.agent_messages SET {','.join(key+'=%s' for key in keys)} "
            "WHERE project_id=%s AND message_id=%s",
            (*(update[key] for key in keys), self.project_id, message_id),
        )