"""Persistence and row locking for agent grants and credential exchange."""

from uuid import UUID

from sqlalchemy import and_, delete, or_, select, update
from sqlalchemy.ext.asyncio import AsyncSession

from src.domain.agent_access.models import (
    AgentClient,
    AgentCredential,
    AgentGrant,
    AgentGrantWorkspace,
    AgentOperation,
    WorkspaceBillingDelegation,
)
from src.domain.auth.models import User
from src.domain.codegen.models import CodegenRun
from src.domain.projects.models import Project


class AgentRepository:
    def __init__(self, db: AsyncSession):
        self.db = db

    async def defer_retry(self, operation_id: UUID, timestamp):
        await self.db.execute(
            update(AgentOperation)
            .where(
                AgentOperation.id == operation_id,
                AgentOperation.status.in_(["queued", "dispatch_pending", "awaiting_payment"]),
            )
            .values(updated_at=timestamp)
        )

    async def add(self, row):
        self.db.add(row)
        await self.db.flush()
        return row

    async def visible_runs(self, workspace_id: UUID, project_id: UUID, offset: int, limit: int):
        return list(
            (
                await self.db.scalars(
                    select(CodegenRun)
                    .join(Project, Project.id == CodegenRun.project_id)
                    .where(
                        Project.workspace_id == workspace_id,
                        Project.deleted_at.is_(None),
                        CodegenRun.project_id == project_id,
                    )
                    .order_by(CodegenRun.created_at.desc(), CodegenRun.id)
                    .offset(offset)
                    .limit(limit)
                )
            ).all()
        )

    async def visible_run(self, workspace_id: UUID, project_id: UUID, run_id: UUID):
        return (
            await self.db.execute(
                select(CodegenRun)
                .join(Project, Project.id == CodegenRun.project_id)
                .where(
                    Project.workspace_id == workspace_id,
                    Project.deleted_at.is_(None),
                    CodegenRun.project_id == project_id,
                    CodegenRun.id == run_id,
                )
            )
        ).scalar_one_or_none()

    async def visible_projects(self, access_rows, offset: int, limit: int):
        """Projects across the given AgentGrantWorkspace rows (one workspace or several)."""
        if not access_rows:
            return []
        clauses = []
        for row in access_rows:
            clause = Project.workspace_id == row.workspace_id
            if not row.workspace_wide:
                clause = and_(clause, Project.id.in_([UUID(value) for value in row.project_ids]))
            clauses.append(clause)
        query = select(Project).where(Project.deleted_at.is_(None), or_(*clauses))
        return list(
            (
                await self.db.scalars(
                    query.order_by(Project.created_at, Project.id).offset(offset).limit(limit)
                )
            ).all()
        )

    async def visible_project(self, workspace_id: UUID, project_id: UUID):
        return (
            await self.db.execute(
                select(Project).where(
                    Project.workspace_id == workspace_id,
                    Project.id == project_id,
                    Project.deleted_at.is_(None),
                )
            )
        ).scalar_one_or_none()

    async def client(self, client_id: str):
        return await self.db.get(AgentClient, client_id)

    async def grant(self, grant_id: UUID, *, lock: bool = False):
        query = select(AgentGrant).where(AgentGrant.id == grant_id)
        if lock:
            query = query.with_for_update().execution_options(populate_existing=True)
        return (await self.db.execute(query)).scalar_one_or_none()

    async def credential(self, fingerprint: str, *, lock: bool = False):
        query = select(AgentCredential).where(AgentCredential.fingerprint == fingerprint)
        if lock:
            query = query.with_for_update().execution_options(populate_existing=True)
        return (await self.db.execute(query)).scalar_one_or_none()

    async def grant_workspaces(self, grant_id: UUID) -> list[AgentGrantWorkspace]:
        return list(
            (
                await self.db.scalars(
                    select(AgentGrantWorkspace)
                    .where(AgentGrantWorkspace.grant_id == grant_id)
                    .order_by(AgentGrantWorkspace.created_at, AgentGrantWorkspace.workspace_id)
                )
            ).all()
        )

    async def grant_workspaces_by_grant(
        self, grant_ids: list[UUID]
    ) -> dict[UUID, list[AgentGrantWorkspace]]:
        by_grant: dict[UUID, list[AgentGrantWorkspace]] = {grant_id: [] for grant_id in grant_ids}
        if not grant_ids:
            return by_grant
        rows = await self.db.scalars(
            select(AgentGrantWorkspace)
            .where(AgentGrantWorkspace.grant_id.in_(grant_ids))
            .order_by(AgentGrantWorkspace.created_at, AgentGrantWorkspace.workspace_id)
        )
        for row in rows.all():
            by_grant[row.grant_id].append(row)
        return by_grant

    async def remove_grant_workspace(self, grant_id: UUID, workspace_id: UUID):
        await self.db.execute(
            delete(AgentGrantWorkspace).where(
                AgentGrantWorkspace.grant_id == grant_id,
                AgentGrantWorkspace.workspace_id == workspace_id,
            )
        )

    async def list_grants(self, workspace_id: UUID, user_id: UUID | None = None):
        query = (
            select(AgentGrant)
            .join(AgentGrantWorkspace, AgentGrantWorkspace.grant_id == AgentGrant.id)
            .where(AgentGrantWorkspace.workspace_id == workspace_id)
        )
        if user_id is not None:
            query = query.where(AgentGrant.created_by == user_id, AgentGrant.kind == "user")
        return list(
            (await self.db.scalars(query.order_by(AgentGrant.created_at.desc()).limit(1000))).all()
        )

    async def list_grants_created_by(self, user_id: UUID):
        return list(
            (
                await self.db.scalars(
                    select(AgentGrant)
                    .where(AgentGrant.created_by == user_id)
                    .order_by(AgentGrant.created_at.desc())
                    .limit(1000)
                )
            ).all()
        )

    async def billing_delegation(self, workspace_id: UUID, member_id: UUID):
        return (
            await self.db.execute(
                select(WorkspaceBillingDelegation).where(
                    WorkspaceBillingDelegation.workspace_id == workspace_id,
                    WorkspaceBillingDelegation.member_id == member_id,
                )
            )
        ).scalar_one_or_none()

    async def billing_delegation_by_id(
        self, delegation_id: UUID, workspace_id: UUID, *, lock: bool = False
    ):
        query = select(WorkspaceBillingDelegation).where(
            WorkspaceBillingDelegation.id == delegation_id,
            WorkspaceBillingDelegation.workspace_id == workspace_id,
        )
        if lock:
            query = query.with_for_update().execution_options(populate_existing=True)
        return (await self.db.execute(query)).scalar_one_or_none()

    async def billing_delegations_for_workspace(self, workspace_id: UUID):
        """(delegation, member email, member display name) for the owner's review list."""
        result = await self.db.execute(
            select(WorkspaceBillingDelegation, User.email, User.display_name)
            .join(User, User.id == WorkspaceBillingDelegation.member_id)
            .where(WorkspaceBillingDelegation.workspace_id == workspace_id)
            .order_by(WorkspaceBillingDelegation.requested_at.desc())
            .limit(200)
        )
        return [(row, email, name) for row, email, name in result.all()]

    async def operation(self, grant_id: UUID, key: str):
        return (
            await self.db.execute(
                select(AgentOperation).where(
                    AgentOperation.grant_id == grant_id,
                    AgentOperation.request_key == key,
                )
            )
        ).scalar_one_or_none()