TorbatYar/backend/services/loyalty/app/repositories/referral.py

"""Referral Engine repositories — Phase 7.5."""
from __future__ import annotations

from uuid import UUID

from sqlalchemy import func, select

from app.models.referral import ReferralAttribution, ReferralCode, ReferralProgram
from app.models.types import ReferralCodeStatus
from app.repositories.base import TenantBaseRepository


class ReferralProgramRepository(TenantBaseRepository[ReferralProgram]):
    model = ReferralProgram

    async def get_by_code(
        self, tenant_id: UUID, program_id: UUID, code: str
    ) -> ReferralProgram | None:
        stmt = select(ReferralProgram).where(
            ReferralProgram.tenant_id == tenant_id,
            ReferralProgram.program_id == program_id,
            ReferralProgram.code == code,
            ReferralProgram.is_deleted.is_(False),
        )
        result = await self.session.execute(stmt)
        return result.scalar_one_or_none()


class ReferralCodeRepository(TenantBaseRepository[ReferralCode]):
    model = ReferralCode

    async def get_by_code(self, tenant_id: UUID, code: str) -> ReferralCode | None:
        stmt = select(ReferralCode).where(
            ReferralCode.tenant_id == tenant_id,
            ReferralCode.code == code,
        )
        result = await self.session.execute(stmt)
        return result.scalar_one_or_none()

    async def get_active_for_member(
        self, tenant_id: UUID, referral_program_id: UUID, member_id: UUID
    ) -> ReferralCode | None:
        stmt = select(ReferralCode).where(
            ReferralCode.tenant_id == tenant_id,
            ReferralCode.referral_program_id == referral_program_id,
            ReferralCode.member_id == member_id,
            ReferralCode.status == ReferralCodeStatus.ACTIVE,
        )
        result = await self.session.execute(stmt)
        return result.scalar_one_or_none()

    async def list_for_program(
        self,
        tenant_id: UUID,
        referral_program_id: UUID | None = None,
        *,
        offset: int = 0,
        limit: int = 50,
    ):
        clauses = [ReferralCode.tenant_id == tenant_id]
        if referral_program_id is not None:
            clauses.append(ReferralCode.referral_program_id == referral_program_id)
        stmt = (
            select(ReferralCode)
            .where(*clauses)
            .order_by(ReferralCode.created_at.desc())
            .offset(offset)
            .limit(limit)
        )
        result = await self.session.execute(stmt)
        return result.scalars().all()

    async def count_for_program(
        self, tenant_id: UUID, referral_program_id: UUID | None = None
    ) -> int:
        clauses = [ReferralCode.tenant_id == tenant_id]
        if referral_program_id is not None:
            clauses.append(ReferralCode.referral_program_id == referral_program_id)
        stmt = select(func.count()).select_from(ReferralCode).where(*clauses)
        result = await self.session.execute(stmt)
        return int(result.scalar_one())


class ReferralAttributionRepository(TenantBaseRepository[ReferralAttribution]):
    model = ReferralAttribution

    async def get_by_idempotency(
        self, tenant_id: UUID, idempotency_key: str | None
    ) -> ReferralAttribution | None:
        if not idempotency_key:
            return None
        stmt = select(ReferralAttribution).where(
            ReferralAttribution.tenant_id == tenant_id,
            ReferralAttribution.idempotency_key == idempotency_key,
        )
        result = await self.session.execute(stmt)
        return result.scalar_one_or_none()

    async def get_by_referee(
        self, tenant_id: UUID, referral_program_id: UUID, referee_member_id: UUID
    ) -> ReferralAttribution | None:
        stmt = select(ReferralAttribution).where(
            ReferralAttribution.tenant_id == tenant_id,
            ReferralAttribution.referral_program_id == referral_program_id,
            ReferralAttribution.referee_member_id == referee_member_id,
        )
        result = await self.session.execute(stmt)
        return result.scalar_one_or_none()

    async def count_for_referrer(
        self,
        tenant_id: UUID,
        referral_program_id: UUID,
        referrer_member_id: UUID,
        *,
        exclude_statuses: tuple = (),
    ) -> int:
        clauses = [
            ReferralAttribution.tenant_id == tenant_id,
            ReferralAttribution.referral_program_id == referral_program_id,
            ReferralAttribution.referrer_member_id == referrer_member_id,
        ]
        if exclude_statuses:
            clauses.append(ReferralAttribution.status.notin_(exclude_statuses))
        stmt = select(func.count()).select_from(ReferralAttribution).where(*clauses)
        result = await self.session.execute(stmt)
        return int(result.scalar_one())

    async def list_for_program(
        self,
        tenant_id: UUID,
        referral_program_id: UUID | None = None,
        *,
        offset: int = 0,
        limit: int = 50,
    ):
        clauses = [ReferralAttribution.tenant_id == tenant_id]
        if referral_program_id is not None:
            clauses.append(
                ReferralAttribution.referral_program_id == referral_program_id
            )
        stmt = (
            select(ReferralAttribution)
            .where(*clauses)
            .order_by(ReferralAttribution.created_at.desc())
            .offset(offset)
            .limit(limit)
        )
        result = await self.session.execute(stmt)
        return result.scalars().all()

    async def count_for_program(
        self, tenant_id: UUID, referral_program_id: UUID | None = None
    ) -> int:
        clauses = [ReferralAttribution.tenant_id == tenant_id]
        if referral_program_id is not None:
            clauses.append(
                ReferralAttribution.referral_program_id == referral_program_id
            )
        stmt = select(func.count()).select_from(ReferralAttribution).where(*clauses)
        result = await self.session.execute(stmt)
        return int(result.scalar_one())