import * as miniflare from "miniflare";
import assert from "node:assert";
import { once } from "node:events";
import http from "node:http";
import { Readable } from "node:stream";
import { WebSocket, WebSocketServer } from "ws";
import { createUpgradeHandler } from "../../util/http.ts";
import { logger } from "../../util/logger.ts";

export interface MiniflareWorkerProxy {
  url: URL;
  close: () => Promise<void>;
}

export async function createMiniflareWorkerProxy(options: {
  port: number;
  transformRequest?: (request: RequestInfo) => void;
  getWorkerName: (request: RequestInfo) => string;
  getMiniflare: () => Promise<miniflare.Miniflare>;
  mode: "local" | "remote";
}) {
  const server = http.createServer();
  const wss = new WebSocketServer({ noServer: true });

  server.on(
    "upgrade",
    createUpgradeHandler({
      wss,
      createServerWebSocket: (req) => createServerWebSocket(req),
    }),
  );

  server.on("request", async (req, res) => {
    try {
      const response = await handleFetch(req, res);
      writeMiniflareResponseToNode(response, res);
    } catch (error) {
      if (!(error instanceof DOMException && error.name === "AbortError")) {
        console.error(error);
      }
      const response = MiniflareWorkerProxyError.fromUnknown(error).toResponse(
        options.mode,
      );
      writeMiniflareResponseToNode(response, res);
    }
  });

  const handleFetch = async (
    req: http.IncomingMessage,
    res: http.ServerResponse,
  ): Promise<miniflare.Response> => {
    const info = parseIncomingMessage(req);
    options.transformRequest?.(info);

    const name = options.getWorkerName(info);
    info.headers.set("MF-Route-Override", name);

    // Handle scheduled events.
    // Wrangler exposes this as /__scheduled, but Miniflare exposes it as /cdn-cgi/handler/scheduled.
    // https://developers.cloudflare.com/workers/runtime-apis/handlers/scheduled/#background
    // https://github.com/cloudflare/workers-sdk/blob/7d53b9ab6b370944b7934ad51ebef43160c3c775/packages/wrangler/templates/middleware/middleware-scheduled.ts#L6
    if (info.url.pathname === "/__scheduled") {
      info.url.pathname = "/cdn-cgi/handler/scheduled";
    }

    const signal = getDisconnectSignal(req, res);
    const request = new miniflare.Request(info.url, {
      method: info.method,
      headers: info.headers,
      body: info.body,
      redirect: info.redirect,
      duplex: info.duplex,
      signal,
    });
    const mf = await options.getMiniflare();
    return await mf.dispatchFetch(request);
  };

  const createServerWebSocket = async (req: http.IncomingMessage) => {
    // Rewrite to Miniflare entry URL
    const mf = await options.getMiniflare();
    const target = await mf.ready;
    const url = new URL(req.url ?? "/", target);
    url.protocol = url.protocol.replace("http", "ws");

    const protocols = req.headers["sec-websocket-protocol"]
      ?.split(",")
      .map((p) => p.trim());
    const userHeaders = Object.fromEntries(
      // filter headers that are set automatically by `ws` (if we set them manually, this causes errors)
      Object.entries(req.headers).filter(
        ([key]) =>
          !key.startsWith("sec-websocket-") &&
          key !== "host" &&
          key !== "connection" &&
          key !== "upgrade",
      ),
    );

    // Get worker name to set MF-Route-Override header
    const info = parseIncomingMessage(req);
    options.transformRequest?.(info);
    const name = options.getWorkerName(info);

    return new WebSocket(url, protocols, {
      headers: {
        "MF-Route-Override": name,
        ...userHeaders,
      },
    });
  };

  await listen(server, options.port);

  return {
    url: getAddress(server),
    close: async () => {
      // If we await this while the `wss` server is open, the server will hang.
      server.close();
    },
  };
}

/**
 * Creates an abort signal that is triggered when the downstream client
 * disconnects before the response is fully written.
 *
 * This signal should be forwarded to Miniflare dispatch so Worker `req.signal`
 * listeners can observe local client cancellation.
 */
function getDisconnectSignal(
  req: http.IncomingMessage,
  res: http.ServerResponse,
): AbortSignal {
  const controller = new AbortController();
  const onAbort = () => {
    if (!controller.signal.aborted) {
      controller.abort();
    }
    cleanup();
  };
  const onClose = () => {
    if (!res.writableEnded) {
      onAbort();
      return;
    }
    cleanup();
  };
  const cleanup = () => {
    req.off("aborted", onAbort);
    req.socket?.off("close", onClose);
    res.off("close", onClose);
    res.off("finish", cleanup);
  };

  req.on("aborted", onAbort);
  req.socket?.on("close", onClose);
  res.on("close", onClose);
  res.on("finish", cleanup);

  return controller.signal;
}

/**
 * Listens on the given port, retrying on the next port if the port is already in use.
 */
async function listen(server: http.Server, port: number) {
  try {
    server.listen(port);
    await once(server, "listening");
  } catch (error) {
    if (
      typeof error === "object" &&
      error !== null &&
      "code" in error &&
      error.code === "EADDRINUSE"
    ) {
      logger.warn(`Port ${port} is already in use, trying ${port + 1}...`);
      await listen(server, port + 1);
    } else {
      throw error;
    }
  }
}

function getAddress(server: http.Server) {
  const address = server.address();
  assert(
    address &&
      typeof address === "object" &&
      "address" in address &&
      "port" in address,
    "Server is not listening",
  );
  const hostname = address.address === "::" ? "localhost" : address.address;
  return new URL(`http://${hostname}:${address.port}`);
}

interface RequestInfo {
  method: string;
  url: URL;
  headers: miniflare.Headers;
  body: miniflare.BodyInit | undefined;
  redirect: "manual";
  duplex: "half" | undefined;
}

const parseIncomingMessage = (req: http.IncomingMessage): RequestInfo => {
  const method = req.method ?? "GET";
  const protocol = req.headers["x-forwarded-proto"] ?? "http";
  const url = new URL(req.url ?? "/", `${protocol}://${req.headers.host}`);
  const headers = new miniflare.Headers();
  for (const [key, value] of Object.entries(req.headers)) {
    if (Array.isArray(value)) {
      for (const v of value) {
        headers.append(key, v);
      }
    } else if (value) {
      headers.set(key, value);
    }
  }
  const body =
    ["GET", "HEAD", "OPTIONS"].includes(method) || !req.readable
      ? undefined
      : Readable.toWeb(req);
  return {
    method,
    url,
    headers,
    body,
    redirect: "manual",
    duplex: body ? "half" : undefined,
  };
};

const writeMiniflareResponseToNode = (
  res: miniflare.Response,
  out: http.ServerResponse,
) => {
  out.statusCode = res.status;
  res.headers.forEach((value, key) => {
    // The `transfer-encoding` header prevents Cloudflare Tunnels from working because of the following error:
    // Request failed error="Unable to reach the origin service. The service may be down or it may not be
    // responding to traffic from cloudflared: net/http: HTTP/1.x transport connection broken: too many
    // transfer encodings: [\"chunked\" \"chunked\"]"
    if (key !== "transfer-encoding") {
      out.appendHeader(key, value);
    }
  });

  if (res.body) {
    Readable.fromWeb(res.body).pipe(out, { end: true });
  } else {
    out.end();
  }
};

export class MiniflareWorkerProxyError extends Error {
  constructor(
    message: string,
    readonly status: number,
    options?: ErrorOptions,
  ) {
    super(message, options);
  }

  toResponse(mode: "local" | "remote"): miniflare.Response {
    let text = `[Alchemy] ${this.message}`;
    if (mode === "local" && this.cause) {
      const cause =
        this.cause instanceof Error
          ? (this.cause.stack ?? this.cause.message)
          : String(this.cause);
      text += `\n\n${cause}`;
    }
    return new miniflare.Response(text, {
      status: this.status,
      headers: {
        "Content-Type": "text/plain",
        "Alchemy-Error": this.message,
      },
    });
  }

  static fromUnknown(error: unknown): MiniflareWorkerProxyError {
    if (error instanceof MiniflareWorkerProxyError) {
      return error;
    }
    return new MiniflareWorkerProxyError("An unknown error occurred.", 500, {
      cause: error,
    });
  }
}