govoplan-identity/src/govoplan_identity/backend/directory.py

from __future__ import annotations

from collections.abc import Sequence

from sqlalchemy import func, or_

from govoplan_core.core.identity import IdentityAccountLinkRef, IdentityDirectory, IdentityRef
from govoplan_core.db.session import get_database
from govoplan_identity.backend.db.models import Identity, IdentityAccountLink


def _status(active: bool) -> str:
    return "active" if active else "inactive"


def _identity_ref(identity: Identity, links: Sequence[IdentityAccountLink]) -> IdentityRef:
    primary_account_id = next((link.account_id for link in links if link.is_primary), None)
    return IdentityRef(
        id=identity.id,
        display_name=identity.display_name,
        external_subject=identity.external_subject,
        source=identity.source,
        primary_account_id=primary_account_id,
        account_ids=tuple(link.account_id for link in links),
        status=_status(identity.is_active),  # type: ignore[arg-type]
    )


def _link_ref(link: IdentityAccountLink) -> IdentityAccountLinkRef:
    return IdentityAccountLinkRef(
        id=link.id,
        identity_id=link.identity_id,
        account_id=link.account_id,
        is_primary=link.is_primary,
        source=link.source,
    )


class SqlIdentityDirectory(IdentityDirectory):
    def get_identity(self, identity_id: str) -> IdentityRef | None:
        with get_database().session() as session:
            identity = session.get(Identity, identity_id)
            if identity is None:
                return None
            links = (
                session.query(IdentityAccountLink)
                .filter(IdentityAccountLink.identity_id == identity.id)
                .order_by(IdentityAccountLink.is_primary.desc(), IdentityAccountLink.created_at.asc())
                .all()
            )
            return _identity_ref(identity, links)

    def identity_for_account(self, account_id: str) -> IdentityRef | None:
        with get_database().session() as session:
            link = (
                session.query(IdentityAccountLink)
                .filter(IdentityAccountLink.account_id == account_id)
                .order_by(IdentityAccountLink.is_primary.desc(), IdentityAccountLink.created_at.asc())
                .first()
            )
            if link is None:
                return None
            identity = session.get(Identity, link.identity_id)
            if identity is None:
                return None
            links = (
                session.query(IdentityAccountLink)
                .filter(IdentityAccountLink.identity_id == identity.id)
                .order_by(IdentityAccountLink.is_primary.desc(), IdentityAccountLink.created_at.asc())
                .all()
            )
            return _identity_ref(identity, links)

    def identities_for_accounts(self, account_ids: Sequence[str]) -> tuple[IdentityRef, ...]:
        ids = sorted({str(account_id) for account_id in account_ids if account_id})
        if not ids:
            return ()
        with get_database().session() as session:
            identity_ids = [
                row[0]
                for row in session.query(IdentityAccountLink.identity_id)
                .filter(IdentityAccountLink.account_id.in_(ids))
                .distinct()
                .all()
            ]
            identities = session.query(Identity).filter(Identity.id.in_(identity_ids)).all() if identity_ids else []
            links = (
                session.query(IdentityAccountLink)
                .filter(IdentityAccountLink.identity_id.in_(identity_ids))
                .order_by(IdentityAccountLink.is_primary.desc(), IdentityAccountLink.created_at.asc())
                .all()
                if identity_ids else []
            )
            links_by_identity: dict[str, list[IdentityAccountLink]] = {}
            for link in links:
                links_by_identity.setdefault(link.identity_id, []).append(link)
            return tuple(_identity_ref(identity, links_by_identity.get(identity.id, [])) for identity in identities)

    def accounts_for_identity(self, identity_id: str) -> tuple[IdentityAccountLinkRef, ...]:
        with get_database().session() as session:
            links = (
                session.query(IdentityAccountLink)
                .filter(IdentityAccountLink.identity_id == identity_id)
                .order_by(IdentityAccountLink.is_primary.desc(), IdentityAccountLink.created_at.asc())
                .all()
            )
            return tuple(_link_ref(link) for link in links)

    def search_identities(
        self,
        query: str | None = None,
        *,
        include_inactive: bool = False,
        limit: int = 25,
    ) -> tuple[IdentityRef, ...]:
        normalized_limit = max(1, min(int(limit), 100))
        with get_database().session() as session:
            identity_query = session.query(Identity)
            if not include_inactive:
                identity_query = identity_query.filter(Identity.is_active.is_(True))
            if query and query.strip():
                pattern = f"%{query.strip().casefold()}%"
                matching_account_links = session.query(IdentityAccountLink.identity_id).filter(
                    func.lower(IdentityAccountLink.account_id).like(pattern)
                )
                identity_query = identity_query.filter(
                    or_(
                        func.lower(Identity.id).like(pattern),
                        func.lower(Identity.display_name).like(pattern),
                        func.lower(Identity.external_subject).like(pattern),
                        Identity.id.in_(matching_account_links),
                    )
                )

            identities = (
                identity_query
                .order_by(Identity.display_name.asc(), Identity.id.asc())
                .limit(normalized_limit)
                .all()
            )
            identity_ids = [identity.id for identity in identities]
            links = (
                session.query(IdentityAccountLink)
                .filter(IdentityAccountLink.identity_id.in_(identity_ids))
                .order_by(IdentityAccountLink.is_primary.desc(), IdentityAccountLink.account_id.asc())
                .all()
                if identity_ids
                else []
            )
            links_by_identity: dict[str, list[IdentityAccountLink]] = {}
            for link in links:
                links_by_identity.setdefault(link.identity_id, []).append(link)
            return tuple(_identity_ref(identity, links_by_identity.get(identity.id, ())) for identity in identities)