smart-org/apps/api/app/infrastructure/repositories/identity.py

"""Repositories for Identity & Access.

Each repository maps between the ORM rows and the pure domain dataclasses, so
the domain layer never sees SQLAlchemy.
"""

from __future__ import annotations

from datetime import UTC, datetime

from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.domain.identity.models import (
    Invite,
    Membership,
    MembershipStatus,
    Organization,
    Role,
    User,
)
from app.infrastructure.db.models import (
    InviteRow,
    MembershipRow,
    OrganizationRow,
    UserRow,
)


def _aware(value: datetime) -> datetime:
    """Re-attach UTC to a timestamp read from the database.

    SQLite drops ``tzinfo`` on round-trip while Postgres keeps it, so comparing
    a freshly computed aware timestamp (``Invite.is_expired`` against
    ``utcnow()``) against one read back raises ``TypeError`` on one backend and
    works on the other. Normalising at the read boundary keeps every consumer —
    domain rules and API output alike — in a single frame.
    """
    if value.tzinfo is None:
        return value.replace(tzinfo=UTC)
    return value


def _user_to_domain(row: UserRow) -> User:
    return User(
        id=row.id,
        email=row.email,
        password_hash=row.password_hash,
        name=row.name,
        avatar_url=row.avatar_url,
        is_platform_owner=row.is_platform_owner,
        is_active=row.is_active,
        created_at=row.created_at,
        updated_at=row.updated_at,
    )


def _org_to_domain(row: OrganizationRow) -> Organization:
    return Organization(
        id=row.id, name=row.name, slug=row.slug,
        owner_user_id=row.owner_user_id, created_at=row.created_at,
    )


def _membership_to_domain(row: MembershipRow) -> Membership:
    return Membership(
        id=row.id,
        organization_id=row.organization_id,
        user_id=row.user_id,
        role=Role(row.role),
        status=MembershipStatus(row.status),
        created_at=row.created_at,
    )


def _invite_to_domain(row: InviteRow) -> Invite:
    return Invite(
        id=row.id,
        organization_id=row.organization_id,
        email=row.email,
        role=Role(row.role),
        token_hash=row.token_hash,
        invited_by_user_id=row.invited_by_user_id,
        expires_at=_aware(row.expires_at),
        accepted_at=_aware(row.accepted_at) if row.accepted_at else None,
        created_at=row.created_at,
    )


class UserRepository:
    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def get(self, user_id: str) -> User | None:
        row = await self.session.get(UserRow, user_id)
        return _user_to_domain(row) if row else None

    async def get_by_email(self, email: str) -> User | None:
        stmt = select(UserRow).where(func.lower(UserRow.email) == email.strip().lower())
        row = (await self.session.execute(stmt)).scalar_one_or_none()
        return _user_to_domain(row) if row else None

    async def create(self, user: User) -> User:
        row = UserRow(
            id=user.id,
            email=user.normalized_email(),
            password_hash=user.password_hash,
            name=user.name,
            avatar_url=user.avatar_url,
            is_platform_owner=user.is_platform_owner,
            is_active=user.is_active,
            created_at=user.created_at,
            updated_at=user.updated_at,
        )
        self.session.add(row)
        await self.session.flush()
        return _user_to_domain(row)

    async def update_password(self, user_id: str, password_hash: str) -> None:
        row = await self.session.get(UserRow, user_id)
        if row:
            row.password_hash = password_hash
            await self.session.flush()

    async def count(self) -> int:
        return int((await self.session.execute(select(func.count(UserRow.id)))).scalar_one())


class OrganizationRepository:
    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def get(self, org_id: str) -> Organization | None:
        row = await self.session.get(OrganizationRow, org_id)
        return _org_to_domain(row) if row else None

    async def get_by_slug(self, slug: str) -> Organization | None:
        stmt = select(OrganizationRow).where(OrganizationRow.slug == slug)
        row = (await self.session.execute(stmt)).scalar_one_or_none()
        return _org_to_domain(row) if row else None

    async def list_all_slugs(self) -> set[str]:
        rows = (await self.session.execute(select(OrganizationRow.slug))).scalars().all()
        return set(rows)

    async def create(self, org: Organization) -> Organization:
        row = OrganizationRow(
            id=org.id, name=org.name, slug=org.slug,
            owner_user_id=org.owner_user_id, created_at=org.created_at,
        )
        self.session.add(row)
        await self.session.flush()
        return _org_to_domain(row)

    async def list_for_user(self, user_id: str) -> list[tuple[Organization, Role]]:
        stmt = (
            select(OrganizationRow, MembershipRow.role)
            .join(MembershipRow, MembershipRow.organization_id == OrganizationRow.id)
            .where(MembershipRow.user_id == user_id, MembershipRow.status == "active")
            .order_by(OrganizationRow.created_at)
        )
        rows = (await self.session.execute(stmt)).all()
        return [(_org_to_domain(o), Role(r)) for o, r in rows]

    async def list_all(self) -> list[Organization]:
        """Every organization. Platform-owner console only — never expose this
        through a tenant-scoped endpoint."""
        stmt = select(OrganizationRow).order_by(OrganizationRow.created_at)
        return [_org_to_domain(r) for r in (await self.session.execute(stmt)).scalars()]

    async def delete(self, org_id: str) -> None:
        row = await self.session.get(OrganizationRow, org_id)
        if row:
            await self.session.delete(row)
            await self.session.flush()

    async def count(self) -> int:
        return int(
            (await self.session.execute(select(func.count(OrganizationRow.id)))).scalar_one()
        )


class MembershipRepository:
    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def get(self, org_id: str, user_id: str) -> Membership | None:
        stmt = select(MembershipRow).where(
            MembershipRow.organization_id == org_id, MembershipRow.user_id == user_id
        )
        row = (await self.session.execute(stmt)).scalar_one_or_none()
        return _membership_to_domain(row) if row else None

    async def get_by_id(self, membership_id: str) -> Membership | None:
        row = await self.session.get(MembershipRow, membership_id)
        return _membership_to_domain(row) if row else None

    async def list_for_org(self, org_id: str) -> list[Membership]:
        stmt = (
            select(MembershipRow)
            .where(MembershipRow.organization_id == org_id)
            .order_by(MembershipRow.created_at)
        )
        return [_membership_to_domain(r) for r in (await self.session.execute(stmt)).scalars()]

    async def add(self, membership: Membership) -> Membership:
        row = MembershipRow(
            id=membership.id,
            organization_id=membership.organization_id,
            user_id=membership.user_id,
            role=membership.role.value,
            status=membership.status.value,
            created_at=membership.created_at,
        )
        self.session.add(row)
        await self.session.flush()
        return _membership_to_domain(row)

    async def set_role(self, membership_id: str, role: Role) -> None:
        row = await self.session.get(MembershipRow, membership_id)
        if row:
            row.role = role.value
            await self.session.flush()

    async def set_status(self, membership_id: str, status: MembershipStatus) -> None:
        row = await self.session.get(MembershipRow, membership_id)
        if row:
            row.status = status.value
            await self.session.flush()

    async def remove(self, membership_id: str) -> None:
        row = await self.session.get(MembershipRow, membership_id)
        if row:
            await self.session.delete(row)
            await self.session.flush()

    async def count_owners(self, org_id: str) -> int:
        stmt = select(func.count(MembershipRow.id)).where(
            MembershipRow.organization_id == org_id,
            MembershipRow.role == Role.OWNER.value,
            MembershipRow.status == MembershipStatus.ACTIVE.value,
        )
        return int((await self.session.execute(stmt)).scalar_one())


class InviteRepository:
    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def get_by_token_hash(self, token_hash: str) -> Invite | None:
        stmt = select(InviteRow).where(InviteRow.token_hash == token_hash)
        row = (await self.session.execute(stmt)).scalar_one_or_none()
        return _invite_to_domain(row) if row else None

    async def list_pending(self, org_id: str) -> list[Invite]:
        stmt = (
            select(InviteRow)
            .where(InviteRow.organization_id == org_id, InviteRow.accepted_at.is_(None))
            .order_by(InviteRow.created_at.desc())
        )
        return [_invite_to_domain(r) for r in (await self.session.execute(stmt)).scalars()]

    async def create(self, invite: Invite) -> Invite:
        row = InviteRow(
            id=invite.id,
            organization_id=invite.organization_id,
            email=invite.email.strip().lower(),
            role=invite.role.value,
            token_hash=invite.token_hash,
            invited_by_user_id=invite.invited_by_user_id,
            expires_at=invite.expires_at,
            accepted_at=invite.accepted_at,
            created_at=invite.created_at,
        )
        self.session.add(row)
        await self.session.flush()
        return _invite_to_domain(row)

    async def mark_accepted(self, invite_id: str, when: datetime) -> None:
        row = await self.session.get(InviteRow, invite_id)
        if row:
            row.accepted_at = when
            await self.session.flush()