"""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())