"""
Enrichment usage tracking, quota enforcement, and idempotency handling.

This service owns the Prevent → Meter → Bill lifecycle:
    * Prevent: quota checks, reservations, spend caps
    * Meter: atomic debit from allowance/packs, append-only ledger writes
    * Bill: expose usage summaries for downstream invoicing jobs
"""

from __future__ import annotations

import logging
from datetime import datetime, timedelta, timezone
from decimal import Decimal, ROUND_HALF_UP
from typing import Any, Dict, Optional, Tuple

from fastapi import HTTPException
from sqlalchemy import delete, func, select
from sqlalchemy.ext.asyncio import AsyncSession
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError

from ..config.enrichment import (
    HEARTBEAT_GRACE_SECONDS,
    IDEMPOTENCY_TTL_HOURS,
    MAX_ACCOUNT_DAILY_SPEND_USD,
    MAX_GLOBAL_DAILY_SPEND_USD,
    RESERVATION_TTL_SECONDS,
)
from ..models import (
    EnrichmentQuota,
    EnrichmentLedgerOutbox,
    EnrichmentReservation,
    EnrichmentReservationStatus,
    EnrichmentUsage,
    EnrichmentUsageIdempotency,
)
from .enrichment_bucket_manager import consume_credits_fefo, InsufficientCreditsError

logger = logging.getLogger(__name__)

LEASE_DEFAULT_SECONDS = RESERVATION_TTL_SECONDS
DECIMAL_TWO = Decimal("0.01")
DECIMAL_THREE = Decimal("0.001")


def _now_utc() -> datetime:
    return datetime.now(timezone.utc)


def _round_money(value: Decimal) -> Decimal:
    return value.quantize(DECIMAL_TWO, rounding=ROUND_HALF_UP)


def _round_duration(value: Decimal) -> Decimal:
    return value.quantize(DECIMAL_THREE, rounding=ROUND_HALF_UP)


def _coerce_decimal(value: Optional[Decimal]) -> Decimal:
    if value is None:
        return Decimal("0")
    if isinstance(value, Decimal):
        return value
    return Decimal(str(value))


def _get_zone(name: Optional[str]) -> ZoneInfo:
    if not name:
        return ZoneInfo("UTC")
    try:
        return ZoneInfo(name)
    except ZoneInfoNotFoundError:
        logger.warning(
            "Unknown timezone '%s'; defaulting to UTC for enrichment quotas.", name
        )
        return ZoneInfo("UTC")


class EnrichmentUsageTracker:
    """Coordinate quota checks, reservations, ledger writes, and billing summaries."""

    def __init__(self, db: AsyncSession):
        self.db = db

    # -------------------------------------------------------------------------
    # Public API
    # -------------------------------------------------------------------------

    def resolve_account(self, request: Any) -> Tuple[str, Optional[str]]:
        """
        Extract (account_id, user_email) from a FastAPI request/Clerk session.
        Raises HTTP 401 when account context is missing.
        """
        state = getattr(request, "state", None)
        account_id = None
        user_email: Optional[str] = None

        if state is not None:
            account_id = getattr(state, "account_id", None) or getattr(
                state, "org_id", None
            )
            user_email = getattr(state, "user_email", None) or getattr(
                state, "email", None
            )
            identity = getattr(state, "identity", None)
            if isinstance(identity, dict):
                account_id = (
                    account_id or identity.get("org_id") or identity.get("account_id")
                )
                user_email = user_email or identity.get("email")

        if account_id is None and hasattr(request, "headers"):
            account_id = request.headers.get("X-Account-Id")

        if not account_id:
            raise HTTPException(
                status_code=401,
                detail="Missing account context for enrichment request.",
            )

        return str(account_id), user_email

    async def preflight(
        self,
        account_id: str,
        estimated_requests: int,
        *,
        idempotency_key: str,
        request_hash: str,
        reserve: bool = True,
        job_id: Optional[str] = None,
        chunk_id: Optional[str] = None,
        lease_seconds: int = LEASE_DEFAULT_SECONDS,
        buffer_ratio: float = 0.0,
    ) -> Dict[str, Any]:
        """
        Check an account's quota/balances before calling the provider.
        Optionally reserves credits atomically when `reserve=True`.
        """
        if estimated_requests < 0:
            raise ValueError("estimated_requests cannot be negative")

        now_utc = _now_utc()
        buffer_multiplier = Decimal("1")
        if buffer_ratio:
            buffer_multiplier += Decimal(str(buffer_ratio))
        buffered_requests = int(
            (Decimal(estimated_requests) * buffer_multiplier).to_integral_value(
                rounding=ROUND_HALF_UP
            )
        )
        buffered_requests = max(buffered_requests, estimated_requests)

        async with self.db.begin():
            idempotency = await self._get_idempotency(
                account_id, idempotency_key, for_update=True
            )
            if (
                idempotency
                and idempotency.request_hash
                and idempotency.request_hash != request_hash
            ):
                raise HTTPException(
                    status_code=409,
                    detail="Idempotency-Key reuse detected with a different payload. Abort to prevent double spend.",
                )

            if idempotency and idempotency.first_usage_id:
                return {
                    "allowed": False,
                    "replay": True,
                    "usage_id": idempotency.first_usage_id,
                    "reason": "Chunk already finalized.",
                }

            quota = await self._get_quota(account_id, for_update=True)
            (
                day_start,
                next_day_start,
                month_start,
                next_month_start,
            ) = self._reset_quota_windows(quota, now_utc)

            credits_per_request = quota.credits_per_request or 0
            if credits_per_request <= 0:
                return self._build_denied_response(
                    quota=quota,
                    reason="Account credits-per-request not configured.",
                    required_requests=buffered_requests,
                    required_credits=0,
                    day_start=day_start,
                    next_day_start=next_day_start,
                )

            required_credits = buffered_requests * credits_per_request
            allowance_remaining = quota.allowance_balance_credits
            pack_remaining = quota.pack_balance_credits
            total_balance = allowance_remaining + pack_remaining
            daily_remaining = self._remaining_limit(
                quota.daily_limit_credits, quota.daily_used_credits
            )
            monthly_remaining = self._remaining_limit(
                quota.monthly_limit_credits, quota.monthly_used_credits
            )

            allowed = True
            reason: Optional[str] = None
            denial_status = 429

            if not quota.enabled:
                allowed = False
                reason = "Enrichment disabled for this account. Contact support to re-enable."
                denial_status = 403
            elif total_balance < required_credits:
                allowed = False
                reason = "Insufficient enrichment credits. Purchase a pack or upgrade your plan."
                denial_status = 402
            elif daily_remaining is not None and required_credits > daily_remaining:
                allowed = False
                reason = "Daily enrichment quota exceeded."
            elif monthly_remaining is not None and required_credits > monthly_remaining:
                allowed = False
                reason = "Monthly enrichment quota exceeded."

            provider_unit_cost = _coerce_decimal(quota.provider_cost_per_request_usd)
            projected_cost = _round_money(
                provider_unit_cost * Decimal(buffered_requests)
            )

            if allowed and MAX_ACCOUNT_DAILY_SPEND_USD > 0:
                account_spend = await self._sum_cost(
                    account_id, day_start, next_day_start
                )
                if account_spend + projected_cost > MAX_ACCOUNT_DAILY_SPEND_USD:
                    allowed = False
                    reason = "Daily spend cap reached for this account."
                    denial_status = 403

            if allowed and MAX_GLOBAL_DAILY_SPEND_USD > 0:
                global_spend = await self._sum_cost(None, day_start, next_day_start)
                if global_spend + projected_cost > MAX_GLOBAL_DAILY_SPEND_USD:
                    allowed = False
                    reason = (
                        "Daily global enrichment spend cap reached. Try again later."
                    )
                    denial_status = 429

            if not allowed:
                return self._build_denied_response(
                    quota=quota,
                    reason=reason or "Quota check failed.",
                    required_requests=buffered_requests,
                    required_credits=required_credits,
                    day_start=day_start,
                    next_day_start=next_day_start,
                    allowance_remaining=allowance_remaining,
                    pack_remaining=pack_remaining,
                    daily_remaining=daily_remaining,
                    monthly_remaining=monthly_remaining,
                    status_code=denial_status,
                )

            if not reserve:
                return self._build_allowed_response(
                    quota=quota,
                    reservation=None,
                    required_requests=buffered_requests,
                    required_credits=required_credits,
                    allowance_remaining=allowance_remaining,
                    pack_remaining=pack_remaining,
                    daily_remaining=daily_remaining,
                    monthly_remaining=monthly_remaining,
                )

            reservation = await self._get_reservation(
                account_id, idempotency_key, for_update=True
            )
            lease_expires_at = now_utc + timedelta(seconds=lease_seconds)

            if reservation is None or reservation.status in {
                EnrichmentReservationStatus.RELEASED,
                EnrichmentReservationStatus.EXPIRED,
            }:
                reservation = EnrichmentReservation(
                    account_id=account_id,
                    idempotency_key=idempotency_key,
                    job_id=job_id,
                    chunk_id=chunk_id,
                    required_credits=required_credits,
                    status=EnrichmentReservationStatus.RESERVED,
                    lease_expires_at=lease_expires_at,
                    last_heartbeat_at=now_utc,
                )
                self.db.add(reservation)
                await self.db.flush()
            else:
                reservation.required_credits = required_credits
                reservation.status = EnrichmentReservationStatus.RESERVED
                reservation.job_id = job_id
                reservation.chunk_id = chunk_id
                reservation.lease_expires_at = lease_expires_at
                reservation.last_heartbeat_at = now_utc

            if idempotency is None:
                idempotency = EnrichmentUsageIdempotency(
                    account_id=account_id,
                    idempotency_key=idempotency_key,
                    request_hash=request_hash,
                )
                self.db.add(idempotency)
            else:
                idempotency.request_hash = request_hash

            return self._build_allowed_response(
                quota=quota,
                reservation=reservation,
                required_requests=buffered_requests,
                required_credits=required_credits,
                allowance_remaining=allowance_remaining,
                pack_remaining=pack_remaining,
                daily_remaining=daily_remaining,
                monthly_remaining=monthly_remaining,
            )

    async def heartbeat(
        self,
        account_id: str,
        idempotency_key: str,
        *,
        lease_seconds: int = LEASE_DEFAULT_SECONDS,
    ) -> bool:
        """Extend a live reservation lease to keep long-running chunks alive."""
        now_utc = _now_utc()
        async with self.db.begin():
            reservation = await self._get_reservation(
                account_id, idempotency_key, for_update=True
            )
            if (
                not reservation
                or reservation.status != EnrichmentReservationStatus.RESERVED
            ):
                return False

            reservation.last_heartbeat_at = now_utc
            reservation.lease_expires_at = now_utc + timedelta(seconds=lease_seconds)
            return True

    async def release(self, account_id: str, idempotency_key: str) -> bool:
        """Release a reservation before any provider spend occurs."""
        async with self.db.begin():
            reservation = await self._get_reservation(
                account_id, idempotency_key, for_update=True
            )
            if (
                not reservation
                or reservation.status != EnrichmentReservationStatus.RESERVED
            ):
                return False

            reservation.status = EnrichmentReservationStatus.RELEASED
            reservation.finalized_at = _now_utc()
            reservation.lease_expires_at = None
            reservation.last_heartbeat_at = None
            return True

    async def sweep_timeouts(self, *, now: Optional[datetime] = None) -> Dict[str, Any]:
        """
        Expire abandoned reservations and purge stale idempotency keys.

        Intended to be invoked from a periodic job (e.g., every 5 minutes).
        """
        sweep_started_at = now or _now_utc()
        grace_cutoff = sweep_started_at - timedelta(seconds=HEARTBEAT_GRACE_SECONDS)

        expired_count = 0
        total_age_seconds = 0.0

        async with self.db.begin():
            stmt = (
                select(EnrichmentReservation)
                .where(
                    EnrichmentReservation.status
                    == EnrichmentReservationStatus.RESERVED,
                    EnrichmentReservation.lease_expires_at.is_not(None),
                    EnrichmentReservation.lease_expires_at < grace_cutoff,
                )
                .with_for_update(skip_locked=True)
            )
            result = await self.db.execute(stmt)
            to_expire = result.scalars().all()

            for reservation in to_expire:
                reservation.status = EnrichmentReservationStatus.EXPIRED
                reservation.finalized_at = None
                expired_count += 1
                if reservation.created_at:
                    age = (sweep_started_at - reservation.created_at).total_seconds()
                    total_age_seconds += max(age, 0.0)

        avg_age_seconds = total_age_seconds / expired_count if expired_count else 0.0

        idempotency_cutoff = sweep_started_at - timedelta(hours=IDEMPOTENCY_TTL_HOURS)
        async with self.db.begin():
            delete_stmt = delete(EnrichmentUsageIdempotency).where(
                EnrichmentUsageIdempotency.created_at < idempotency_cutoff
            )
            delete_result = await self.db.execute(delete_stmt)

        rowcount = delete_result.rowcount or 0
        purged_idempotency = rowcount if rowcount >= 0 else 0

        summary = {
            "expired_reservations": expired_count,
            "avg_reservation_age_seconds": round(avg_age_seconds, 3),
            "idempotency_purged": purged_idempotency,
            "sweep_timestamp": sweep_started_at.isoformat(),
        }

        logger.info(
            "Enrichment reservation sweep completed: expired=%s avg_age_s=%.3f purged_idempotency=%s",
            expired_count,
            avg_age_seconds,
            purged_idempotency,
        )
        return summary

    async def finalize(
        self,
        *,
        account_id: str,
        user_email: Optional[str],
        idempotency_key: str,
        request_hash: str,
        actual_requests: int,
        matches_found: int,
        operation: str,
        enrichment_type: str,
        provider: str,
        job_id: Optional[str],
        chunk_id: Optional[str],
        meta: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """
        Debit credits, append to the ledger, and finalize the reservation.
        """
        if actual_requests < 0:
            raise ValueError("actual_requests cannot be negative")

        now_utc = _now_utc()
        meta = meta or {}

        async with self.db.begin():
            quota = await self._get_quota(account_id, for_update=True)
            reservation = await self._get_reservation(
                account_id, idempotency_key, for_update=True
            )
            idempotency = await self._get_idempotency(
                account_id, idempotency_key, for_update=True
            )

            if reservation is None:
                raise HTTPException(
                    status_code=404, detail="No reservation found to finalize."
                )
            if reservation.status == EnrichmentReservationStatus.FINALIZED:
                if idempotency and idempotency.first_usage_id:
                    return {
                        "replay": True,
                        "usage_id": idempotency.first_usage_id,
                    }
                raise HTTPException(
                    status_code=409,
                    detail="Reservation already finalized without usage record.",
                )
            if reservation.status != EnrichmentReservationStatus.RESERVED:
                raise HTTPException(
                    status_code=409, detail="Reservation is not in a reservable state."
                )

            (
                day_start,
                _next_day_start,
                _month_start,
                _next_month_start,
            ) = self._reset_quota_windows(quota, now_utc)

            credits_per_request = quota.credits_per_request or 0
            if credits_per_request <= 0:
                raise HTTPException(
                    status_code=409,
                    detail="Account credits-per-request not configured.",
                )

            credits_needed = actual_requests * credits_per_request
            if credits_needed > reservation.required_credits:
                shortfall = credits_needed - reservation.required_credits
                available = quota.allowance_balance_credits + quota.pack_balance_credits
                if shortfall > available:
                    raise HTTPException(
                        status_code=402,
                        detail="Insufficient credits to cover provider calls. Reservation under-estimated usage.",
                    )

            # 🔄 NEW: Use FEFO (First-Expiring-First-Out) bucket consumption
            try:
                consumed_buckets = await consume_credits_fefo(
                    account_id=account_id,
                    quantity_needed=credits_needed,
                    job_id=job_id or "unknown",
                    user_email=user_email or "system",
                    db=self.db,
                )
                logger.info(
                    "[FEFO] Consumed %s credits from %s bucket(s) for account %s (job=%s)",
                    credits_needed,
                    len(consumed_buckets),
                    account_id,
                    job_id,
                )
            except InsufficientCreditsError as e:
                logger.error(
                    "[FEFO] Insufficient credits for account %s: %s",
                    account_id,
                    str(e),
                )
                raise HTTPException(
                    status_code=402,
                    detail=f"Insufficient enrichment credits to finalize usage: {e}",
                )

            # Update legacy quota balances for backward compatibility with existing reporting
            # Note: The source of truth is now the bucket system, but we maintain these for analytics
            allowance_debit = min(quota.allowance_balance_credits, credits_needed)
            quota.allowance_balance_credits -= allowance_debit
            remaining_to_debit = credits_needed - allowance_debit
            pack_debit = min(quota.pack_balance_credits, remaining_to_debit)
            quota.pack_balance_credits -= pack_debit

            quota.daily_used_credits += credits_needed
            quota.monthly_used_credits += credits_needed

            provider_unit_cost = _coerce_decimal(quota.provider_cost_per_request_usd)
            price_unit = _coerce_decimal(quota.price_per_request_usd)
            cost_usd = _round_money(provider_unit_cost * Decimal(actual_requests))
            price_usd = _round_money(price_unit * Decimal(actual_requests))
            margin_usd = _round_money(price_usd - cost_usd)

            usage = EnrichmentUsage(
                account_id=account_id,
                user_email=user_email,
                operation=operation,
                provider=provider,
                enrichment_type=enrichment_type,
                sfdc_object=meta.get("sfdc_object"),
                job_id=job_id,
                chunk_id=chunk_id,
                idempotency_key=idempotency_key,
                request_hash=request_hash,
                records_scanned=int(meta.get("records_scanned", 0) or 0),
                candidates_found=int(meta.get("candidates_found", 0) or 0),
                pdl_requests=actual_requests,
                matches_found=matches_found,
                credits_debited=credits_needed,
                records_written=int(meta.get("records_written", 0) or 0),
                fields_updated=int(meta.get("fields_updated", 0) or 0),
                cost_usd=cost_usd,
                price_usd=price_usd,
                margin_usd=margin_usd,
                duration_ms=self._normalize_duration(meta.get("duration_ms")),
                error=meta.get("error"),
            )
            self.db.add(usage)
            await self.db.flush()

            reservation.status = EnrichmentReservationStatus.FINALIZED
            reservation.finalized_at = now_utc
            reservation.lease_expires_at = None
            reservation.last_heartbeat_at = None

            if idempotency is None:
                idempotency = EnrichmentUsageIdempotency(
                    account_id=account_id,
                    idempotency_key=idempotency_key,
                    request_hash=request_hash,
                    first_usage_id=usage.id,
                )
                self.db.add(idempotency)
            else:
                idempotency.request_hash = request_hash
                if idempotency.first_usage_id is None:
                    idempotency.first_usage_id = usage.id

            return {
                "usage_id": usage.id,
                "credits_debited": credits_needed,
                "allowance_remaining": quota.allowance_balance_credits,
                "pack_remaining": quota.pack_balance_credits,
                "cost_usd": float(cost_usd),
                "price_usd": float(price_usd),
                "margin_usd": float(margin_usd),
            }

    async def enqueue_ledger_outbox(
        self,
        *,
        account_id: str,
        user_email: Optional[str],
        idempotency_key: str,
        request_hash: str,
        actual_requests: int,
        matches_found: int,
        operation: str,
        enrichment_type: str,
        provider: str,
        job_id: Optional[str],
        chunk_id: Optional[str],
        meta: Optional[Dict[str, Any]] = None,
    ) -> Optional[int]:
        """
        Persist a retryable outbox record so finalize can be replayed asynchronously.
        """
        payload = {
            "account_id": account_id,
            "user_email": user_email,
            "idempotency_key": idempotency_key,
            "request_hash": request_hash,
            "actual_requests": actual_requests,
            "matches_found": matches_found,
            "operation": operation,
            "enrichment_type": enrichment_type,
            "provider": provider,
            "job_id": job_id,
            "chunk_id": chunk_id,
            "meta": meta or {},
        }

        now_utc = _now_utc()

        async with self.db.begin():
            reservation = await self._get_reservation(
                account_id, idempotency_key, for_update=True
            )
            if reservation is None:
                logger.error(
                    "Failed to enqueue ledger outbox: reservation missing (account=%s, idempotency=%s).",
                    account_id,
                    idempotency_key,
                )
                return None

            outbox = EnrichmentLedgerOutbox(
                reservation_id=reservation.id,
                account_id=account_id,
                idempotency_key=idempotency_key,
                request_hash=request_hash,
                operation=operation,
                enrichment_type=enrichment_type,
                provider=provider,
                payload=payload,
            )
            self.db.add(outbox)

            # Keep the reservation alive while the outbox is retried.
            reservation.last_heartbeat_at = now_utc
            reservation.lease_expires_at = now_utc + timedelta(
                seconds=LEASE_DEFAULT_SECONDS
            )

            await self.db.flush()
            outbox_id = outbox.id

        logger.warning(
            "Queued enrichment ledger outbox item %s (account=%s, reservation=%s) after finalize failure.",
            outbox_id,
            account_id,
            reservation.id,
        )
        return outbox_id

    async def balances(self, account_id: str) -> Dict[str, Any]:
        """Return a snapshot of wallet balances and quota limits."""
        quota = await self._get_quota(account_id, for_update=False)
        now_utc = _now_utc()
        day_start, next_day_start, _, _ = self._day_month_bounds(quota, now_utc)
        return self._build_allowed_response(
            quota=quota,
            reservation=None,
            required_requests=0,
            required_credits=0,
            allowance_remaining=quota.allowance_balance_credits,
            pack_remaining=quota.pack_balance_credits,
            daily_remaining=self._remaining_limit(
                quota.daily_limit_credits, quota.daily_used_credits
            ),
            monthly_remaining=self._remaining_limit(
                quota.monthly_limit_credits, quota.monthly_used_credits
            ),
            day_start=day_start,
            next_day_start=next_day_start,
        )

    async def get_usage_summary(
        self, account_id: str, days: int = 30
    ) -> Dict[str, Any]:
        """Aggregate usage for analytics and billing."""
        since = _now_utc() - timedelta(days=days)
        stmt = select(
            func.coalesce(func.sum(EnrichmentUsage.credits_debited), 0),
            func.coalesce(func.sum(EnrichmentUsage.cost_usd), Decimal("0")),
            func.coalesce(func.sum(EnrichmentUsage.price_usd), Decimal("0")),
            func.coalesce(func.sum(EnrichmentUsage.margin_usd), Decimal("0")),
            func.coalesce(func.sum(EnrichmentUsage.records_written), 0),
            func.coalesce(func.sum(EnrichmentUsage.fields_updated), 0),
        ).where(
            EnrichmentUsage.account_id == account_id,
            EnrichmentUsage.created_at >= since,
        )

        result = await self.db.execute(stmt)
        row = result.first()
        total_credits = int(row[0] or 0)
        total_cost = _coerce_decimal(row[1])
        total_price = _coerce_decimal(row[2])
        total_margin = _coerce_decimal(row[3])
        records_enriched = int(row[4] or 0)
        fields_filled = int(row[5] or 0)

        unbilled_stmt = select(
            func.coalesce(func.sum(EnrichmentUsage.price_usd), Decimal("0"))
        ).where(
            EnrichmentUsage.account_id == account_id,
            EnrichmentUsage.billed.is_(False),
        )
        unbilled_result = await self.db.execute(unbilled_stmt)
        unbilled_amount = _coerce_decimal(unbilled_result.scalar())

        return {
            "total_credits": total_credits,
            "total_cost_usd": float(_round_money(total_cost)),
            "total_price_usd": float(_round_money(total_price)),
            "total_margin_usd": float(_round_money(total_margin)),
            "records_enriched": records_enriched,
            "fields_filled": fields_filled,
            "unbilled_amount_usd": float(_round_money(unbilled_amount)),
        }

    async def mark_as_billed(self, account_id: str, stripe_invoice_id: str) -> int:
        """Mark all unbilled ledger rows as billed with a Stripe invoice ID."""
        stmt = select(EnrichmentUsage).where(
            EnrichmentUsage.account_id == account_id,
            EnrichmentUsage.billed.is_(False),
        )
        result = await self.db.execute(stmt)
        records = result.scalars().all()

        now = _now_utc()
        for record in records:
            record.billed = True
            record.billed_at = now
            record.stripe_invoice_id = stripe_invoice_id

        if records:
            await self.db.commit()

        logger.info(
            "Marked %s enrichment ledger rows as billed for account %s (invoice %s)",
            len(records),
            account_id,
            stripe_invoice_id,
        )
        return len(records)

    # -------------------------------------------------------------------------
    # Internal helpers
    # -------------------------------------------------------------------------

    async def _get_quota(self, account_id: str, *, for_update: bool) -> EnrichmentQuota:
        stmt = select(EnrichmentQuota).where(EnrichmentQuota.account_id == account_id)
        if for_update:
            stmt = stmt.with_for_update()
        result = await self.db.execute(stmt)
        quota = result.scalar_one_or_none()

        if quota is None:
            quota = EnrichmentQuota(
                account_id=account_id,
                enabled=False,
            )
            self.db.add(quota)
            await self.db.flush()

        return quota

    async def _get_reservation(
        self,
        account_id: str,
        idempotency_key: str,
        *,
        for_update: bool,
    ) -> Optional[EnrichmentReservation]:
        stmt = select(EnrichmentReservation).where(
            EnrichmentReservation.account_id == account_id,
            EnrichmentReservation.idempotency_key == idempotency_key,
        )
        if for_update:
            stmt = stmt.with_for_update()
        result = await self.db.execute(stmt)
        return result.scalar_one_or_none()

    async def _get_idempotency(
        self,
        account_id: str,
        idempotency_key: str,
        *,
        for_update: bool,
    ) -> Optional[EnrichmentUsageIdempotency]:
        stmt = select(EnrichmentUsageIdempotency).where(
            EnrichmentUsageIdempotency.account_id == account_id,
            EnrichmentUsageIdempotency.idempotency_key == idempotency_key,
        )
        if for_update:
            stmt = stmt.with_for_update()
        result = await self.db.execute(stmt)
        return result.scalar_one_or_none()

    def _reset_quota_windows(
        self,
        quota: EnrichmentQuota,
        now_utc: datetime,
    ) -> Tuple[datetime, datetime, datetime, datetime]:
        day_start, next_day_start, month_start, next_month_start = (
            self._day_month_bounds(quota, now_utc)
        )

        if quota.daily_reset_at is None or now_utc >= quota.daily_reset_at:
            quota.daily_used_credits = 0
            quota.daily_reset_at = next_day_start

        if quota.monthly_reset_at is None or now_utc >= quota.monthly_reset_at:
            quota.monthly_used_credits = 0
            quota.monthly_reset_at = next_month_start

        return day_start, next_day_start, month_start, next_month_start

    def _day_month_bounds(
        self,
        quota: EnrichmentQuota,
        now_utc: datetime,
    ) -> Tuple[datetime, datetime, datetime, datetime]:
        tz = _get_zone(quota.account_timezone)
        now_local = now_utc.astimezone(tz)

        day_start_local = now_local.replace(hour=0, minute=0, second=0, microsecond=0)
        next_day_local = day_start_local + timedelta(days=1)

        month_start_local = now_local.replace(
            day=1, hour=0, minute=0, second=0, microsecond=0
        )
        next_month_local = (month_start_local + timedelta(days=32)).replace(day=1)

        return (
            day_start_local.astimezone(timezone.utc),
            next_day_local.astimezone(timezone.utc),
            month_start_local.astimezone(timezone.utc),
            next_month_local.astimezone(timezone.utc),
        )

    @staticmethod
    def _remaining_limit(limit: Optional[int], used: Optional[int]) -> Optional[int]:
        if limit is None or limit <= 0:
            return None
        used_value = used or 0
        return max(limit - used_value, 0)

    async def _sum_cost(
        self,
        account_id: Optional[str],
        start: datetime,
        end: datetime,
    ) -> Decimal:
        stmt = select(
            func.coalesce(func.sum(EnrichmentUsage.cost_usd), Decimal("0"))
        ).where(
            EnrichmentUsage.created_at >= start,
            EnrichmentUsage.created_at < end,
        )
        if account_id is not None:
            stmt = stmt.where(EnrichmentUsage.account_id == account_id)

        result = await self.db.execute(stmt)
        return _coerce_decimal(result.scalar())

    @staticmethod
    def _normalize_duration(value: Any) -> Optional[Decimal]:
        if value is None:
            return None
        try:
            return _round_duration(Decimal(str(value)))
        except Exception:  # pragma: no cover - defensive
            logger.warning(
                "Failed to normalize enrichment duration '%s'; dropping metric.", value
            )
            return None

    def _build_allowed_response(
        self,
        *,
        quota: EnrichmentQuota,
        required_requests: int,
        required_credits: int,
        allowance_remaining: int,
        pack_remaining: int,
        daily_remaining: Optional[int],
        monthly_remaining: Optional[int],
        reservation: Optional[EnrichmentReservation],
        day_start: Optional[datetime] = None,
        next_day_start: Optional[datetime] = None,
    ) -> Dict[str, Any]:
        lease_expires_at = (
            reservation.lease_expires_at.isoformat()
            if reservation and reservation.lease_expires_at
            else None
        )
        return {
            "allowed": True,
            "replay": False,
            "status_code": 200,
            "required_requests": required_requests,
            "required_credits": required_credits,
            "credits_per_request": quota.credits_per_request,
            "allowance_remaining": allowance_remaining,
            "pack_remaining": pack_remaining,
            "daily_remaining": daily_remaining,
            "monthly_remaining": monthly_remaining,
            "reservation_id": reservation.id if reservation else None,
            "lease_expires_at": lease_expires_at,
            "provider_cost_per_request_usd": float(
                _round_money(_coerce_decimal(quota.provider_cost_per_request_usd))
            ),
            "price_per_request_usd": float(
                _round_money(_coerce_decimal(quota.price_per_request_usd))
            ),
            "day_start_utc": day_start.isoformat() if day_start else None,
            "next_day_start_utc": next_day_start.isoformat()
            if next_day_start
            else None,
        }

    def _build_denied_response(
        self,
        *,
        quota: EnrichmentQuota,
        reason: str,
        required_requests: int,
        required_credits: int,
        allowance_remaining: int = 0,
        pack_remaining: int = 0,
        daily_remaining: Optional[int] = None,
        monthly_remaining: Optional[int] = None,
        day_start: Optional[datetime] = None,
        next_day_start: Optional[datetime] = None,
        status_code: int = 429,
    ) -> Dict[str, Any]:
        return {
            "allowed": False,
            "replay": False,
            "reason": reason,
            "status_code": status_code,
            "required_requests": required_requests,
            "required_credits": required_credits,
            "credits_per_request": quota.credits_per_request,
            "allowance_remaining": allowance_remaining,
            "pack_remaining": pack_remaining,
            "daily_remaining": daily_remaining,
            "monthly_remaining": monthly_remaining,
            "provider_cost_per_request_usd": float(
                _round_money(_coerce_decimal(quota.provider_cost_per_request_usd))
            ),
            "price_per_request_usd": float(
                _round_money(_coerce_decimal(quota.price_per_request_usd))
            ),
            "day_start_utc": day_start.isoformat() if day_start else None,
            "next_day_start_utc": next_day_start.isoformat()
            if next_day_start
            else None,
        }
