export class RequestBodyLimitError extends Error {
  constructor(readonly limitBytes: number) {
    super(`Request body exceeds the ${limitBytes}-byte route limit.`);
    this.name = "RequestBodyLimitError";
  }
}

export async function readLimitedRequestBytes(request: Request, maxBytes: number): Promise<Uint8Array> {
  assertValidRequestBodyLimit(maxBytes);

  await rejectOversizedDeclaredBody(request, maxBytes);
  if (!request.body) return new Uint8Array();

  const reader = request.body.getReader();
  const chunks: Uint8Array[] = [];
  let totalBytes = 0;
  while (true) {
    const chunk = await reader.read();
    if (chunk.done) break;
    totalBytes += chunk.value.byteLength;
    if (totalBytes > maxBytes) {
      await reader.cancel().catch(() => undefined);
      throw new RequestBodyLimitError(maxBytes);
    }
    chunks.push(chunk.value);
  }

  const body = new Uint8Array(totalBytes);
  let offset = 0;
  for (const chunk of chunks) {
    body.set(chunk, offset);
    offset += chunk.byteLength;
  }
  return body;
}

/**
 * Parse multipart or URL-encoded form data while enforcing the bytes actually
 * consumed by the parser. `Content-Length` remains a fast rejection path, but
 * missing or understated headers cannot bypass the limit.
 */
export async function readLimitedRequestFormData(request: Request, maxBytes: number): Promise<FormData> {
  assertValidRequestBodyLimit(maxBytes);
  await rejectOversizedDeclaredBody(request, maxBytes);
  if (!request.body) return request.formData();

  const limitedBody = createLimitedRequestBodyStream(request.body, maxBytes);
  const contentType = request.headers.get("content-type");
  return new Response(limitedBody, {
    headers: contentType ? { "content-type": contentType } : undefined,
  }).formData();
}

function assertValidRequestBodyLimit(maxBytes: number): void {
  if (!Number.isSafeInteger(maxBytes) || maxBytes <= 0) {
    throw new TypeError("Request body limit must be a positive safe integer.");
  }
}

async function rejectOversizedDeclaredBody(request: Request, maxBytes: number): Promise<void> {
  const contentLength = request.headers.get("content-length");
  if (contentLength && /^\d+$/.test(contentLength) && Number(contentLength) > maxBytes) {
    await request.body?.cancel().catch(() => undefined);
    throw new RequestBodyLimitError(maxBytes);
  }
}

function createLimitedRequestBodyStream(body: ReadableStream<Uint8Array>, maxBytes: number): ReadableStream<Uint8Array> {
  const reader = body.getReader();
  let totalBytes = 0;
  return new ReadableStream<Uint8Array>({
    async pull(controller) {
      try {
        const chunk = await reader.read();
        if (chunk.done) {
          controller.close();
          return;
        }
        totalBytes += chunk.value.byteLength;
        if (totalBytes > maxBytes) {
          await reader.cancel().catch(() => undefined);
          controller.error(new RequestBodyLimitError(maxBytes));
          return;
        }
        controller.enqueue(chunk.value);
      } catch (error) {
        controller.error(error);
      }
    },
    cancel(reason) {
      return reader.cancel(reason);
    },
  });
}