import { env } from "cloudflare:workers";
import { all, first } from "../db";
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;
}

/**
 * Meters a request once its response has been fully sent or the client went away (so aborting a download can't
 * dodge the meter), counting bytes without buffering. Upgrades and empty bodies are metered immediately.
 */
export function meteredResponse(
  response: Response,
  meta: Pick<SiteMeta, "id" | "ownerId" | "scratch" | "repoId">,
  ctx: ExecutionContext,
  rayId: string,
): Response {
  if (response.status === HTTP_SWITCHING || response.webSocket || !response.body) {
    meterSiteRequest(ctx, meta, { bytes: 0, rayId });
    return response;
  }
  const reader = response.body.getReader();
  let bytes = 0;
  let metered = false;
  const finish = () => {
    if (metered) return;
    metered = true;
    meterSiteRequest(ctx, meta, { bytes, rayId });
  };
  const body = new ReadableStream<Uint8Array>({
    async pull(controller) {
      let next: Awaited<ReturnType<typeof reader.read>>;
      try {
        next = await reader.read();
      } catch (error) {
        finish();
        controller.error(error);
        return;
      }
      if (next.done) {
        finish();
        controller.close();
        return;
      }
      bytes += next.value.byteLength;
      controller.enqueue(next.value);
    },
    cancel(reason) {
      finish();
      return reader.cancel(reason);
    },
  });
  return new Response(body, response);
}

/** 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 rows = await 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,
  );
  return rows.map((row) => ({ day: row.day, requests: row.requests, cpuMs: row.cpu_ms, bytesOut: row.bytes_out }));
}

/** 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> {
  const row = await 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,
  );
  return row?.bytes ?? 0;
}

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 });
}