import { fetchHeld, RESPONSE_BODY_LIMIT_CANCELLATION } from "./http.js";

const HOP_BY_HOP = new Set([
  "connection",
  "keep-alive",
  "proxy-authenticate",
  "proxy-authorization",
  "proxy-connection",
  "te",
  "trailer",
  "transfer-encoding",
  "upgrade",
  "host",
]);
const FIELD_NAME = /^[!#$%&'*+\-.^_`|~0-9a-z]+$/;

const hopByHopHeaders = (headers: Headers): Set<string> => {
  const names = new Set(HOP_BY_HOP);
  for (const raw of (headers.get("connection") ?? "").split(",")) {
    const name = raw.trim().toLowerCase();
    if (FIELD_NAME.test(name)) names.add(name);
  }
  return names;
};

export const PROXY_SECRET_HEADER = "pinboard-proxy-secret";

export const isReservedPinboardHeader = (name: string): boolean =>
  name.toLowerCase().startsWith("pinboard-");

const stripReservedPinboardHeaders = (headers: Headers): boolean => {
  const reserved: string[] = [];
  headers.forEach((_value, name) => {
    if (isReservedPinboardHeader(name)) reserved.push(name);
  });
  for (const name of reserved) headers.delete(name);
  return reserved.length > 0;
};

export const proxySecretHeaders = (secret: string): Record<string, string> => ({
  [PROXY_SECRET_HEADER]: secret,
});

/** In-band placement NACK: a pinned worker marks responses that refused to
 * start a new activation; stripped before anything reaches the client. */
export const REFUSAL_HEADER = "pinned-refusal";
export const LEASE_UNHEALTHY_REFUSAL = "lease-unhealthy";

export const FROM_HEADER = "pinboard-from";
/** Placement epoch of the row a forward resolved against. */
export const EPOCH_HEADER = "pinboard-epoch";
export const SESSION_HEADER = "pinboard-session";
export const SESSION_SECRET_HEADER = "pinboard-session-secret";
export const DEPTH_HEADER = "pinboard-subrequest-depth";

// Node socket error codes plus Bun's names for the same connect failures.
// UND_ERR_CONNECT_TIMEOUT: a timeout before TCP established never reached the
// worker, unlike a post-connect timeout.
const CONNECT_ERROR_CODES = new Set([
  "ECONNREFUSED",
  "EHOSTUNREACH",
  "ENETUNREACH",
  "ConnectionRefused",
  "FailedToOpenSocket",
  "UND_ERR_CONNECT_TIMEOUT",
]);

const TIMEOUT_ERROR_NAMES = new Set([
  "HeadersTimeoutError",
  "ConnectTimeoutError",
]);

const hasConnectCode = (err: unknown): boolean => {
  if (!(err instanceof Error)) return false;
  const code = (err as { code?: unknown }).code;
  return typeof code === "string" && CONNECT_ERROR_CODES.has(code);
};

const causes = function* (err: unknown): Generator<Error> {
  let current: unknown = err;
  while (current instanceof Error) {
    yield current;
    if (current instanceof AggregateError) yield* current.errors;
    current = current.cause;
  }
};

export const isConnectError = (err: unknown): boolean => {
  for (const error of causes(err)) {
    if (hasConnectCode(error)) return true;
  }
  return false;
};

export const isTimeoutError = (err: unknown): boolean => {
  for (const error of causes(err)) {
    if (TIMEOUT_ERROR_NAMES.has(error.name)) return true;
  }
  return false;
};

export type ForwardOutcome = "forwarded" | "aborted" | "unreachable" | "failed";

export namespace Transport {
  export type FetchLike = (request: Request) => Promise<Response>;

  export type UpgradeHooks = {
    onFailure(outcome: Exclude<ForwardOutcome, "forwarded">): void;
    onClose(): void;
  };

  /** Performs the upgrade forward; the returned Response completes the
   * client's upgrade (on Workers, a 101 carrying `webSocket`). */
  export type Upgrade = (
    req: Request,
    outbound: { url: string; headers: Headers },
    hooks: UpgradeHooks,
  ) => Promise<Response>;

  export type Options = {
    fetch?: FetchLike | undefined;
    upgrade?: Upgrade | undefined;
  };
}

/** Outbound headers: hop-by-hop and reserved Pinboard-* stripped, the stamp
 * applied, X-Forwarded-For reduced to the observed peer address. */
export const outboundHeaders = (
  req: Request,
  address: string | undefined,
  stamp: Record<string, string>,
): Headers => {
  const headers = new Headers();
  const hopByHop = hopByHopHeaders(req.headers);
  req.headers.forEach((value, name) => {
    if (hopByHop.has(name)) return;
    // The Pinboard-* prefix is reserved for proxy-attested headers.
    if (isReservedPinboardHeader(name)) return;
    headers.set(name, value);
  });
  // The client-supplied chain is spoofable; only the observed remote rides.
  headers.delete("x-forwarded-for");
  if (address !== undefined) headers.set("x-forwarded-for", address);
  for (const [name, value] of Object.entries(stamp)) headers.set(name, value);
  return headers;
};

const trackedBody = (
  upstream: ReadableStream<Uint8Array>,
  signal: AbortSignal,
  settle: (outcome: ForwardOutcome) => void,
): ReadableStream<Uint8Array> => {
  const reader = upstream.getReader();
  return new ReadableStream({
    async pull(controller) {
      let result: Awaited<ReturnType<typeof reader.read>>;
      try {
        result = await reader.read();
      } catch (err) {
        settle(signal.aborted ? "aborted" : "failed");
        controller.error(err);
        return;
      }
      if (result.done) {
        settle("forwarded");
        controller.close();
      } else {
        controller.enqueue(result.value);
      }
    },
    async cancel(reason) {
      settle(
        reason === RESPONSE_BODY_LIMIT_CANCELLATION ? "forwarded" : "aborted",
      );
      await reader.cancel(reason).catch(() => {});
    },
  });
};

export const forward = async (
  fetchImpl: Transport.FetchLike,
  req: Request,
  target: string,
  headers: Headers,
  onOutcome: (outcome: ForwardOutcome) => void,
  /** Cap on the wait for upstream response headers; streaming bodies are
   * never cut. Exceeding answers 504. */
  timeoutMs?: number,
): Promise<Response> => {
  let settled = false;
  const settle = (outcome: ForwardOutcome): void => {
    if (settled) return;
    settled = true;
    onOutcome(outcome);
  };

  if (req.signal.aborted) {
    settle("aborted");
    return new Response(null, { status: 502 });
  }

  let signal = req.signal;
  let timedOut = false;
  let timer: ReturnType<typeof setTimeout> | undefined;
  if (timeoutMs !== undefined) {
    const controller = new AbortController();
    timer = setTimeout(() => {
      timedOut = true;
      controller.abort();
    }, timeoutMs);
    req.signal.addEventListener("abort", () => controller.abort(), {
      once: true,
    });
    signal = controller.signal;
  }

  let upstream: Response;
  try {
    upstream = await fetchHeld(
      fetchImpl,
      new Request(target, {
        method: req.method,
        headers,
        ...(req.body !== null && {
          body: req.body,
          duplex: "half",
        }),
        redirect: "manual",
        signal,
      } as RequestInit),
    );
  } catch (err) {
    if (req.signal.aborted) {
      settle("aborted");
      return new Response(null, { status: 502 });
    }
    settle(isConnectError(err) ? "unreachable" : "failed");
    const upstreamTimeout = timedOut || isTimeoutError(err);
    return new Response(upstreamTimeout ? "upstream timeout" : "bad gateway", {
      status: upstreamTimeout ? 504 : 502,
      headers: { "content-type": "text/plain" },
    });
  } finally {
    clearTimeout(timer);
  }

  const responseHeaders = new Headers(upstream.headers);
  for (const name of hopByHopHeaders(responseHeaders)) {
    responseHeaders.delete(name);
  }
  stripReservedPinboardHeaders(responseHeaders);
  responseHeaders.set("x-accel-buffering", "no");
  if (upstream.body === null) {
    settle("forwarded");
    return new Response(null, {
      status: upstream.status,
      statusText: upstream.statusText,
      headers: responseHeaders,
    });
  }
  return new Response(trackedBody(upstream.body, req.signal, settle), {
    status: upstream.status,
    statusText: upstream.statusText,
    headers: responseHeaders,
  });
};

/** Workers-style upgrade: the platform fetch returns the 101 Response
 * (carrying the socket) directly. */
export const fetchUpgrade =
  (fetchImpl: Transport.FetchLike): Transport.Upgrade =>
  async (req, outbound, hooks) => {
    outbound.headers.set("connection", "Upgrade");
    outbound.headers.set("upgrade", "websocket");
    try {
      const upstream = await fetchImpl(
        new Request(outbound.url, {
          method: req.method,
          headers: outbound.headers,
          redirect: "manual",
        }),
      );
      const headers = new Headers(upstream.headers);
      if (!stripReservedPinboardHeaders(headers)) return upstream;
      const platform = upstream as Response & {
        readonly webSocket?: unknown;
        readonly cf?: unknown;
      };
      return new Response(upstream.body, {
        status: upstream.status,
        statusText: upstream.statusText,
        headers,
        ...(platform.webSocket !== undefined && {
          webSocket: platform.webSocket,
        }),
        ...(platform.cf !== undefined && { cf: platform.cf }),
      } as ResponseInit);
    } catch (err) {
      hooks.onFailure(isConnectError(err) ? "unreachable" : "failed");
      hooks.onClose();
      return new Response("bad gateway", {
        status: 502,
        headers: { "content-type": "text/plain" },
      });
    }
  };