"""Persistent MCP context and receipt storage in the selected project's BizDB."""

from datetime import datetime, timezone
from psycopg import AsyncConnection
from psycopg.types.json import Jsonb
from app.mcp.models import BoundContext


class McpQueryRepository:
    def __init__(self, conn: AsyncConnection, project_id):
        self.conn, self.project_id = conn, project_id

    async def create(self, context: BoundContext):
        async with self.conn.transaction(), self.conn.cursor() as cur:
            # Bounded maintenance; expired/revoked handles remain invalid even if cleanup fails.
            await cur.execute(
                """DELETE FROM workbench.mcp_query_contexts WHERE context_handle IN
                (SELECT context_handle FROM workbench.mcp_query_contexts
                 WHERE project_id=%s AND (expires_at < now() OR revoked_at IS NOT NULL)
                 ORDER BY expires_at LIMIT 100 FOR UPDATE SKIP LOCKED)""",
                (self.project_id,),
            )
            await cur.execute(
                """INSERT INTO workbench.mcp_query_contexts
                (context_handle,project_id,subject,client_id,document,expires_at)
                VALUES (%s,%s,%s,%s,%s,%s)""",
                (
                    context.context_handle,
                    self.project_id,
                    context.subject,
                    context.client_id,
                    Jsonb(context.model_dump(mode="json")),
                    context.expires_at,
                ),
            )

    async def get(self, handle: str) -> BoundContext | None:
        async with self.conn.transaction(), self.conn.cursor() as cur:
            await cur.execute(
                """SELECT document,revoked_at FROM workbench.mcp_query_contexts
                WHERE context_handle=%s AND project_id=%s""",
                (handle, self.project_id),
            )
            row = await cur.fetchone()
        if not row:
            return None
        document = dict(row["document"])
        document["revoked_at"] = row["revoked_at"]
        return BoundContext.model_validate(document)

    async def revoke(self, handle: str, subject: str, client_id: str):
        async with self.conn.transaction(), self.conn.cursor() as cur:
            await cur.execute(
                """UPDATE workbench.mcp_query_contexts SET revoked_at=COALESCE(revoked_at,%s)
                WHERE context_handle=%s AND project_id=%s AND subject=%s AND client_id=%s""",
                (
                    datetime.now(timezone.utc),
                    handle,
                    self.project_id,
                    subject,
                    client_id,
                ),
            )

    async def receipt(self, scope_key: str, call_id: str) -> dict | None:
        async with self.conn.transaction(), self.conn.cursor() as cur:
            await cur.execute(
                """SELECT parameters,result FROM workbench.mcp_query_receipts
                WHERE project_id=%s AND scope_key=%s AND call_id=%s""",
                (self.project_id, scope_key, call_id),
            )
            return await cur.fetchone()

    async def save_receipt(
        self, scope_key: str, call_id: str, parameters: dict, result: dict
    ):
        # Caller owns the short evidence + receipt transaction.
        async with self.conn.cursor() as cur:
            await cur.execute(
                """INSERT INTO workbench.mcp_query_receipts
                (project_id,scope_key,call_id,parameters,result) VALUES (%s,%s,%s,%s,%s)
                ON CONFLICT (project_id,scope_key,call_id) DO NOTHING""",
                (self.project_id, scope_key, call_id, Jsonb(parameters), Jsonb(result)),
            )

    async def lock_receipt(self, scope_key: str, call_id: str):
        async with self.conn.cursor() as cur:
            await cur.execute(
                "SELECT pg_advisory_xact_lock(hashtextextended(%s,0))",
                (f"mcp:{self.project_id}:{scope_key}:{call_id}",),
            )