// JSON request/response plumbing: bounded body reads and the json() reply.

export const json = (
  status: number,
  body: unknown,
  headers?: Record<string, string>,
): Response =>
  new Response(JSON.stringify(body), {
    status,
    headers: { "content-type": "application/json", ...headers },
  });

export const MAX_BODY_BYTES = 1024 * 1024;
export const MAX_CONTROL_RESPONSE_BYTES = 1024 * 1024;
export const RESPONSE_BODY_LIMIT_CANCELLATION = Symbol(
  "response body limit reached",
);
// Register carries `instances` (a full host's held ids); size for 10k+
// max-length ids across several namespaces.
export const MAX_REGISTER_BODY_BYTES = 8 * 1024 * 1024;

class BodyTooLargeError extends Error {}

const declaredTooLarge = (headers: Headers, maxBytes: number): boolean => {
  const raw = headers.get("content-length");
  return raw !== null && Number(raw) > maxBytes;
};

const consumeBody = async (
  body: ReadableStream<Uint8Array> | null,
  maxBytes: number,
  collect: boolean,
): Promise<Uint8Array> => {
  if (body === null) return new Uint8Array();
  const reader = body.getReader();
  const chunks: Uint8Array[] = [];
  let size = 0;
  for (;;) {
    const { done, value } = await reader.read();
    if (done) break;
    size += value.length;
    if (size > maxBytes) {
      await reader.cancel(RESPONSE_BODY_LIMIT_CANCELLATION).catch(() => {});
      throw new BodyTooLargeError();
    }
    if (collect) chunks.push(value);
  }
  if (!collect) return new Uint8Array();
  const merged = new Uint8Array(size);
  let offset = 0;
  for (const chunk of chunks) {
    merged.set(chunk, offset);
    offset += chunk.length;
  }
  return merged;
};

const readJsonBody = async (
  req: Request,
  maxBytes: number,
): Promise<unknown> => {
  if (declaredTooLarge(req.headers, maxBytes)) throw new BodyTooLargeError();
  return JSON.parse(
    new TextDecoder().decode(await consumeBody(req.body, maxBytes, true)),
  );
};

const inflight = new Set<Request>();

// undici follows init.signal through a weak ref: a Request collected before
// its fetch settles never aborts, so the request stays reachable until then.
export const fetchHeld = async (
  fetchImpl: (request: Request) => Promise<Response>,
  request: Request,
): Promise<Response> => {
  inflight.add(request);
  try {
    return await fetchImpl(request);
  } finally {
    inflight.delete(request);
  }
};

export const readResponseJson = async (
  response: Response,
  maxBytes: number = MAX_CONTROL_RESPONSE_BYTES,
): Promise<unknown> => {
  if (declaredTooLarge(response.headers, maxBytes)) {
    await response.body
      ?.cancel(RESPONSE_BODY_LIMIT_CANCELLATION)
      .catch(() => {});
    throw new BodyTooLargeError();
  }
  return JSON.parse(
    new TextDecoder().decode(await consumeBody(response.body, maxBytes, true)),
  );
};

export const discardResponseBody = async (
  response: Response,
  maxBytes: number = MAX_CONTROL_RESPONSE_BYTES,
): Promise<void> => {
  if (declaredTooLarge(response.headers, maxBytes)) {
    await response.body
      ?.cancel(RESPONSE_BODY_LIMIT_CANCELLATION)
      .catch(() => {});
    throw new BodyTooLargeError();
  }
  await consumeBody(response.body, maxBytes, false);
};

export const readBodyOrReject = async (
  req: Request,
  maxBytes: number = MAX_BODY_BYTES,
): Promise<{ body: unknown } | { response: Response }> => {
  try {
    return { body: await readJsonBody(req, maxBytes) };
  } catch (err) {
    if (err instanceof BodyTooLargeError) {
      return { response: json(413, { error: "body too large" }) };
    }
    return { response: json(400, { error: "invalid json" }) };
  }
};