"""Billing domain - data access layer."""
from __future__ import annotations
from datetime import datetime, timezone
from typing import Any, Optional
from uuid import UUID

from sqlalchemy import select, func as sa_func, text
from sqlalchemy.ext.asyncio import AsyncSession

from src.domain.billing.models import (
    EnterpriseInquiry,
    Invoice,
    StripeCustomer,
    StripeEvent,
    Subscription,
    UsageCreditPackage,
    UsageCreditPurchase,
    UsageLedger,
    WorkspaceCreditLedger,
    WorkspaceUsageAllowance,
)


def _stripe_timestamp_to_utc_datetime(value: Any) -> Optional[datetime]:
    """Stripe JSON webhooks use Unix seconds for period fields; asyncpg expects datetime."""
    if value is None:
        return None
    if isinstance(value, datetime):
        if value.tzinfo is None:
            return value.replace(tzinfo=timezone.utc)
        return value
    if isinstance(value, (int, float)):
        return datetime.fromtimestamp(int(value), tz=timezone.utc)
    return None


class StripeCustomerRepository:
    def __init__(self, db: AsyncSession):
        self._db = db

    async def get_by_workspace(self, workspace_id: UUID) -> Optional[StripeCustomer]:
        result = await self._db.execute(
            select(StripeCustomer).where(StripeCustomer.workspace_id == workspace_id)
        )
        return result.scalar_one_or_none()

    async def get_by_stripe_id(self, stripe_customer_id: str) -> Optional[StripeCustomer]:
        result = await self._db.execute(
            select(StripeCustomer).where(StripeCustomer.stripe_customer_id == stripe_customer_id)
        )
        return result.scalar_one_or_none()

    async def create(self, workspace_id: UUID, stripe_customer_id: str) -> StripeCustomer:
        c = StripeCustomer(
            workspace_id=workspace_id,
            stripe_customer_id=stripe_customer_id,
        )
        self._db.add(c)
        await self._db.flush()
        await self._db.refresh(c)
        return c


class SubscriptionRepository:
    def __init__(self, db: AsyncSession):
        self._db = db

    async def get_by_workspace(self, workspace_id: UUID) -> Optional[Subscription]:
        """Get active/trialing/past_due/complimentary subscription for workspace."""
        result = await self._db.execute(
            select(Subscription)
            .where(Subscription.workspace_id == workspace_id)
            .where(Subscription.status.in_(("active", "trialing", "past_due", "complimentary")))
            .order_by(Subscription.created_at.desc())
            .limit(1)
        )
        return result.scalar_one_or_none()

    async def create_for_no_charge_plan(self, workspace_id: UUID, plan_slug: str) -> Subscription:
        """Create an internal subscription record for free/staff plans (no Stripe counterpart)."""
        sub = Subscription(
            workspace_id=workspace_id,
            stripe_subscription_id=None,
            stripe_customer_id=None,
            status="active",
            plan_slug=plan_slug,
        )
        self._db.add(sub)
        await self._db.flush()
        await self._db.refresh(sub)
        return sub

    async def get_latest_for_workspace(self, workspace_id: UUID) -> Optional[Subscription]:
        """Get latest subscription for workspace (any status) for billing status display."""
        result = await self._db.execute(
            select(Subscription)
            .where(Subscription.workspace_id == workspace_id)
            .order_by(Subscription.created_at.desc())
            .limit(1)
        )
        return result.scalar_one_or_none()

    async def get_by_stripe_id(self, stripe_subscription_id: str) -> Optional[Subscription]:
        result = await self._db.execute(
            select(Subscription).where(Subscription.stripe_subscription_id == stripe_subscription_id)
        )
        return result.scalar_one_or_none()

    async def upsert_from_stripe(self, workspace_id: UUID, data: dict) -> Subscription:
        result = await self._db.execute(
            select(Subscription).where(Subscription.stripe_subscription_id == data["id"])
        )
        sub = result.scalar_one_or_none()
        if sub:
            sub.status = data.get("status", sub.status)
            sub.plan_id = data.get("plan_id", sub.plan_id)
            sub.plan_name = data.get("plan_name", sub.plan_name)
            sub.plan_slug = data.get("plan_slug") or sub.plan_slug
            sub.stripe_product_id = data.get("stripe_product_id") or sub.stripe_product_id
            sub.stripe_price_id = data.get("stripe_price_id") or data.get("plan_id") or sub.stripe_price_id
            sub.billing_interval = data.get("billing_interval") or sub.billing_interval
            if data.get("current_period_start") is not None:
                sub.current_period_start = _stripe_timestamp_to_utc_datetime(
                    data.get("current_period_start")
                )
            if data.get("current_period_end") is not None:
                sub.current_period_end = _stripe_timestamp_to_utc_datetime(
                    data.get("current_period_end")
                )
            sub.cancel_at_period_end = data.get("cancel_at_period_end", sub.cancel_at_period_end)
            if data.get("ended_at") is not None:
                sub.ended_at = _stripe_timestamp_to_utc_datetime(data.get("ended_at"))
            sub.subscription_plan_id = data.get("subscription_plan_id")
            if data.get("canceled_at") is not None:
                sub.canceled_at = _stripe_timestamp_to_utc_datetime(data.get("canceled_at"))
            if data.get("trial_start") is not None:
                sub.trial_start = _stripe_timestamp_to_utc_datetime(data.get("trial_start"))
            if data.get("trial_end") is not None:
                sub.trial_end = _stripe_timestamp_to_utc_datetime(data.get("trial_end"))
            if data.get("metered_item_id"):
                sub.metered_subscription_item_id = data["metered_item_id"]
            if data.get("metadata_json") is not None:
                sub.metadata_json = data["metadata_json"]
        else:
            sub = Subscription(
                workspace_id=workspace_id,
                stripe_subscription_id=data["id"],
                stripe_customer_id=data.get("customer", ""),
                status=data.get("status", "active"),
                plan_id=data.get("plan_id"),
                plan_name=data.get("plan_name"),
                plan_slug=data.get("plan_slug"),
                stripe_product_id=data.get("stripe_product_id"),
                stripe_price_id=data.get("stripe_price_id") or data.get("plan_id"),
                billing_interval=data.get("billing_interval", "month"),
                cancel_at_period_end=data.get("cancel_at_period_end", False),
                subscription_plan_id=data.get("subscription_plan_id"),
                metered_subscription_item_id=data.get("metered_item_id"),
                current_period_start=_stripe_timestamp_to_utc_datetime(data.get("current_period_start")),
                current_period_end=_stripe_timestamp_to_utc_datetime(data.get("current_period_end")),
                canceled_at=_stripe_timestamp_to_utc_datetime(data.get("canceled_at")),
                trial_start=_stripe_timestamp_to_utc_datetime(data.get("trial_start")),
                trial_end=_stripe_timestamp_to_utc_datetime(data.get("trial_end")),
                ended_at=_stripe_timestamp_to_utc_datetime(data.get("ended_at")),
            )
            self._db.add(sub)
        await self._db.flush()
        await self._db.refresh(sub)
        return sub


class InvoiceRepository:
    def __init__(self, db: AsyncSession):
        self._db = db

    async def list_for_workspace(self, workspace_id: UUID) -> list[Invoice]:
        result = await self._db.execute(
            select(Invoice)
            .where(Invoice.workspace_id == workspace_id)
            .order_by(Invoice.created_at.desc())
        )
        return list(result.scalars().all())

    async def upsert_from_stripe(self, workspace_id: UUID, data: dict) -> Invoice:
        result = await self._db.execute(
            select(Invoice).where(Invoice.stripe_invoice_id == data["id"])
        )
        inv = result.scalar_one_or_none()
        if inv:
            inv.status = data.get("status", inv.status)
            if data.get("amount_due") is not None:
                inv.amount_due = int(data["amount_due"])
            if data.get("amount_paid") is not None:
                inv.amount_paid = int(data["amount_paid"])
            if data.get("hosted_invoice_url"):
                inv.invoice_url = data["hosted_invoice_url"]
            if data.get("invoice_pdf"):
                inv.invoice_pdf = data["invoice_pdf"]
            ps = _stripe_timestamp_to_utc_datetime(data.get("period_start"))
            pe = _stripe_timestamp_to_utc_datetime(data.get("period_end"))
            if ps is not None:
                inv.period_start = ps
            if pe is not None:
                inv.period_end = pe
        else:
            inv = Invoice(
                workspace_id=workspace_id,
                stripe_invoice_id=data["id"],
                stripe_customer_id=data.get("customer", ""),
                status=data.get("status", "draft"),
                amount_due=int(data["amount_due"]) if data.get("amount_due") is not None else 0,
                amount_paid=int(data["amount_paid"]) if data.get("amount_paid") is not None else 0,
                currency=data.get("currency", "usd"),
                invoice_url=data.get("hosted_invoice_url"),
                invoice_pdf=data.get("invoice_pdf"),
                period_start=_stripe_timestamp_to_utc_datetime(data.get("period_start")),
                period_end=_stripe_timestamp_to_utc_datetime(data.get("period_end")),
            )
            self._db.add(inv)
        await self._db.flush()
        await self._db.refresh(inv)
        return inv


class StripeEventRepository:
    def __init__(self, db: AsyncSession):
        self._db = db

    async def already_processed(self, stripe_event_id: str) -> bool:
        result = await self._db.execute(
            select(StripeEvent.id).where(StripeEvent.stripe_event_id == stripe_event_id)
        )
        return result.scalar_one_or_none() is not None

    async def record(self, stripe_event_id: str, event_type: str, payload: Optional[dict] = None) -> StripeEvent:
        e = StripeEvent(stripe_event_id=stripe_event_id, event_type=event_type, payload=payload)
        self._db.add(e)
        await self._db.flush()
        return e


class UsageLedgerRepository:
    def __init__(self, db: AsyncSession):
        self._db = db

    async def record(
        self,
        workspace_id: UUID,
        event_type: str,
        quantity: int = 1,
        unit_label: Optional[str] = None,
        project_id: Optional[UUID] = None,
        metadata: Optional[dict] = None,
        *,
        reported_to_stripe: bool = False,
    ) -> UsageLedger:
        entry = UsageLedger(
            workspace_id=workspace_id,
            project_id=project_id,
            event_type=event_type,
            quantity=quantity,
            unit_label=unit_label,
            metadata_json=metadata,
            reported_to_stripe=reported_to_stripe,
        )
        self._db.add(entry)
        await self._db.flush()
        return entry

    async def find_by_cloud_job_id(
        self,
        workspace_id: UUID,
        job_id: str,
    ) -> Optional[UsageLedger]:
        """Idempotency for Cloud usage finalize/release (metadata contains job_id)."""
        result = await self._db.execute(
            select(UsageLedger)
            .where(UsageLedger.workspace_id == workspace_id)
            .where(UsageLedger.metadata_json.contains({"job_id": job_id}))
            .limit(1)
        )
        return result.scalar_one_or_none()

    async def get_unreported(self, workspace_id: UUID) -> list[UsageLedger]:
        result = await self._db.execute(
            select(UsageLedger)
            .where(UsageLedger.workspace_id == workspace_id)
            .where(UsageLedger.reported_to_stripe.is_(False))
            .order_by(UsageLedger.recorded_at)
        )
        return list(result.scalars().all())

    async def mark_reported(self, entry_ids: list[UUID]) -> None:
        if not entry_ids:
            return
        result = await self._db.execute(
            select(UsageLedger).where(UsageLedger.id.in_(entry_ids))
        )
        for entry in result.scalars().all():
            entry.reported_to_stripe = True
        await self._db.flush()

    async def summarize(self, workspace_id: UUID, since: Optional[datetime] = None) -> list[dict]:
        q = (
            select(
                UsageLedger.event_type,
                sa_func.sum(UsageLedger.quantity).label("total"),
                UsageLedger.unit_label,
            )
            .where(UsageLedger.workspace_id == workspace_id)
            .group_by(UsageLedger.event_type, UsageLedger.unit_label)
        )
        if since:
            q = q.where(UsageLedger.recorded_at >= since)
        result = await self._db.execute(q)
        return [
            {"event_type": row.event_type, "total_quantity": row.total, "unit_label": row.unit_label}
            for row in result.all()
        ]


class UsageCreditPackageRepository:
    def __init__(self, db: AsyncSession):
        self._db = db

    async def get_by_id(self, package_id: UUID) -> Optional[UsageCreditPackage]:
        result = await self._db.execute(
            select(UsageCreditPackage).where(UsageCreditPackage.id == package_id)
        )
        return result.scalar_one_or_none()

    async def get_by_slug(self, slug: str) -> Optional[UsageCreditPackage]:
        result = await self._db.execute(
            select(UsageCreditPackage).where(UsageCreditPackage.slug == slug)
        )
        return result.scalar_one_or_none()

    async def get_by_stripe_price_id(self, stripe_price_id: str) -> Optional[UsageCreditPackage]:
        result = await self._db.execute(
            select(UsageCreditPackage).where(UsageCreditPackage.stripe_price_id == stripe_price_id)
        )
        return result.scalar_one_or_none()

    async def list_all(self, active_only: bool = False) -> list[UsageCreditPackage]:
        q = select(UsageCreditPackage).order_by(UsageCreditPackage.sort_order)
        if active_only:
            q = q.where(UsageCreditPackage.is_active.is_(True))
        result = await self._db.execute(q)
        return list(result.scalars().all())

    async def create(self, **kwargs) -> UsageCreditPackage:
        pkg = UsageCreditPackage(**kwargs)
        self._db.add(pkg)
        await self._db.flush()
        await self._db.refresh(pkg)
        return pkg

    async def update(self, package_id: UUID, **kwargs) -> Optional[UsageCreditPackage]:
        pkg = await self.get_by_id(package_id)
        if not pkg:
            return None
        for k, v in kwargs.items():
            if v is not None and hasattr(pkg, k):
                setattr(pkg, k, v)
        await self._db.flush()
        await self._db.refresh(pkg)
        return pkg


class UsageCreditPurchaseRepository:
    def __init__(self, db: AsyncSession):
        self._db = db

    async def get_by_id(self, purchase_id: UUID) -> Optional[UsageCreditPurchase]:
        result = await self._db.execute(
            select(UsageCreditPurchase).where(UsageCreditPurchase.id == purchase_id)
        )
        return result.scalar_one_or_none()

    async def get_by_stripe_session_id(self, session_id: str) -> Optional[UsageCreditPurchase]:
        result = await self._db.execute(
            select(UsageCreditPurchase).where(
                UsageCreditPurchase.stripe_checkout_session_id == session_id
            )
        )
        return result.scalar_one_or_none()

    async def list_for_workspace(self, workspace_id: UUID) -> list[UsageCreditPurchase]:
        result = await self._db.execute(
            select(UsageCreditPurchase)
            .where(UsageCreditPurchase.workspace_id == workspace_id)
            .order_by(UsageCreditPurchase.created_at.desc())
        )
        return list(result.scalars().all())

    async def create(self, **kwargs) -> UsageCreditPurchase:
        purchase = UsageCreditPurchase(**kwargs)
        self._db.add(purchase)
        await self._db.flush()
        await self._db.refresh(purchase)
        return purchase

    async def update_status(
        self, purchase_id: UUID, status: str, paid_at: Optional[datetime] = None
    ) -> Optional[UsageCreditPurchase]:
        purchase = await self.get_by_id(purchase_id)
        if not purchase:
            return None
        purchase.status = status
        if paid_at is not None:
            purchase.paid_at = paid_at
        await self._db.flush()
        await self._db.refresh(purchase)
        return purchase


class WorkspaceCreditLedgerRepository:
    def __init__(self, db: AsyncSession):
        self._db = db

    async def add_entry(
        self,
        workspace_id: UUID,
        entry_type: str,
        units_delta: int,
        *,
        balance_after: Optional[int] = None,
        source_bucket: Optional[str] = None,
        source_type: Optional[str] = None,
        source_id: Optional[str] = None,
        subscription_id: Optional[UUID] = None,
        usage_credit_purchase_id: Optional[UUID] = None,
        project_id: Optional[UUID] = None,
        codegen_job_id: Optional[UUID] = None,
        artifact_id: Optional[UUID] = None,
        revision_id: Optional[UUID] = None,
        description: Optional[str] = None,
        metadata: Optional[dict] = None,
        allow_zero: bool = False,
    ) -> WorkspaceCreditLedger:
        """Add ledger entry. Zero-unit entries rejected unless allow_zero=True.

        balance_after: Optional snapshot of purchased_topup balance after this entry.
        Used for purchased_topup consumption entries; total = allowance + purchased.
        """
        if units_delta == 0 and not allow_zero:
            raise ValueError("units_delta must be non-zero unless allow_zero=True")
        if units_delta < 0 and entry_type in ("purchase", "bonus"):
            raise ValueError(
                f"Negative units_delta invalid for entry_type={entry_type}; "
                "use consumption/refund/expiration for negative deltas"
            )
        entry = WorkspaceCreditLedger(
            workspace_id=workspace_id,
            entry_type=entry_type,
            units_delta=units_delta,
            balance_after=balance_after,
            source_bucket=source_bucket,
            source_type=source_type,
            source_id=source_id,
            subscription_id=subscription_id,
            usage_credit_purchase_id=usage_credit_purchase_id,
            project_id=project_id,
            codegen_job_id=codegen_job_id,
            artifact_id=artifact_id,
            revision_id=revision_id,
            description=description,
            metadata_json=metadata,
        )
        self._db.add(entry)
        await self._db.flush()
        return entry

    async def list_by_source(
        self,
        workspace_id: UUID,
        *,
        source_type: str,
        source_id: str,
    ) -> list[WorkspaceCreditLedger]:
        result = await self._db.execute(
            select(WorkspaceCreditLedger)
            .where(WorkspaceCreditLedger.workspace_id == workspace_id)
            .where(WorkspaceCreditLedger.source_type == source_type)
            .where(WorkspaceCreditLedger.source_id == source_id)
            .order_by(WorkspaceCreditLedger.created_at, WorkspaceCreditLedger.id)
        )
        return list(result.scalars().all())

    async def sum_balance(self, workspace_id: UUID) -> int:
        """Sum all ledger deltas for workspace (legacy; prefer sum_purchased_balance for purchased units)."""
        result = await self._db.execute(
            select(sa_func.coalesce(sa_func.sum(WorkspaceCreditLedger.units_delta), 0)).where(
                WorkspaceCreditLedger.workspace_id == workspace_id
            )
        )
        return int(result.scalar() or 0)

    async def acquire_consumption_lock(self, workspace_id: UUID) -> None:
        """
        Acquire transaction-scoped advisory lock for workspace consumption.
        Serializes concurrent consume_units for the same workspace.
        Call at start of consume_units; lock released when transaction commits/rolls back.
        """
        # Serialize consumption per workspace; lock released when transaction ends
        await self._db.execute(
            text("SELECT pg_advisory_xact_lock(hashtext(:ws)::bigint)"),
            {"ws": str(workspace_id)},
        )

    async def sum_purchased_balance(self, workspace_id: UUID) -> int:
        """
        Sum ledger deltas that contribute to purchased top-up balance only.
        Excludes source_bucket='subscription_allowance' (those are allowance consumption).
        """
        result = await self._db.execute(
            select(sa_func.coalesce(sa_func.sum(WorkspaceCreditLedger.units_delta), 0)).where(
                WorkspaceCreditLedger.workspace_id == workspace_id
            ).where(
                (WorkspaceCreditLedger.source_bucket.is_(None))
                | (WorkspaceCreditLedger.source_bucket != "subscription_allowance")
            )
        )
        return int(result.scalar() or 0)

    async def list_recent(
        self, workspace_id: UUID, limit: int = 50
    ) -> list[WorkspaceCreditLedger]:
        result = await self._db.execute(
            select(WorkspaceCreditLedger)
            .where(WorkspaceCreditLedger.workspace_id == workspace_id)
            .order_by(WorkspaceCreditLedger.created_at.desc())
            .limit(limit)
        )
        return list(result.scalars().all())

    async def sum_purchased_topup_units(self, workspace_id: UUID) -> int:
        """Total units ever purchased via top-up (sum of purchase entry deltas)."""
        result = await self._db.execute(
            select(sa_func.coalesce(sa_func.sum(WorkspaceCreditLedger.units_delta), 0)).where(
                WorkspaceCreditLedger.workspace_id == workspace_id
            ).where(WorkspaceCreditLedger.entry_type == "purchase")
        )
        return int(result.scalar() or 0)

    async def sum_consumed_topup_units(self, workspace_id: UUID) -> int:
        """Total units consumed from purchased top-up (absolute sum of negative deltas)."""
        result = await self._db.execute(
            select(sa_func.coalesce(sa_func.sum(sa_func.abs(WorkspaceCreditLedger.units_delta)), 0)).where(
                WorkspaceCreditLedger.workspace_id == workspace_id
            ).where(WorkspaceCreditLedger.entry_type == "consumption")
            .where(WorkspaceCreditLedger.source_bucket == "purchased_topup")
        )
        return int(result.scalar() or 0)

    async def sum_allowance_consumed(self, workspace_id: UUID) -> int:
        """Total units consumed from subscription allowance."""
        result = await self._db.execute(
            select(sa_func.coalesce(sa_func.sum(sa_func.abs(WorkspaceCreditLedger.units_delta)), 0)).where(
                WorkspaceCreditLedger.workspace_id == workspace_id
            ).where(WorkspaceCreditLedger.source_bucket == "subscription_allowance")
        )
        return int(result.scalar() or 0)


class WorkspaceUsageAllowanceRepository:
    def __init__(self, db: AsyncSession):
        self._db = db

    async def reconcile_plan_limit(self, workspace_id: UUID, included_units: int) -> None:
        """Preserve usage and the allowance window when a subscription changes plan.

        Caller holds the workspace consumption lock in the surrounding transaction.
        Repeated subscription notifications apply the same absolute limit.
        """
        allowance = await self.get_current_for_workspace_for_update(workspace_id)
        if allowance:
            allowance.included_units = max(0, included_units)
            allowance.remaining_units = max(
                0, allowance.included_units - allowance.consumed_units
            )
            await self._db.flush()

    async def get_current_for_workspace(
        self, workspace_id: UUID
    ) -> Optional[WorkspaceUsageAllowance]:
        now = datetime.now(timezone.utc)
        result = await self._db.execute(
            select(WorkspaceUsageAllowance)
            .where(WorkspaceUsageAllowance.workspace_id == workspace_id)
            .where(WorkspaceUsageAllowance.period_end >= now)
            .where(WorkspaceUsageAllowance.period_start <= now)
            .order_by(WorkspaceUsageAllowance.period_end.desc())
            .limit(1)
        )
        return result.scalar_one_or_none()

    async def get_current_for_workspace_for_update(
        self, workspace_id: UUID
    ) -> Optional[WorkspaceUsageAllowance]:
        """Get current allowance with row lock for consumption (must be in transaction)."""
        now = datetime.now(timezone.utc)
        result = await self._db.execute(
            select(WorkspaceUsageAllowance)
            .where(WorkspaceUsageAllowance.workspace_id == workspace_id)
            .where(WorkspaceUsageAllowance.period_end >= now)
            .where(WorkspaceUsageAllowance.period_start <= now)
            .order_by(WorkspaceUsageAllowance.period_end.desc())
            .limit(1)
            .with_for_update()
        )
        return result.scalar_one_or_none()

    async def create(
        self,
        workspace_id: UUID,
        period_start: datetime,
        period_end: datetime,
        included_units: int,
        subscription_id: Optional[UUID] = None,
    ) -> WorkspaceUsageAllowance:
        allowance = WorkspaceUsageAllowance(
            workspace_id=workspace_id,
            subscription_id=subscription_id,
            period_start=period_start,
            period_end=period_end,
            included_units=included_units,
            consumed_units=0,
            remaining_units=included_units,
        )
        self._db.add(allowance)
        await self._db.flush()
        await self._db.refresh(allowance)
        return allowance

    async def consume_units(
        self, allowance_id: UUID, units: int
    ) -> Optional[WorkspaceUsageAllowance]:
        result = await self._db.execute(
            select(WorkspaceUsageAllowance).where(WorkspaceUsageAllowance.id == allowance_id)
        )
        allowance = result.scalar_one_or_none()
        if not allowance or allowance.remaining_units < units:
            return None
        allowance.consumed_units += units
        allowance.remaining_units -= units
        await self._db.flush()
        await self._db.refresh(allowance)
        return allowance

    async def restore_units(
        self, allowance_id: UUID, units: int
    ) -> Optional[WorkspaceUsageAllowance]:
        if units <= 0:
            raise ValueError("Restored units must be positive")
        result = await self._db.execute(
            select(WorkspaceUsageAllowance)
            .where(WorkspaceUsageAllowance.id == allowance_id)
            .with_for_update()
        )
        allowance = result.scalar_one_or_none()
        if not allowance:
            return None
        allowance.consumed_units = max(0, allowance.consumed_units - units)
        allowance.remaining_units = max(0, allowance.included_units - allowance.consumed_units)
        await self._db.flush()
        await self._db.refresh(allowance)
        return allowance


class EnterpriseInquiryRepository:
    def __init__(self, db: AsyncSession):
        self._db = db

    async def create(
        self,
        *,
        workspace_id: Optional[UUID],
        user_id: Optional[UUID],
        first_name: str,
        last_name: str,
        email: str,
        phone: str,
        website_count: int,
        ai_credit_minutes_desired: int,
        is_agency_account: bool,
        employee_count: int,
        additional_features: Optional[str],
        plan_slug: Optional[str],
    ) -> EnterpriseInquiry:
        row = EnterpriseInquiry(
            workspace_id=workspace_id,
            user_id=user_id,
            first_name=first_name,
            last_name=last_name,
            email=email,
            phone=phone,
            website_count=website_count,
            ai_credit_minutes_desired=ai_credit_minutes_desired,
            is_agency_account=is_agency_account,
            employee_count=employee_count,
            additional_features=additional_features,
            plan_slug=plan_slug,
            status="new",
        )
        self._db.add(row)
        await self._db.flush()
        await self._db.refresh(row)
        return row