import { parseHeaderObject, ProxyError, readBareHeaders, setBareHeaders, validateTarget } from "./policy";

const MAX_FRAME = 4 * 1024 * 1024;
export async function relayHttp(request: Request, allowed: string): Promise<Response> {
  const target = validateTarget(request.headers.get("x-bare-url") ?? "", allowed);
  const headers = readBareHeaders(request.headers);
  // Only these ambient browser headers are eligible for forwarding.
  for (const name of ["accept-language", "user-agent"]) {
    if (!headers.has(name) && request.headers.has(name)) headers.set(name, request.headers.get(name)!);
  }
  headers.set("accept-encoding", "identity");
  const upstream = await fetch(target, {
    method: request.method, headers, body: request.body,
    redirect: "manual", signal: AbortSignal.timeout(30_000),
  });
  if (upstream.status >= 400) console.warn(JSON.stringify({
    event:"upstream_http_error", host:target.hostname, status:upstream.status,
  }));
  const output = new Headers({
    "content-type": "application/octet-stream",
    "cache-control": "no-store",
    "x-content-type-options": "nosniff",
    "x-bare-status": String(upstream.status),
    "x-bare-status-text": upstream.statusText,
  });
  setBareHeaders(output, upstream.headers);
  // Fetch/Response must agree about encoded bytes; do not transform the streamed body.
  for (const name of ["content-encoding", "content-length", "last-modified"]) {
    const value = upstream.headers.get(name);
    if (value !== null) output.set(name, value);
  }
  return new Response(request.method === "HEAD" ? null : upstream.body, {
    status: 200, headers: output, encodeBody: "manual",
  });
}

type Handshake = {type: "connect"; remote: string; protocols: string[]; headers: unknown};
function handshake(data: string | ArrayBuffer, allowed: string): {message: Handshake; target: URL; headers: Headers} {
  if (typeof data !== "string" || data.length > 64 * 1024) throw new Error("Invalid handshake.");
  const value = JSON.parse(data) as Partial<Handshake>;
  if (value.type !== "connect" || typeof value.remote !== "string" || !Array.isArray(value.protocols)
      || value.protocols.length > 16 || value.protocols.some(p => typeof p !== "string" || !/^[!#$%&'*+.^_`|~0-9A-Za-z-]{1,128}$/.test(p))) {
    throw new Error("Invalid handshake.");
  }
  const target = validateTarget(value.remote, allowed, true);
  const headers = parseHeaderObject(value.headers);
  headers.set("upgrade", "websocket");
  if (value.protocols.length) headers.set("sec-websocket-protocol", value.protocols.join(", "));
  return {message: value as Handshake, target, headers};
}
export function relayWebSocket(request: Request, allowed: string, ctx: ExecutionContext): Response {
  if (request.method !== "GET") throw new ProxyError(405, "METHOD_NOT_ALLOWED", "WebSocket upgrades require GET.");
  const pair = new WebSocketPair();
  const client = pair[0], server = pair[1];
  server.binaryType = "arraybuffer";
  server.accept();
  let upstream: WebSocket | null = null;
  let state: "waiting" | "connecting" | "open" | "closed" = "waiting";
  const isClosed = () => state === "closed";
  const close = (code = 1000, reason = "") => {
    if (state === "closed") return;
    state = "closed";
    clearTimeout(timeout);
    try { server.close(code, reason); } catch { /* already closed */ }
    try { upstream?.close(code, reason); } catch { /* already closed */ }
  };
  const timeout = setTimeout(() => close(1008, "Handshake timed out."), 15_000);
  server.addEventListener("close", () => close());
  server.addEventListener("error", () => close(1011, "Tunnel error."));
  server.addEventListener("message", event => {
    if (state === "closed") return;
    const size = typeof event.data === "string" ? event.data.length : event.data.byteLength;
    if (size > MAX_FRAME) { close(1009, "Frame too large."); return; }
    if (state === "open") {
      try { upstream!.send(event.data); } catch { close(1011, "Upstream unavailable."); }
      return;
    }
    if (state !== "waiting") { close(1008, "Wait for the open message."); return; }
    state = "connecting";
    ctx.waitUntil((async () => {
      try {
        const {target, headers} = handshake(event.data, allowed);
        // Workers use HTTPS + Upgrade to open an outgoing secure WebSocket.
        target.protocol = "https:";
        // A handshake deadline must not abort the upgraded socket ten seconds later.
        const controller = new AbortController();
        const connectingTimeout = setTimeout(() => controller.abort(),10_000);
        let response: Response;
        try {
          response = await fetch(target, {headers, redirect: "manual", signal: controller.signal});
        } finally { clearTimeout(connectingTimeout); }
        if (response.status !== 101 || !response.webSocket) throw new Error("Upgrade rejected.");
        const socket = response.webSocket;
        socket.binaryType = "arraybuffer";
        upstream = socket;
        if (isClosed()) { socket.accept(); socket.close(); return; }
        socket.addEventListener("message", event => {
          if (state !== "open") return;
          try { server.send(event.data); } catch { close(); }
        });
        socket.addEventListener("close", () => close());
        socket.addEventListener("error", () => close(1011, "Upstream error."));
        // Metadata goes first, before accepting incoming upstream messages.
        server.send(JSON.stringify({
          type: "open", protocol: response.headers.get("sec-websocket-protocol") ?? "",
          setCookies: response.headers.getSetCookie(),
        }));
        state = "open";
        clearTimeout(timeout);
        socket.accept();
      } catch {
        close(1008, "Destination blocked or connection failed.");
      }
    })());
  });
  return new Response(null, {status: 101, webSocket: client});
}