import { env } from "cloudflare:workers";
import { accountDb, repoDb, type ScopedStatement } from "../scoped-db";
import { claimScratchUsage, eventStatements, periodEnd, periodOf, resolveAccount } from "../metering/record";
import { billingSettings, hardCap, siteCapKey } from "../metering/quota";
import { scratchSiteAccount } from "../metering/attribution";
import type { MeterId } from "~/lib/meters";
import { overFreeTier, recordUsage, usageDay } from "./usage";
import type { SiteMeta } from "./config";

/**
 * The only place Sites talks to metering. Per-request counters stay in `site_usage` (Durable Object counters,
 * flushed to D1) so the hot path never writes D1; `rollupScopeSiteUsage` folds each finished UTC day into the shared
 * meters, billed to the owner or to `scratch:site:<id>` (redirected to whoever claimed it).
 */
type Metered = Pick<SiteMeta, "id" | "ownerId" | "scratch" | "repoId">;

const ROLLUP_PAGE = 500;
const ROLLUP_MAX_PAGES = 40;
const CAP_FLAG_CACHE_SECONDS = 60;
/** Every hosting meter whose hard cap stops an owner's Sites from serving. */
export const HOSTING_METERS: MeterId[] = ["compute.site_requests", "compute.site_cpu", "network.sites"];
export const ROLLUP_DAYS = 3;

/** One request and its response bytes, after the response finished streaming. */
export function meterSiteRequest(ctx: ExecutionContext, site: Metered, input: { bytes: number; rayId: string }): void {
  ctx.waitUntil(recordUsage(site, { requests: 1, bytesOut: input.bytes }));
}

/** CPU time of a dynamic site's isolate, from its tail events. */
export async function meterSiteCpu(site: Metered, cpuMs: number): Promise<void> {
  await recordUsage(site, { cpuMs });
}

/** Before serving a scratch site: false once its free tier is used up (the dispatcher answers 429). */
export async function scratchWithinQuota(site: Metered): Promise<boolean> {
  return !site.scratch || !(await overFreeTier(site.id));
}

/** Set by the admin "suspend hosting" action; unlike the cap flag, plan changes never clear it. */
export function hostingSuspendedKey(accountId: string): string {
  return `hosting:suspended:${accountId}`;
}

/** An owner's Sites stop serving once a hosting meter reached its hard cap, or an admin suspended them. */
export async function ownerHostingAllowed(ownerId: string | null): Promise<boolean> {
  if (!ownerId) return true;
  const options = { cacheTtl: CAP_FLAG_CACHE_SECONDS };
  const [capped, suspended] = await Promise.all([
    env.KV.get(siteCapKey(ownerId), options),
    env.KV.get(hostingSuspendedKey(ownerId), options),
  ]);
  return capped === null && suspended === null;
}

export async function ownerWithinSiteQuota(site: Metered): Promise<boolean> {
  if (site.scratch) return true;
  return ownerHostingAllowed(site.ownerId);
}

/** After a successful claim: usage recorded against the scratch account moves to the claimer. Never fails the claim. */
export async function onScratchClaimed(siteId: string, userId: string): Promise<void> {
  await claimScratchUsage(scratchSiteAccount(siteId), userId).catch((error: unknown) =>
    console.error("Scratch usage fold failed", siteId, error),
  );
}

type UsageRow = {
  scope: string;
  account_id: string | null;
  repo_id: string | null;
  requests: number;
  cpu_ms: number;
  bytes_out: number;
};

type RollupSource = {
  /** Prefix of each event's idempotency key and of the per-day done flag. */
  name: "site-usage";
  sql: string;
  meters: [MeterId, (row: UsageRow) => number][];
  account(row: UsageRow): string | null;
  attributes(row: UsageRow): Record<string, string>;
};

export const ROLLUP_SOURCES: RollupSource[] = [
  {
    name: "site-usage",
    sql: `SELECT u.site_id AS scope, u.account_id, s.repo_id, u.requests, u.cpu_ms, u.bytes_out
          FROM site_usage u LEFT JOIN sites s ON s.id = u.site_id
          WHERE u.day = ? ORDER BY u.site_id LIMIT ? OFFSET ?`,
    meters: [
      ["compute.site_requests", (row) => row.requests],
      ["compute.site_cpu", (row) => row.cpu_ms],
      ["network.sites", (row) => row.bytes_out],
    ],
    account: (row) => row.account_id ?? scratchSiteAccount(row.scope),
    attributes: (row) => ({ siteId: row.scope }),
  },
];

/** A finished day is folded this long after midnight UTC, once the last counters have flushed. */
export const ROLLUP_GRACE_MS = 10 * 60 * 1000;

/**
 * Folds every finished UTC day of this scope's site counters into the billed accounts' meters, once: the scope's
 * `site_rollups` remembers folded days and idempotency keys make a retried day safe. Run by the scope's RepoState
 * alarm after each day ends (repo-state/housekeeping.ts), then sets the serving stop for any account it billed
 * that reached a hosting hard cap.
 */
export async function rollupScopeSiteUsage(scopeId: string, now = Date.now()): Promise<void> {
  const db = repoDb(scopeId);
  const today = usageDay(now);
  const days = (
    await db.all<{ day: string }>(
      "SELECT DISTINCT day FROM site_usage WHERE day < ? AND day NOT IN (SELECT day FROM site_rollups) ORDER BY day LIMIT 31",
      today,
    )
  ).map((row) => row.day);
  const billed = new Set<string>();
  for (const source of ROLLUP_SOURCES) {
    for (const day of days) {
      const at = Date.parse(`${day}T23:59:59.999Z`);
      const period = periodOf(at);
      const byAccount = new Map<string, ScopedStatement[]>();
      for (let page = 0; page < ROLLUP_MAX_PAGES; page += 1) {
        const rows = await db.all<UsageRow>(source.sql, day, ROLLUP_PAGE, page * ROLLUP_PAGE);
        for (const row of rows) {
          const owner = source.account(row);
          if (!owner) continue;
          const account = await resolveAccount(owner);
          const statements = byAccount.get(account) ?? [];
          for (const [meter, quantity] of source.meters) {
            const value = quantity(row);
            if (value <= 0) continue;
            statements.push(
              ...eventStatements(
                {
                  accountId: account,
                  meter,
                  quantity: value,
                  repoId: row.repo_id,
                  idempotencyKey: `${source.name}:${row.scope}:${day}:${meter}`,
                  attributes: source.attributes(row),
                },
                account,
                period,
                at,
              ),
            );
          }
          byAccount.set(account, statements);
        }
        if (rows.length < ROLLUP_PAGE) break;
      }
      // Each account's events are one transaction in its own object; idempotency keys make a retried day safe.
      for (const [account, statements] of byAccount) {
        if (statements.length > 0) await accountDb(account).batch(statements);
        billed.add(account);
      }
    }
  }
  for (const day of days) await db.run("INSERT OR IGNORE INTO site_rollups (day, at) VALUES (?, ?)", day, now);
  await flagOverCapOwners([...billed], now);
}

/** Sets the serving stop for each non-scratch account whose site meters reached a plan hard cap this period. */
async function flagOverCapOwners(accounts: string[], now: number): Promise<void> {
  const period = periodOf(now);
  const ttl = Math.max(60, Math.ceil((periodEnd(period) - now) / 1000));
  for (const account of accounts) {
    if (account.startsWith("scratch:")) continue;
    const rows = await accountDb(account).all<{ meter: MeterId; quantity: number }>(
      `SELECT meter, quantity FROM usage_periods
       WHERE period = ? AND meter IN (${HOSTING_METERS.map(() => "?").join(", ")})`,
      period,
      ...HOSTING_METERS,
    );
    const settings = await billingSettings(account);
    const over = rows.find((row) => {
      const cap = hardCap(settings, row.meter);
      return cap !== null && row.quantity >= cap;
    });
    if (over) await env.KV.put(siteCapKey(account, period), over.meter, { expirationTtl: ttl });
  }
}