#!/usr/bin/env node

import { createServer, request as httpRequest } from "node:http";
import path from "node:path";
import process from "node:process";
import { fileURLToPath } from "node:url";
import dotenv from "dotenv";
import WebSocket, { WebSocketServer } from "ws";

const root = path.resolve(path.dirname(fileURLToPath(import.meta.url)), "..");
dotenv.config({ path: path.join(root, ".env.local") });

const host = "127.0.0.1";
const port = Number(process.env.HICAD_ZOO_WS_DEV_PROXY_PORT || 3010);
const appPort = Number(process.env.HICAD_APP_DEV_PORT || 3003);
const apiBaseUrl = String(
  process.env.KITTYCAD_BASE_URL ||
    process.env.ZOO_BASE_URL ||
    "https://api.zoo.dev"
).replace(/\/+$/, "");
const apiKey = String(
  process.env.KITTYCAD_API_TOKEN ||
    process.env.KITTYCAD_API_KEY ||
    process.env.ZOO_API_TOKEN ||
    ""
)
  .replace(/\s+/g, "")
  .trim();

if (!apiKey) {
  throw new Error("Missing KITTYCAD_API_TOKEN in .env.local");
}

const allowedPaths = new Map([
  ["/api/ai-cad/zoo-ws/ml/copilot", "/ws/ml/copilot"],
  ["/api/ai-cad/zoo-ws/modeling/commands", "/ws/modeling/commands"],
]);
const browserSockets = new WebSocketServer({ noServer: true });
const server = createServer((request, response) => {
  const upstreamRequest = httpRequest(
    {
      headers: request.headers,
      host,
      method: request.method,
      path: request.url,
      port: appPort,
    },
    (upstreamResponse) => {
      response.writeHead(
        upstreamResponse.statusCode || 502,
        upstreamResponse.headers
      );
      upstreamResponse.pipe(response);
    }
  );
  upstreamRequest.on("error", (error) => {
    if (!response.headersSent) {
      response.writeHead(502, {
        "Content-Type": "text/plain; charset=utf-8",
      });
    }
    response.end(`HiCAD development proxy failed: ${error.message}`);
  });
  request.pipe(upstreamRequest);
});

server.on("upgrade", (request, socket, head) => {
  const incomingUrl = new URL(request.url || "/", `http://${host}:${port}`);
  const upstreamPath = allowedPaths.get(incomingUrl.pathname);
  if (!upstreamPath) {
    socket.write("HTTP/1.1 404 Not Found\r\nConnection: close\r\n\r\n");
    socket.destroy();
    return;
  }

  browserSockets.handleUpgrade(request, socket, head, (browserSocket) => {
    const base = new URL(apiBaseUrl);
    base.protocol = base.protocol === "http:" ? "ws:" : "wss:";
    base.pathname = upstreamPath;
    base.search = incomingUrl.search;
    const upstreamSocket = new WebSocket(base, {
      headers: {
        Authorization: `Bearer ${apiKey}`,
        Origin: "https://app.zoo.dev",
      },
    });

    const closeBoth = (code = 1011, reason = "Zoo connection closed") => {
      if (browserSocket.readyState < WebSocket.CLOSING) {
        browserSocket.close(code, reason);
      }
      if (upstreamSocket.readyState < WebSocket.CLOSING) {
        upstreamSocket.close(code, reason);
      }
    };

    upstreamSocket.on("open", () => {
      browserSocket.on("message", (data, isBinary) => {
        if (upstreamSocket.readyState === WebSocket.OPEN) {
          upstreamSocket.send(data, { binary: isBinary });
        }
      });
      upstreamSocket.on("message", (data, isBinary) => {
        if (browserSocket.readyState === WebSocket.OPEN) {
          browserSocket.send(data, { binary: isBinary });
        }
      });
    });
    browserSocket.on("close", () => closeBoth(1000, "Browser disconnected"));
    browserSocket.on("error", () => closeBoth());
    upstreamSocket.on("close", () => closeBoth(1000, "Zoo disconnected"));
    upstreamSocket.on("error", (error) => {
      console.error("[HiCAD Zoo WS proxy]", error.message);
      closeBoth();
    });
  });
});

server.listen(port, host, () => {
  console.log(`HiCAD development proxy: http://${host}:${port}`);
});