import { env } from "cloudflare:workers";
import { stripeEnabled } from "../billing/config";
import { MARKUP, METER_PRICES } from "~/lib/pricing";
import {
  commitUsage,
  PlatformLimitError,
  QuotaExceededError,
  releaseUsage,
  remainingQuota,
  reserveUsage,
} from "./quota";
import type { Reservation } from "./quota";
import { isRecorded, recordUsage } from "./record";
import { captureServerEvent } from "../analytics";
import { addCounter, counterTotals } from "../counters";

export type AiMeterContext = {
  accountId: string;
  repoId?: string | null;
  /** The feature making the calls, e.g. "mergefix" or "clef". */
  feature: string;
  /** A durable id for the unit of work (run id); call keys are `ai:<scope>:<index>`. */
  scope: string;
};

export type AiTokens = { inputTokens: number; outputTokens: number };

/** Tokens held for a call before it starts (or what is left, if less); afterCall settles it at the real count. */
export const AI_TOKEN_HOLD = 8_000;
const AI_TOKEN_HOLD_TTL_MS = 15 * 60 * 1000;
const DEFAULT_DAILY_CEILING_CENTS = 10_000;

/** The ceiling in cents, null when off, or NaN when misconfigured (which refuses every call). */
export function dailyAiCeilingCents(raw = env.AI_DAILY_SPEND_CEILING_CENTS): number | null {
  // Unset: on for Stripe-billed instances, off for self-hosted ones that never opted in.
  if (raw === undefined || raw.trim() === "") return stripeEnabled() ? DEFAULT_DAILY_CEILING_CENTS : null;
  if (raw.trim().toLowerCase() === "off") return null;
  const cents = Number(raw);
  return Number.isFinite(cents) && cents >= 0 ? cents : Number.NaN;
}

/** Platform cost (Cloudflare's price, before markup) of one unit, in microcents. */
function costMicro(meter: "ai.calls" | "ai.tokens"): number {
  return METER_PRICES[meter] / MARKUP;
}

/**
 * Platform-wide safety net, independent of any account's budget: once today's (UTC) platform-billed AI cost
 * across every account, in-flight token holds included, reaches AI_DAILY_SPEND_CEILING_CENTS, no new call
 * starts. The total is one running counter (`platform-ai`), not a sum over accounts. It fails closed: a bad
 * setting or an unreadable total refuses the call.
 */
export async function assertPlatformAiBudget(now = Date.now()): Promise<void> {
  const ceiling = dailyAiCeilingCents();
  if (ceiling === null) return;
  const refuse = () =>
    new PlatformLimitError("ai.calls", "AI is paused for today: the platform's daily AI budget is used up");
  if (Number.isNaN(ceiling)) {
    console.error("AI_DAILY_SPEND_CEILING_CENTS is not a number of cents; refusing AI calls");
    throw refuse();
  }
  let spentMicro = 0;
  try {
    spentMicro = (await counterTotals("platform-ai", PLATFORM_SCOPE, utcDay(now)))[0];
  } catch (error) {
    console.error("Could not read today's AI spend; refusing AI calls", error);
    throw refuse();
  }
  if (!Number.isFinite(spentMicro) || spentMicro >= ceiling * 1e6) throw refuse();
}

const PLATFORM_SCOPE = "global";

function utcDay(now: number): string {
  return new Date(now).toISOString().slice(0, 10);
}

/**
 * Moves today's platform AI total. A call adds its worst case (the call plus a full token hold) when it starts
 * and the difference to its real tokens when it settles, so in-flight calls count against the ceiling. Best
 * effort: a lost update only skews the safety net, never an account's bill.
 */
function addPlatformAiSpend(micro: number, now = Date.now()): Promise<void> {
  if (micro === 0 || dailyAiCeilingCents() === null) return Promise.resolve();
  return addCounter("platform-ai", PLATFORM_SCOPE, utcDay(now), [Math.round(micro), 0, 0]).then(
    () => undefined,
    (error: unknown) => console.warn("Platform AI spend update failed", error),
  );
}

/**
 * The single metering path for AI Gateway calls. Call `beforeCall` before every model call: it checks the
 * platform's daily AI ceiling, atomically reserves one `ai.calls` (counted at once, so calls that time out
 * still count) and holds AI_TOKEN_HOLD `ai.tokens` against the account's cap and spending limit, throwing
 * QuotaExceededError when either does not fit. Call `afterCall` with the reported tokens, which settles the
 * hold at the real count. Bring-your-own-key calls are recorded as `ai.byok_calls` and `ai.byok_tokens`:
 * free, never capped, and never drawing on plan allowances or prepaid credit.
 * Keys derive from `scope` and the call index, so a retried workflow step never double counts.
 */
export function aiMeter(context: AiMeterContext) {
  let index = 0;
  const holds = new Map<number, Reservation>();
  const holdId = (callIndex: number) => `ai:${context.scope}:${callIndex}:tokens-hold`;
  const reportUsage = (model: string, usage: AiTokens, byok: boolean) => {
    // Bring-your-own-key model ids are free text chosen by the account, so only platform model ids are sent.
    captureServerEvent(
      "ai_usage",
      context.accountId,
      {
        feature: context.feature,
        provider_kind: byok ? "byok" : "platform",
        model: byok ? "byok" : model,
        input_tokens: Math.max(0, Math.round(usage.inputTokens)),
        output_tokens: Math.max(0, Math.round(usage.outputTokens)),
        calls: 1,
      },
      { ownerAccountId: context.accountId },
    );
  };
  return {
    async beforeCall(model: string, byok = false): Promise<number> {
      const callIndex = index++;
      const key = `ai:${context.scope}:${callIndex}`;
      const attributes = { model, feature: context.feature };
      if (byok) {
        await recordUsage({ ...context, meter: "ai.byok_calls", quantity: 1, attributes, idempotencyKey: key });
        return callIndex;
      }
      if (await isRecorded(context.accountId, key)) return callIndex;
      await assertPlatformAiBudget();
      // Hold what is left when less than a full hold remains: the call may still start, its overrun is bounded
      // by one call, and concurrent calls cannot all claim the same last tokens.
      const left = await remainingQuota(context.accountId, "ai.tokens");
      if (left !== null && left <= 0) {
        throw new QuotaExceededError("ai.tokens", 0, undefined, context.accountId);
      }
      const tokens = await reserveUsage({
        id: holdId(callIndex),
        accountId: context.accountId,
        meter: "ai.tokens",
        quantity: left === null ? AI_TOKEN_HOLD : Math.min(AI_TOKEN_HOLD, left),
        ttlMs: AI_TOKEN_HOLD_TTL_MS,
      });
      try {
        const reservation = await reserveUsage({
          id: key,
          accountId: context.accountId,
          meter: "ai.calls",
          quantity: 1,
        });
        await commitUsage(reservation, { quantity: 1, repoId: context.repoId, attributes, idempotencyKey: key });
      } catch (error) {
        await releaseUsage(tokens.id, tokens.accountId);
        throw error;
      }
      holds.set(callIndex, tokens);
      await addPlatformAiSpend(costMicro("ai.calls") + tokens.quantity * costMicro("ai.tokens"));
      return callIndex;
    },
    async afterCall(callIndex: number, model: string, usage: AiTokens, byok = false): Promise<void> {
      const quantity = Math.max(0, Math.round(usage.inputTokens)) + Math.max(0, Math.round(usage.outputTokens));
      const attributes = {
        model,
        feature: context.feature,
        input: usage.inputTokens,
        output: usage.outputTokens,
        byok,
      };
      const idempotencyKey = `ai:${context.scope}:${callIndex}:tokens`;
      const hold = byok ? undefined : holds.get(callIndex);
      if (hold) {
        holds.delete(callIndex);
        await commitUsage(hold, { quantity, repoId: context.repoId, attributes, idempotencyKey });
        await addPlatformAiSpend((quantity - hold.quantity) * costMicro("ai.tokens"));
        reportUsage(model, usage, byok);
        return;
      }
      await recordUsage({
        accountId: context.accountId,
        repoId: context.repoId,
        meter: byok ? "ai.byok_tokens" : "ai.tokens",
        quantity,
        attributes,
        idempotencyKey,
      });
      // A retried step that lost its hold: its start already counted a full hold, which this replaces.
      if (!byok) await addPlatformAiSpend((quantity - AI_TOKEN_HOLD) * costMicro("ai.tokens"));
      reportUsage(model, usage, byok);
      // A retried step that lost its in-memory hold still frees it rather than waiting for the TTL.
      if (!byok) await releaseUsage(holdId(callIndex), context.accountId);
    },
  };
}