import { env } from "cloudflare:workers";
import { repoDb } from "../scoped-db";
import { ownerSiteScopes } from "./store";
import { addCounter, counterTotals } from "../counters";
import { SCRATCH_LIMITS } from "./config";
import { meterSiteRequest } from "./metering";
import type { SiteMeta } from "./config";

const FLAG_TTL_SECONDS = 2 * 24 * 60 * 60;
const HTTP_SWITCHING = 101;

export type UsageDelta = { requests?: number; cpuMs?: number; bytesOut?: number };

export type AccountUsage = {
  day: string;
  requests: number;
  cpuMs: number;
  bytesOut: number;
};

export function usageDay(now = Date.now()): string {
  return new Date(now).toISOString().slice(0, 10);
}

function overKey(siteId: string, day: string): string {
  return `sites:over:${siteId}:${day}`;
}

/**
 * Adds to a site's daily usage through its Counters object (D1's `site_usage` row follows within seconds). Unclaimed scratch sites (no account) are the free tier: once a daily cap is
 * crossed a KV flag makes the dispatcher refuse further requests until the next UTC day.
 */
export async function recordUsage(
  meta: Pick<SiteMeta, "id" | "ownerId" | "scratch">,
  delta: UsageDelta,
): Promise<void> {
  const day = usageDay();
  const [requests, cpuMs, bytesOut] = await addCounter(
    "site-usage",
    meta.id,
    day,
    [delta.requests ?? 0, Math.round(delta.cpuMs ?? 0), delta.bytesOut ?? 0],
    meta.scratch ? null : meta.ownerId,
  );
  if (!meta.scratch) return;
  const over =
    requests > SCRATCH_LIMITS.requestsPerDay ||
    cpuMs > SCRATCH_LIMITS.cpuMsPerDay ||
    bytesOut > SCRATCH_LIMITS.bytesOutPerDay;
  if (over) await env.KV.put(overKey(meta.id, day), "1", { expirationTtl: FLAG_TTL_SECONDS });
}

export async function overFreeTier(siteId: string): Promise<boolean> {
  return (await env.KV.get(overKey(siteId, usageDay()))) !== null;
}

/**
 * Calls `finish` with the body's byte count once the response has been fully sent or the client went away (so
 * aborting a download can't dodge the meter), without buffering. Upgrades and empty bodies finish immediately.
 */
export function countedResponse(response: Response, finish: (bytes: number) => void): Response {
  if (response.status === HTTP_SWITCHING || response.webSocket || !response.body) {
    finish(0);
    return response;
  }
  const reader = response.body.getReader();
  let bytes = 0;
  let done = false;
  const once = () => {
    if (done) return;
    done = true;
    finish(bytes);
  };
  const body = new ReadableStream<Uint8Array>({
    async pull(controller) {
      let next: Awaited<ReturnType<typeof reader.read>>;
      try {
        next = await reader.read();
      } catch (error) {
        once();
        controller.error(error);
        return;
      }
      if (next.done) {
        once();
        controller.close();
        return;
      }
      bytes += next.value.byteLength;
      controller.enqueue(next.value);
    },
    cancel(reason) {
      once();
      return reader.cancel(reason);
    },
  });
  return new Response(body, response);
}

/** Meters a site request and its response bytes once the body is done. */
export function meteredResponse(
  response: Response,
  meta: Pick<SiteMeta, "id" | "ownerId" | "scratch" | "repoId">,
  ctx: ExecutionContext,
  rayId: string,
): Response {
  return countedResponse(response, (bytes) => meterSiteRequest(ctx, meta, { bytes, rayId }));
}

/** Per-day totals for an account's claimed sites, for billing and the metering dashboard. */
export async function accountUsage(accountId: string, sinceDay: string): Promise<AccountUsage[]> {
  const days = new Map<string, AccountUsage>();
  for (const scope of await ownerSiteScopes(accountId)) {
    const rows = await repoDb(scope).all<{ day: string; requests: number; cpu_ms: number; bytes_out: number }>(
      `SELECT day, SUM(requests) AS requests, SUM(cpu_ms) AS cpu_ms, SUM(bytes_out) AS bytes_out
       FROM site_usage WHERE account_id = ? AND day >= ? GROUP BY day ORDER BY day LIMIT 400`,
      accountId,
      sinceDay,
    );
    for (const row of rows) {
      const day = days.get(row.day) ?? { day: row.day, requests: 0, cpuMs: 0, bytesOut: 0 };
      day.requests += row.requests;
      day.cpuMs += row.cpu_ms;
      day.bytesOut += row.bytes_out;
      days.set(row.day, day);
    }
  }
  return [...days.values()].sort((left, right) => left.day.localeCompare(right.day)).slice(0, 400);
}

/** Bytes held by an account's live deployments (content-addressed objects may be shared, so this is an upper bound). */
export async function accountStorageBytes(accountId: string): Promise<number> {
  let bytes = 0;
  for (const scope of await ownerSiteScopes(accountId)) {
    const row = await repoDb(scope).first<{ bytes: number | null }>(
      `SELECT SUM(d.total_bytes) AS bytes FROM site_deployments d JOIN sites s ON s.id = d.site_id
       WHERE s.owner_id = ? AND s.scratch = 0 AND d.status = 'active'`,
      accountId,
    );
    bytes += row?.bytes ?? 0;
  }
  return bytes;
}

export async function siteUsageToday(siteId: string): Promise<AccountUsage> {
  const day = usageDay();
  const [requests, cpuMs, bytesOut] = await counterTotals("site-usage", siteId, day);
  return { day, requests, cpuMs, bytesOut };
}

const USAGE_DAYS = 31;

/** `GET /api/v1/user/site-usage`: the caller's daily site usage for the last month plus current storage. */
export async function getSiteUsage(viewer: { id: string }): Promise<Response> {
  const since = usageDay(Date.now() - USAGE_DAYS * 24 * 60 * 60 * 1000);
  const [days, storageBytes] = await Promise.all([accountUsage(viewer.id, since), accountStorageBytes(viewer.id)]);
  return Response.json({ since, storage_bytes: storageBytes, days });
}