/**
 * AAO Gateway Phase 2b — budget enforcement.
 *
 * Post-hoc check: before serving a request we compare the matched key's
 * cumulative usage for the current UTC month against its `tokens_budget`
 * column. If the running total already meets / exceeds the budget, the
 * request is rejected with 402 `{error: 'budget_exhausted', detail: ...}`.
 *
 * This is "post-hoc" because:
 *  1. We don't pre-estimate the tokens an outstanding request will use
 *     (estimation is noisy + LLM-dependent).
 *  2. The very request that pushes the counter past the limit still
 *     completes — but the NEXT request is rejected.
 *
 * Concurrency caveat (no pre-reservation):
 * ─────────────────────────────────────────
 * With N simultaneous in-flight requests on the same key, ALL N can
 * pass the pre-flight check while `used < budget`. Each one then adds
 * its tokens at completion, so the budget can be exceeded by up to
 *   N × (max completion tokens per request)
 * before the (N+1)-th request is rejected. Effective max overshoot
 * equals the key's in-flight concurrency × max_completion_tokens.
 *
 * Phase 2b accepts this trade-off because the alternative — pre-reserving
 * an estimated token cost up front, then refunding the delta — adds:
 *  - estimation accuracy noise (we'd reserve based on guesses)
 *  - extra DB round-trip on every request
 *  - reservation TTL management (orphaned reservations on crash)
 *
 * For Phase 2b operators should size budgets with headroom or limit
 * concurrency upstream. Phase 3 candidate: a small pre-reservation
 * (e.g., max_completion_tokens from the request body) that's decremented
 * to the real value at completion — would tighten the overshoot bound
 * to "at most one request's worth" as originally desired.
 *
 * The actual counter increment happens after upstream replies; see
 * stream-proxy.ts.
 *
 * Pure function design: this module never reads the system clock outside
 * `deps.now` and never touches the DB except via the injected Repository
 * slice. That makes it trivially unit-testable.
 */
import type { Request, Response, NextFunction, RequestHandler } from 'express';
import type { AuthenticatedRequest } from './auth.js';
import type { GatewayKeyUsage } from '../db/repository.js';
import { currentPeriodYearMonth } from './period.js';
import { logger } from '../logger.js';

/** Repository slice used by the budget check. */
export interface BudgetCheckRepo {
  getGatewayKeyUsage(keyId: string, periodStart: string): GatewayKeyUsage | null;
  findGatewayVirtualKeyById(id: string): { tokensBudget: number | null } | null;
}

export interface BudgetCheckResult {
  allowed: boolean;
  used: number;
  budget: number | null;
  period: string;
}

export interface BudgetCheckDeps {
  repo: BudgetCheckRepo;
  /** Test hook for the period clock. Defaults to `new Date()` per call. */
  now?: () => Date;
  /**
   * Phase 3b: optional metrics handle. When set, the 402 reject path
   * increments `requestsTotal{status="budget_exhausted"}` so SLO
   * dashboards can show budget-driven rejections separately from rate
   * limits and upstream errors.
   */
  metrics?: import('../metrics/gateway-metrics.js').GatewayMetrics;
}

/**
 * Pure budget check, easy to call from unit tests. Returns the verdict
 * without performing any HTTP work.
 *
 * `tokensBudget=null` is the unlimited case → always allowed.
 *
 * `used` is the sum of `tokens_in + tokens_out` for the current UTC
 * month; we treat input + output tokens as fungible because the design
 * doc commits to a single tokens counter (no cost-tier split until
 * Phase 3).
 */
export function buildBudgetCheck(deps: BudgetCheckDeps): (keyId: string, tokensBudget: number | null) => BudgetCheckResult {
  const nowFn = deps.now ?? ((): Date => new Date());
  return (keyId, tokensBudget) => {
    const period = currentPeriodYearMonth(nowFn());
    if (tokensBudget === null || tokensBudget <= 0) {
      return { allowed: true, used: 0, budget: null, period };
    }
    const usage = deps.repo.getGatewayKeyUsage(keyId, period);
    const used = (usage?.tokensIn ?? 0) + (usage?.tokensOut ?? 0);
    return { allowed: used < tokensBudget, used, budget: tokensBudget, period };
  };
}

/**
 * Express middleware wrapper. Skips the check when:
 *  - There's no gatewayAuth on the request (auth middleware will have
 *    already 401'd; defensive guard).
 *  - The matched key came from `config` (Phase 1 fallback path — those
 *    keys have no DB row id to attach usage to).
 *  - The key has no DB id (defensive — shouldn't happen but guard
 *    against future code paths).
 *
 * When the check rejects, the middleware terminates the request with
 * HTTP 402 and a structured `detail` string suitable for an SDK to
 * surface to the user. The 402 status is borrowed from LiteLLM's
 * `budget_exhausted` convention so client libraries that already
 * handle that case keep working.
 */
export function buildBudgetCheckMiddleware(deps: BudgetCheckDeps): RequestHandler {
  const check = buildBudgetCheck(deps);
  return function gatewayBudgetCheckMiddleware(req: Request, res: Response, next: NextFunction): void {
    const auth = (req as AuthenticatedRequest).gatewayAuth;
    if (!auth || auth.source !== 'db' || !auth.keyId) {
      // No DB-backed key → nothing to bill against. Allow through and
      // log once-per-process via the auth deprecation warning path
      // (already handled in auth.ts).
      next();
      return;
    }
    const row = deps.repo.findGatewayVirtualKeyById(auth.keyId);
    const tokensBudget = row?.tokensBudget ?? null;
    const verdict = check(auth.keyId, tokensBudget);
    if (verdict.allowed) {
      next();
      return;
    }
    logger.warn(
      `[gateway-budget] reject team=${auth.team} key=${auth.keyId} used=${verdict.used} budget=${verdict.budget ?? 'null'} period=${verdict.period}`,
    );
    res.status(402).json({
      error: 'budget_exhausted',
      detail: `team=${auth.team} used=${verdict.used} budget=${verdict.budget} period=${verdict.period}`,
    });
    if (deps.metrics) {
      try {
        deps.metrics.requestsTotal
          .labels({ team: auth.team, backend: 'none', model: 'none', status: 'budget_exhausted' })
          .inc();
      } catch {
        /* metrics never affect the response */
      }
    }
  };
}
