import http from "node:http";
import { randomUUID } from "node:crypto";
import { WebSocketServer } from "ws";

function parseJsonBody(req) {
  return new Promise((resolve, reject) => {
    let raw = "";
    req.on("data", (chunk) => {
      raw += chunk.toString("utf8");
    });
    req.on("end", () => {
      try {
        resolve(raw === "" ? {} : JSON.parse(raw));
      } catch (error) {
        reject(error);
      }
    });
    req.on("error", reject);
  });
}

function jsonResponse(res, statusCode, body) {
  res.writeHead(statusCode, { "content-type": "application/json" });
  res.end(JSON.stringify(body));
}

export function createBrokerServer(options = {}) {
  const authMode = options.authMode || process.env.EDGE_AUTH_MODE || "strict";
  const sharedSecret = options.sharedSecret ?? process.env.EDGE_BROKER_SHARED_SECRET ?? "";
  const commandTimeoutMs = options.commandTimeoutMs ?? 10000;
  const agents = new Map();
  const pendingCommands = new Map();
  const browserSessions = new Map();
  const isStubMode = authMode === "stub";

  const validateAgent = options.validateAgent || (isStubMode
    ? (async ({ gatewayId }) => ({ id: gatewayId }))
    : (async () => {
      throw new Error("Unauthorized");
    }));
  const validateShellSession = options.validateShellSession || (isStubMode
    ? (async ({ token }) => ({ id: token, gateway_id: 1, reason: "stub" }))
    : (async () => {
      throw new Error("Unauthorized");
    }));
  const closeShellSession = options.closeShellSession || (async () => ({}));

  const markBrowserSessionsClosed = (gatewayId, reason) => {
    for (const sessionRecord of browserSessions.values()) {
      if (String(sessionRecord.session.gateway_id) !== String(gatewayId)) {
        continue;
      }

      sessionRecord.closedReason = reason;

      if (sessionRecord.ws.readyState < 2) {
        sessionRecord.ws.close();
      }
    }
  };

  const server = http.createServer(async (req, res) => {
    try {
      const url = new URL(req.url, "http://localhost");
      if (req.method === "POST" && /^\/api\/gateways\/\d+\/commands$/.test(url.pathname)) {
        if (sharedSecret && req.headers["x-edge-broker-secret"] !== sharedSecret) {
          jsonResponse(res, 403, { error: "Forbidden" });
          return;
        }

        const gatewayId = url.pathname.split("/")[3];
        const agent = agents.get(String(gatewayId));
        if (!agent || agent.readyState !== 1) {
          jsonResponse(res, 503, { error: "Gateway agent is offline" });
          return;
        }

        const body = await parseJsonBody(req);
        const commandId = randomUUID();
        const promise = new Promise((resolve, reject) => {
          const timeout = setTimeout(() => {
            pendingCommands.delete(commandId);
            reject(new Error("Agent command timed out"));
          }, commandTimeoutMs);
          pendingCommands.set(commandId, {
            resolve,
            reject,
            timeout,
          });
        });

        agent.send(JSON.stringify({
          type: "COMMAND",
          commandId,
          commandType: body.commandType,
          payload: body.payload || {},
        }));

        try {
          const result = await promise;
          jsonResponse(res, 200, result);
        } catch (error) {
          jsonResponse(res, 504, { ok: false, error: error instanceof Error ? error.message : String(error) });
        }
        return;
      }

      jsonResponse(res, 404, { error: "Not found" });
    } catch (error) {
      jsonResponse(res, 500, { error: error instanceof Error ? error.message : String(error) });
    }
  });

  const wss = new WebSocketServer({ noServer: true });
  server.on("upgrade", async (req, socket, head) => {
    const url = new URL(req.url, "http://localhost");
    try {
      if (url.pathname === "/ws/agent") {
        const gatewayId = String(url.searchParams.get("gatewayId") || "");
        const token = String(url.searchParams.get("token") || "");
        if (gatewayId === "" || token === "") {
          socket.destroy();
          return;
        }
        await validateAgent({ gatewayId, token, headers: req.headers });
        wss.handleUpgrade(req, socket, head, (ws) => {
          ws.gatewayId = gatewayId;
          agents.set(gatewayId, ws);
          wss.emit("connection", ws, req);
        });
        return;
      }

      if (url.pathname === "/ws/browser-shell") {
        const token = String(url.searchParams.get("token") || "");
        if (token === "") {
          socket.destroy();
          return;
        }
        const session = await validateShellSession({ token, headers: req.headers });

        wss.handleUpgrade(req, socket, head, (ws) => {
          ws.sessionToken = token;
          ws.sessionInfo = session;
          browserSessions.set(String(session.id), {
            ws,
            session,
            transcript: "",
            closedReason: null,
          });
          const agent = agents.get(String(session.gateway_id));
          if (agent && agent.readyState === 1) {
            agent.send(JSON.stringify({
              type: "OPEN_ROOT_SHELL",
              payload: {
                sessionId: String(session.id),
                reason: session.reason,
              },
            }));
          } else {
            const sessionRecord = browserSessions.get(String(session.id));
            if (sessionRecord) {
              sessionRecord.closedReason = "agent_offline";
            }
            ws.close();
          }
          wss.emit("connection", ws, req);
        });
        return;
      }
    } catch {
      socket.destroy();
      return;
    }

    socket.destroy();
  });

  wss.on("connection", (ws) => {
    ws.on("message", async (raw) => {
      const message = JSON.parse(raw.toString());

      if (ws.gatewayId) {
        if (message.type === "COMMAND_RESULT") {
          const pending = pendingCommands.get(message.commandId);
          if (!pending) {
            return;
          }
          clearTimeout(pending.timeout);
          pendingCommands.delete(message.commandId);
          pending.resolve({
            ok: Boolean(message.ok),
            payload: message.payload,
            error: message.error,
          });
          return;
        }

        if (["SHELL_OUTPUT", "SHELL_OPENED", "SHELL_EXIT"].includes(message.type)) {
          const sessionRecord = browserSessions.get(String(message.sessionId));
          if (!sessionRecord) {
            return;
          }
          if (message.type === "SHELL_OUTPUT") {
            sessionRecord.transcript += String(message.data || "");
            sessionRecord.ws.send(JSON.stringify({ type: "output", data: String(message.data || "") }));
          }
          if (message.type === "SHELL_OPENED") {
            sessionRecord.ws.send(JSON.stringify({ type: "opened" }));
          }
          if (message.type === "SHELL_EXIT") {
            sessionRecord.ws.send(JSON.stringify({ type: "closed", code: message.code ?? 0 }));
            await closeShellSession(sessionRecord.session.id, sessionRecord.ws.sessionToken, sessionRecord.transcript, "agent_exit");
            browserSessions.delete(String(message.sessionId));
            if (sessionRecord.ws.readyState < 2) {
              sessionRecord.ws.close();
            }
          }
        }
        return;
      }

      if (ws.sessionInfo) {
        const sessionId = String(ws.sessionInfo.id);
        const agent = agents.get(String(ws.sessionInfo.gateway_id));
        if (!agent || agent.readyState !== 1) {
          return;
        }
        if (message.type === "input") {
          agent.send(JSON.stringify({
            type: "SHELL_INPUT",
            payload: {
              sessionId,
              data: String(message.data || ""),
            },
          }));
        }
        if (message.type === "resize") {
          agent.send(JSON.stringify({
            type: "RESIZE_ROOT_SHELL",
            payload: {
              sessionId,
              cols: Number(message.cols || 0),
              rows: Number(message.rows || 0),
            },
          }));
        }
        if (message.type === "close") {
          agent.send(JSON.stringify({
            type: "CLOSE_ROOT_SHELL",
            payload: { sessionId },
          }));
        }
      }
    });

    ws.on("close", async () => {
      if (ws.gatewayId) {
        agents.delete(String(ws.gatewayId));
        markBrowserSessionsClosed(String(ws.gatewayId), "agent_disconnected");
        return;
      }

      if (ws.sessionInfo) {
        const sessionId = String(ws.sessionInfo.id);
        const agent = agents.get(String(ws.sessionInfo.gateway_id));
        if (agent && agent.readyState === 1) {
          agent.send(JSON.stringify({
            type: "CLOSE_ROOT_SHELL",
            payload: { sessionId },
          }));
        }
        const sessionRecord = browserSessions.get(sessionId);
        if (sessionRecord) {
          await closeShellSession(
            sessionRecord.session.id,
            ws.sessionToken,
            sessionRecord.transcript,
            sessionRecord.closedReason || "browser_closed"
          );
          browserSessions.delete(sessionId);
        }
      }
    });
  });

  return {
    server,
    listen(port = Number(process.env.PORT || 4300)) {
      return new Promise((resolve) => {
        server.listen(port, () => resolve(server.address()));
      });
    },
    close() {
      return new Promise((resolve, reject) => {
        for (const agent of agents.values()) {
          agent.terminate();
        }
        for (const session of browserSessions.values()) {
          session.ws.terminate();
        }
        for (const pending of pendingCommands.values()) {
          clearTimeout(pending.timeout);
          pending.reject(new Error("Broker shutting down"));
        }
        pendingCommands.clear();

        wss.close(() => {
          server.close((error) => {
            if (error) {
              reject(error);
              return;
            }
            resolve();
          });
        });
      });
    },
    state: {
      agents,
      browserSessions,
      pendingCommands,
    },
  };
}

if (import.meta.url === `file://${process.argv[1]}`) {
  const broker = createBrokerServer();
  broker.listen().then(() => {
    console.log("TruckWash edge broker listening");
  }).catch((error) => {
    console.error(error instanceof Error ? error.message : String(error));
    process.exitCode = 1;
  });
}