lab/bastion/src/labd/src/routes/bastions.ts

// Bastion management routes.
// GET  /api/bastions          — list connected bastions
// GET  /api/machines          — aggregated machines from all bastions
// POST /api/machines/install  — queue install on correct bastion
// DELETE /api/machines/:mac   — forget machine on correct bastion
// POST /api/machines/role     — update role on correct bastion
// GET  /api/machines/:mac/logs — get provision logs from correct bastion

import type { FastifyInstance } from "fastify";
import type { DbClient } from "../server.js";
import { bastionRegistry } from "../services/bastion-registry.js";
import { generateRequestId } from "@lab/shared";
import type { VyosInstallSpec } from "@lab/shared";

const COMMAND_TIMEOUT_MS = 15_000;

/** Send a command to a bastion and wait for the response. */
function sendCommand(
  bastionId: string,
  msg: Record<string, unknown>,
): Promise<{ status: string; data?: unknown; error?: string | undefined }> {
  const bastion = bastionRegistry.getById(bastionId);
  if (!bastion) {
    return Promise.reject(new Error(`Bastion ${bastionId} not connected`));
  }

  const requestId = generateRequestId();
  const fullMsg = { ...msg, requestId };

  return new Promise((resolve, reject) => {
    const timeout = setTimeout(() => {
      cleanup();
      reject(new Error("Command timed out"));
    }, COMMAND_TIMEOUT_MS);

    const handler = (data: Buffer) => {
      try {
        const parsed = JSON.parse(data.toString()) as { type: string; requestId?: string; status?: string; data?: unknown; error?: string };
        if (parsed.type === "command-response" && parsed.requestId === requestId) {
          cleanup();
          resolve({ status: parsed.status ?? "ok", data: parsed.data, error: parsed.error });
        }
      } catch { /* not our message */ }
    };

    const cleanup = () => {
      clearTimeout(timeout);
      bastion.socket.off("message", handler);
    };

    bastion.socket.on("message", handler);
    bastion.socket.send(JSON.stringify(fullMsg));
  });
}

export function registerBastionRoutes(app: FastifyInstance, db: DbClient): void {
  // List all bastions (DB records enriched with online status from registry)
  app.get("/api/bastions", async () => {
    const dbBastions = await db.bastion.findMany() as Array<{
      id: string; hostname: string; network: string; serverIp: string;
      status: string; lastHeartbeat: Date | null; createdAt: Date;
    }>;

    return dbBastions.map((b) => {
      const connected = bastionRegistry.getById(b.id);
      return {
        id: b.id,
        hostname: b.hostname,
        network: b.network,
        serverIp: b.serverIp,
        status: connected ? "online" : "offline",
        lastHeartbeat: connected?.lastHeartbeat ?? b.lastHeartbeat,
        connectedAt: connected?.connectedAt,
        machineCount: connected
          ? Object.keys(connected.state.discovered).length +
            Object.keys(connected.state.install_queue).length +
            Object.keys(connected.state.installed).length
          : 0,
        createdAt: b.createdAt,
      };
    });
  });

  // Aggregated machines from all connected bastions + DB fallback
  app.get("/api/machines", async () => {
    const live = bastionRegistry.getAggregatedState();

    try {
      const dbServers = (await db.server.findMany({})) as Array<{
        mac: string | null; hostname: string; role: string; ip: string | null;
        status: string; labels: Record<string, unknown>;
      }>;
      for (const s of dbServers) {
        if (!s.mac) continue;
        const mac = s.mac.toLowerCase();

        // DB knows this machine has been installed at some point if it has a real
        // hostname+role (not just product-name-as-hostname and role="unknown").
        // Status alone is unreliable: a rediscovery can re-set it without erasing the
        // install identity. If the bastion restarted and lost its installed map, the
        // machine will only show up in live.discovered — promote it here so the CLI
        // still sees hostname/role/IP.
        const dbKnowsInstalled =
          s.role !== "unknown" && s.role !== "" &&
          s.hostname !== "" && s.hostname !== s.mac;

        if (dbKnowsInstalled && !(mac in live.installed) && !(mac in live.install_queue)) {
          const hw = live.discovered[mac];
          live.installed[mac] = {
            hostname: s.hostname,
            role: s.role,
            ip: s.ip ?? "",
            installed_at: "",
            bastionId: hw?.bastionId ?? "db",
            ...(hw ? {
              product: hw.product,
              manufacturer: hw.manufacturer,
              cpu_model: hw.cpu_model,
              cpu_cores: hw.cpu_cores,
              memory_gb: hw.memory_gb,
              arch: hw.arch,
            } : {}),
          };
          delete live.discovered[mac];
          continue;
        }

        // Unknown-to-live MAC: fall back to whatever the DB says.
        if (!(mac in live.discovered) && !(mac in live.install_queue) && !(mac in live.installed)) {
          if (s.status === "online" || s.status === "offline") {
            live.installed[mac] = {
              hostname: s.hostname,
              role: s.role,
              ip: s.ip ?? "",
              installed_at: "",
              bastionId: "db",
            };
          } else {
            live.discovered[mac] = {
              mac,
              product: String(s.labels?.product ?? "unknown"),
              board: "unknown",
              serial: "unknown",
              manufacturer: String(s.labels?.manufacturer ?? "unknown"),
              cpu_model: String(s.labels?.cpu ?? "unknown"),
              cpu_cores: Number(s.labels?.cores ?? 0),
              memory_gb: Number(s.labels?.memory_gb ?? 0),
              arch: String(s.labels?.arch ?? "unknown"),
              disks: [],
              nics: [],
              first_seen: "",
              last_seen: "",
              bastionId: "db",
            };
          }
        }
      }
    } catch {
      // DB unavailable — return live state only
    }

    return live;
  });

  // Queue install — route to correct bastion by MAC
  app.post<{
    Body: { mac?: string; hostname?: string; disk?: string; role?: string; os?: string; vyos?: VyosInstallSpec };
  }>("/api/machines/install", async (request, reply) => {
    const { mac, hostname, disk, role, os, vyos } = request.body ?? {};
    if (!mac || !hostname) {
      return reply.code(400).send({ error: "mac and hostname are required" });
    }

    // Find bastion that knows this MAC, or let caller specify
    const bastion = bastionRegistry.findBastionByMac(mac);
    if (!bastion) {
      // If only one bastion is connected, use it
      const all = bastionRegistry.getAll();
      if (all.length === 0) {
        return reply.code(503).send({ error: "No bastions connected" });
      }
      if (all.length === 1) {
        try {
          const result = await sendCommand(all[0]!.bastionId, {
            type: "command-install",
            mac, hostname, disk: disk ?? "", role: role ?? "infra", os: os ?? "fedora-43",
            ...(vyos ? { vyos } : {}),
          });
          return reply.code(result.status === "ok" ? 200 : 500).send(result);
        } catch (err) {
          return reply.code(500).send({ error: err instanceof Error ? err.message : String(err) });
        }
      }
      return reply.code(404).send({ error: `MAC ${mac} not found on any bastion` });
    }

    try {
      const result = await sendCommand(bastion.bastionId, {
        type: "command-install",
        mac, hostname, disk: disk ?? "", role: role ?? "infra", os: os ?? "fedora-43",
        ...(vyos ? { vyos } : {}),
      });
      return reply.code(result.status === "ok" ? 200 : 500).send(result);
    } catch (err) {
      return reply.code(500).send({ error: err instanceof Error ? err.message : String(err) });
    }
  });

  // Register an already-installed machine — route to correct bastion (or single bastion)
  app.post<{
    Body: { mac?: string; hostname?: string; role?: string; ip?: string };
  }>("/api/machines/register", async (request, reply) => {
    const { mac, hostname, role, ip } = request.body ?? {};
    if (!mac || !hostname) {
      return reply.code(400).send({ error: "mac and hostname are required" });
    }

    const normalized = mac.toLowerCase().replace(/-/g, ":");

    // Find bastion that knows this MAC, or use single connected bastion
    const bastion = bastionRegistry.findBastionByMac(normalized);
    const target = bastion ?? (bastionRegistry.getAll().length === 1 ? bastionRegistry.getAll()[0] : null);

    if (!target) {
      const all = bastionRegistry.getAll();
      if (all.length === 0) {
        return reply.code(503).send({ error: "No bastions connected" });
      }
      return reply.code(404).send({ error: `MAC ${normalized} not found on any bastion and multiple bastions connected` });
    }

    try {
      const result = await sendCommand(target.bastionId, {
        type: "command-register",
        mac: normalized,
        hostname,
        role: role ?? "worker",
        ip: ip ?? "",
      });
      return reply.code(result.status === "ok" ? 200 : 500).send(result);
    } catch (err) {
      return reply.code(500).send({ error: err instanceof Error ? err.message : String(err) });
    }
  });

  // Queue debug/rescue mode — route to correct bastion by MAC
  app.post<{
    Body: { mac?: string; pxeBoot?: boolean };
  }>("/api/machines/debug", async (request, reply) => {
    const mac = (request.body?.mac ?? "").toLowerCase().replace(/-/g, ":");
    const pxeBoot = request.body?.pxeBoot ?? false;
    if (!mac) {
      return reply.code(400).send({ error: "mac is required" });
    }

    const bastion = bastionRegistry.findBastionByMac(mac);
    if (!bastion) {
      const all = bastionRegistry.getAll();
      if (all.length === 0) {
        return reply.code(503).send({ error: "No bastions connected" });
      }
      if (all.length === 1) {
        try {
          const result = await sendCommand(all[0]!.bastionId, { type: "command-debug", mac, pxeBoot });
          return reply.code(result.status === "ok" ? 200 : 500).send(result);
        } catch (err) {
          return reply.code(500).send({ error: err instanceof Error ? err.message : String(err) });
        }
      }
      return reply.code(404).send({ error: `MAC ${mac} not found on any bastion` });
    }

    try {
      const result = await sendCommand(bastion.bastionId, { type: "command-debug", mac, pxeBoot });
      return reply.code(result.status === "ok" ? 200 : 500).send(result);
    } catch (err) {
      return reply.code(500).send({ error: err instanceof Error ? err.message : String(err) });
    }
  });

  // Forget machine
  app.delete<{ Params: { mac: string } }>("/api/machines/:mac", async (request, reply) => {
    const mac = request.params.mac.toLowerCase().replace(/-/g, ":");
    const bastion = bastionRegistry.findBastionByMac(mac);
    if (!bastion) {
      return reply.code(404).send({ error: `MAC ${mac} not found on any bastion` });
    }

    try {
      const result = await sendCommand(bastion.bastionId, { type: "command-forget", mac });
      return reply.send(result);
    } catch (err) {
      return reply.code(500).send({ error: err instanceof Error ? err.message : String(err) });
    }
  });

  // Update hardware info (discovery data) for a machine
  app.post<{
    Body: {
      mac?: string; product?: string; board?: string; serial?: string;
      manufacturer?: string; cpu_model?: string; cpu_cores?: number;
      memory_gb?: number; arch?: string;
      disks?: Array<{ name: string; size_gb: number; model: string }>;
      nics?: Array<{ name: string; mac: string; state: string }>;
    };
  }>("/api/machines/discover", async (request, reply) => {
    const data = request.body ?? {};
    const mac = (data.mac ?? "").toLowerCase().replace(/-/g, ":");
    if (!mac) {
      return reply.code(400).send({ error: "mac is required" });
    }

    const bastion = bastionRegistry.findBastionByMac(mac);
    const target = bastion ?? (bastionRegistry.getAll().length === 1 ? bastionRegistry.getAll()[0] : null);

    if (!target) {
      return reply.code(503).send({ error: "No bastion found for this MAC" });
    }

    try {
      const result = await sendCommand(target.bastionId, { type: "command-discover", ...data, mac });
      return reply.code(result.status === "ok" ? 200 : 500).send(result);
    } catch (err) {
      return reply.code(500).send({ error: err instanceof Error ? err.message : String(err) });
    }
  });

  // Update role
  app.post<{
    Body: { mac?: string; role?: string };
  }>("/api/machines/role", async (request, reply) => {
    const { mac, role } = request.body ?? {};
    if (!mac || !role) {
      return reply.code(400).send({ error: "mac and role are required" });
    }

    const normalized = mac.toLowerCase().replace(/-/g, ":");
    const bastion = bastionRegistry.findBastionByMac(normalized);
    if (!bastion) {
      return reply.code(404).send({ error: `MAC ${normalized} not found on any bastion` });
    }

    try {
      const result = await sendCommand(bastion.bastionId, { type: "command-role-update", mac: normalized, role });
      return reply.send(result);
    } catch (err) {
      return reply.code(500).send({ error: err instanceof Error ? err.message : String(err) });
    }
  });

  // Machine logs (snapshot from bastion's state)
  app.get<{ Params: { mac: string } }>("/api/machines/:mac/logs", async (request, reply) => {
    const mac = request.params.mac.toLowerCase().replace(/-/g, ":");
    const bastion = bastionRegistry.findBastionByMac(mac);
    if (!bastion) {
      return reply.code(404).send({ error: `MAC ${mac} not found` });
    }

    const queued = bastion.state.install_queue[mac];
    const installed = bastion.state.installed[mac];

    // Active install takes priority over old installed state (reprovision case)
    if (queued) {
      return {
        mac,
        hostname: queued.hostname,
        status: queued.progress ? "installing" : "queued",
        progress: queued.progress,
        progress_detail: queued.progress_detail,
        progress_at: queued.progress_at,
        role: queued.role,
        os: queued.os,
        log: queued.log,
      };
    }

    if (installed) {
      return {
        mac,
        hostname: installed.hostname,
        status: "installed",
        role: installed.role,
        ip: installed.ip,
        installed_at: installed.installed_at,
      };
    }

    return reply.code(404).send({ error: `MAC ${mac} not found in install queue or installed` });
  });
}