import { Readable, type Transform } from "node:stream";
import { ReadableStream } from "node:stream/web";
import { promisify } from "node:util";
import * as zlib from "node:zlib";
import { $atom, $hook, $inject, $state, Alepha, type Static, z } from "alepha";
import type { ServerResponse } from "alepha/server";

const gzip = promisify(zlib.gzip);
const createGzip = zlib.createGzip;
const brotli = promisify(zlib.brotliCompress);
const createBrotliCompress = zlib.createBrotliCompress;
const zstd = zlib.zstdCompress ? promisify(zlib.zstdCompress) : undefined;
const createZstdCompress = zstd ? zlib.createZstdCompress : undefined;

// ---------------------------------------------------------------------------------------------------------------------

/**
 * Compression configuration atom.
 */
export const compressOptions = $atom({
  name: "alepha.server.compress.options",
  schema: z.object({
    disabled: z
      .boolean()
      .describe("Disable response compression entirely.")
      .optional(),
    allowedContentTypes: z
      .array(z.string())
      .describe("Content types eligible for compression."),
  }),
  default: {
    allowedContentTypes: [
      "application/json",
      "text/html",
      "application/javascript",
      "text/plain",
      "text/css",
    ],
  },
});

export type CompressOptions = Static<typeof compressOptions.schema>;

declare module "alepha" {
  interface State {
    [compressOptions.key]: CompressOptions;
  }
}

// ---------------------------------------------------------------------------------------------------------------------

export class ServerCompressProvider {
  static compressors: Record<
    string,
    | {
        compress: (...args: any[]) => Promise<Buffer>;
        stream: (options?: any) => Transform;
      }
    | undefined
  > = {
    gzip: {
      compress: gzip,
      stream: createGzip,
    },
    br: {
      compress: brotli,
      stream: createBrotliCompress,
    },
    zstd:
      zstd && createZstdCompress
        ? {
            compress: zstd,
            stream: createZstdCompress,
          }
        : undefined,
  };

  protected readonly alepha = $inject(Alepha);
  protected readonly options = $state(compressOptions);

  public readonly onResponse = $hook({
    on: "server:onResponse",
    handler: async ({ request, response }) => {
      if (this.options.disabled) {
        return;
      }

      // In serverless (Cloudflare Workers), skip compression entirely:
      // Cloudflare's edge network automatically compresses responses
      if (this.alepha.isServerless()) {
        return;
      }

      // skip if already compressed
      if (response.headers["content-encoding"]) {
        return;
      }

      const acceptEncoding = request.headers["accept-encoding"]; // skip if no accept-encoding header
      if (!acceptEncoding) {
        return;
      }

      // skip if not json or html (for now)
      if (!this.isAllowedContentType(response.headers["content-type"])) {
        return;
      }

      for (const encoding of ["zstd", "br", "gzip"] as const) {
        if (
          acceptEncoding.includes(encoding) &&
          ServerCompressProvider.compressors[encoding]
        ) {
          await this.compress(encoding, response);
          return;
        }
      }
    },
  });

  protected isAllowedContentType(contentType: string | undefined): boolean {
    if (!contentType) {
      return false;
    }

    const lowerContentType = contentType.toLowerCase();

    return !!this.options.allowedContentTypes.find((it) =>
      lowerContentType.includes(it),
    );
  }

  protected async compress(
    encoding: keyof typeof ServerCompressProvider.compressors,
    response: ServerResponse,
  ): Promise<void> {
    const body = response.body; // can be string or Buffer or ArrayBuffer or Readable

    const compressor = ServerCompressProvider.compressors[encoding];
    if (!compressor) {
      return;
    }

    const params = this.getParams(encoding);

    if (
      typeof body === "string" ||
      Buffer.isBuffer(body) ||
      body instanceof ArrayBuffer
    ) {
      const compressed = await compressor.compress(body, {
        params,
      });
      this.setHeaders(response, encoding);
      response.headers["content-length"] = compressed.length.toString();
      response.body = compressed;
      return;
    }

    if (typeof body === "object" && body instanceof Readable) {
      this.setHeaders(response, encoding);
      response.body = body.pipe(compressor.stream({ params }));
      return;
    }

    if (typeof body === "object" && body instanceof ReadableStream) {
      this.setHeaders(response, encoding);
      // For streaming responses, use flush mode to avoid buffering
      response.body = this.createFlushingCompressStream(
        body,
        compressor.stream,
        encoding,
        params,
      );
    }
  }

  /**
   * Create a compressed stream that flushes after each chunk.
   * This is essential for streaming SSR - ensures each chunk is sent immediately.
   */
  protected createFlushingCompressStream(
    input: ReadableStream,
    createCompressor: (options?: any) => Transform,
    encoding: string,
    params: Record<number, any>,
  ): ReadableStream<Uint8Array> {
    const compressor = createCompressor({
      params,
      flush:
        encoding === "br"
          ? zlib.constants.BROTLI_OPERATION_FLUSH
          : zlib.constants.Z_SYNC_FLUSH,
    });
    const reader = Readable.fromWeb(input);

    return new ReadableStream<Uint8Array>({
      start(controller) {
        compressor.on("data", (chunk: Buffer) => {
          controller.enqueue(new Uint8Array(chunk));
        });

        compressor.on("end", () => {
          controller.close();
        });

        compressor.on("error", (err) => {
          controller.error(err);
        });

        reader.on("data", (chunk: Buffer) => {
          compressor.write(chunk);
          // Force flush after each chunk for streaming
          // Cast to any because flush() exists on zlib streams but not in Transform type
          const zlibStream = compressor as any;
          if (encoding === "gzip") {
            zlibStream.flush(zlib.constants.Z_SYNC_FLUSH);
          } else if (encoding === "br") {
            zlibStream.flush(zlib.constants.BROTLI_OPERATION_FLUSH);
          } else if (encoding === "zstd") {
            zlibStream.flush();
          }
        });

        reader.on("end", () => {
          compressor.end();
        });

        reader.on("error", (err) => {
          controller.error(err);
        });
      },
    });
  }

  protected getParams(
    encoding: keyof typeof ServerCompressProvider.compressors,
  ): Record<number, any> {
    if (encoding === "zstd") {
      return {
        [zlib.constants.ZSTD_c_compressionLevel]: 3, // default compression level for zstd
      };
    }
    if (encoding === "br") {
      return {};
    }
    if (encoding === "gzip") {
      return {};
    }
    return {};
  }

  protected setHeaders(
    response: ServerResponse,
    encoding: keyof typeof ServerCompressProvider.compressors,
  ): void {
    response.headers.vary = "accept-encoding";
    response.headers["content-encoding"] = encoding;
  }
}