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

"""Repositories for Billing: plans, subscriptions, wallets, invoices, payments.

The usage ledger lives in ``repositories/audit.py`` because it is shared with the
observability layer; billing reads it but never writes to it directly.
"""

from __future__ import annotations

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

from app.domain.billing.models import (
    CreditWallet,
    Invoice,
    InvoiceStatus,
    Payment,
    PaymentStatus,
    Plan,
    Subscription,
    SubscriptionStatus,
)
from app.infrastructure.db.models import (
    CreditWalletRow,
    InvoiceRow,
    PaymentRow,
    PlanRow,
    SubscriptionRow,
)


def _plan(row: PlanRow) -> Plan:
    return Plan(
        id=row.id, name=row.name, monthly_price=row.monthly_price,
        included_credit=row.included_credit, features_json=dict(row.features_json or {}),
        is_active=row.is_active, created_at=row.created_at,
    )


def _subscription(row: SubscriptionRow) -> Subscription:
    return Subscription(
        id=row.id, organization_id=row.organization_id, plan_id=row.plan_id,
        status=SubscriptionStatus(row.status),
        external_subscription_id=row.external_subscription_id,
        current_period_end=row.current_period_end, created_at=row.created_at,
    )


def _wallet(row: CreditWalletRow) -> CreditWallet:
    return CreditWallet(
        id=row.id, organization_id=row.organization_id, balance=row.balance,
        currency=row.currency, low_balance_threshold=row.low_balance_threshold,
        updated_at=row.updated_at,
    )


def _invoice(row: InvoiceRow) -> Invoice:
    return Invoice(
        id=row.id, organization_id=row.organization_id, amount=row.amount,
        currency=row.currency, status=InvoiceStatus(row.status),
        external_invoice_id=row.external_invoice_id, period_start=row.period_start,
        period_end=row.period_end, created_at=row.created_at,
    )


def _payment(row: PaymentRow) -> Payment:
    return Payment(
        id=row.id, organization_id=row.organization_id, amount=row.amount,
        status=PaymentStatus(row.status), provider=row.provider,
        external_payment_id=row.external_payment_id, invoice_id=row.invoice_id,
        created_at=row.created_at,
    )


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

    async def list(self) -> list[Plan]:
        stmt = select(PlanRow).order_by(PlanRow.monthly_price)
        return [_plan(r) for r in (await self.session.execute(stmt)).scalars()]

    async def get(self, plan_id: str) -> Plan | None:
        row = await self.session.get(PlanRow, plan_id)
        return _plan(row) if row else None


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

    async def get_for_org(self, organization_id: str) -> Subscription | None:
        stmt = select(SubscriptionRow).where(
            SubscriptionRow.organization_id == organization_id
        )
        row = (await self.session.execute(stmt)).scalars().first()
        return _subscription(row) if row else None

    async def upsert(self, subscription: Subscription) -> Subscription:
        existing = await self.get_for_org(subscription.organization_id)
        if existing:
            row = await self.session.get(SubscriptionRow, existing.id)
            assert row is not None
            row.plan_id = subscription.plan_id
            row.status = str(subscription.status)
            row.external_subscription_id = subscription.external_subscription_id
            row.current_period_end = subscription.current_period_end
            await self.session.flush()
            return _subscription(row)

        row = SubscriptionRow(
            id=subscription.id, organization_id=subscription.organization_id,
            plan_id=subscription.plan_id, status=str(subscription.status),
            external_subscription_id=subscription.external_subscription_id,
            current_period_end=subscription.current_period_end,
            created_at=subscription.created_at,
        )
        self.session.add(row)
        await self.session.flush()
        return _subscription(row)


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

    async def get_for_org(self, organization_id: str) -> CreditWallet | None:
        stmt = select(CreditWalletRow).where(
            CreditWalletRow.organization_id == organization_id
        )
        row = (await self.session.execute(stmt)).scalar_one_or_none()
        return _wallet(row) if row else None

    async def ensure(self, organization_id: str) -> CreditWallet:
        """Return the org's wallet, creating a zero-balance one if absent."""
        existing = await self.get_for_org(organization_id)
        if existing:
            return existing
        row = CreditWalletRow(organization_id=organization_id, balance=0.0, currency="USD")
        self.session.add(row)
        await self.session.flush()
        return _wallet(row)

    async def adjust(self, *, organization_id: str, delta: float) -> CreditWallet:
        """Add (or subtract) credit. Callers must validate the delta's sign."""
        wallet = await self.ensure(organization_id)
        row = await self.session.get(CreditWalletRow, wallet.id)
        assert row is not None
        row.balance = float(row.balance) + float(delta)
        await self.session.flush()
        return _wallet(row)


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

    async def list_for_org(self, organization_id: str) -> list[Invoice]:
        stmt = (
            select(InvoiceRow)
            .where(InvoiceRow.organization_id == organization_id)
            .order_by(InvoiceRow.created_at.desc())
        )
        return [_invoice(r) for r in (await self.session.execute(stmt)).scalars()]

    async def create(self, invoice: Invoice) -> Invoice:
        row = InvoiceRow(
            id=invoice.id, organization_id=invoice.organization_id,
            amount=invoice.amount, currency=invoice.currency,
            status=str(invoice.status),
            external_invoice_id=invoice.external_invoice_id,
            period_start=invoice.period_start, period_end=invoice.period_end,
            created_at=invoice.created_at,
        )
        self.session.add(row)
        await self.session.flush()
        return _invoice(row)


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

    async def list_for_org(self, organization_id: str) -> list[Payment]:
        stmt = (
            select(PaymentRow)
            .where(PaymentRow.organization_id == organization_id)
            .order_by(PaymentRow.created_at.desc())
        )
        return [_payment(r) for r in (await self.session.execute(stmt)).scalars()]

    async def create(self, payment: Payment) -> Payment:
        row = PaymentRow(
            id=payment.id, organization_id=payment.organization_id,
            amount=payment.amount, status=str(payment.status),
            provider=payment.provider, external_payment_id=payment.external_payment_id,
            invoice_id=payment.invoice_id, created_at=payment.created_at,
        )
        self.session.add(row)
        await self.session.flush()
        return _payment(row)

    async def set_status(
        self, payment_id: str, status: PaymentStatus, *, external_id: str | None = None
    ) -> Payment | None:
        row = await self.session.get(PaymentRow, payment_id)
        if row is None:
            return None
        row.status = str(status)
        if external_id is not None:
            row.external_payment_id = external_id
        await self.session.flush()
        return _payment(row)