import { Elysia } from "elysia";
import { JwtVerificationError, UnauthorizedError } from "@continuum/ts-core";
import type { AppPrinterState, AppTelemetryMessage } from "@continuum/ts-types";
import { db } from "../../db";
import { farms, printers } from "../../db/schema";
import { visibleFarms, visiblePrinters } from "../../auth";
import { jwtVerifier } from "../../auth/verifier";
import { tokenFromSubprotocol } from "../../auth/ws_subprotocol";
import { maySee, subscribe } from "../../telemetry/broadcast";

/**
 * The operator console's live fleet feed.
 *
 * Read-only and push-only: the browser sends nothing after the handshake,
 * and a frame arriving from it is ignored. Everything that changes a printer
 * goes over HTTP, where the role checks and the audit trail already live -
 * `POST /printers/:id/commands`. A socket that can also issue commands is a
 * second authorization surface to get right for no gain.
 *
 * ## Not /ws/edge/v1
 *
 * That socket authenticates as a *farm*, with an edge JWT that any gateway
 * for that farm holds, and it accepts telemetry. This one authenticates as a
 * *user*, with a Supabase access token, and accepts nothing. Pointing a
 * browser at the gateway socket would mean handing the browser a credential
 * that can drive printers.
 *
 * ## The wire format
 *
 * `continuum/v1/app/telemetry_message.schema.json`, mirrored in
 * `@continuum/ts-types` as `AppTelemetryMessageSchema`. camelCase, matching
 * this backend's HTTP routes rather than the gateway protocol - see that
 * file. The payload is the printer resource itself, so a client applies a
 * push onto the row it already fetched from `GET /printers`.
 */

/**
 * The farms a user can see, resolved once when their socket opens.
 *
 * Once, not per message: a farm's telemetry arrives every few seconds per
 * printer, and re-deriving visibility each time would put a database round
 * trip on every push to every viewer. The cost is staleness in one
 * direction - someone added to an organisation mid-connection does not see
 * its printers until they reconnect, which a page reload does.
 *
 * Staleness in the *other* direction is the one that would matter, and it is
 * bounded by the same thing: someone removed from an organisation keeps
 * receiving its telemetry until their socket drops. That is a real gap worth
 * knowing about. Closing it properly means invalidating open sockets on a
 * membership change, which needs a signal this backend does not have yet; a
 * periodic re-resolve would narrow the window but not close it, so it is
 * written down rather than half-done. Telemetry is also the least sensitive
 * thing in the system - nozzle temperatures, not credentials - and a removed
 * member can be cut off immediately by anything that matters, because every
 * HTTP route re-checks on each request.
 */
async function visibleFarmIds(userId: string): Promise<Set<string>> {
  const rows = await db.select({ id: farms.id }).from(farms).where(visibleFarms(userId));
  return new Set(rows.map((row) => row.id));
}

/**
 * How many printers this user can see, for the `ready` message.
 *
 * Counting farms would be the cheaper query and the wrong answer: a farm
 * with no printers claimed yet is exactly the first-run state this message
 * exists to make legible.
 */
async function countVisiblePrinters(userId: string): Promise<number> {
  const rows = await db.select({ id: printers.id }).from(printers).where(visiblePrinters(userId));
  return rows.length;
}

/** One unsubscribe per open socket, removed when it closes. */
const unsubscribers = new Map<string | number, () => void>();

export const telemetryWsRoute = new Elysia()
  /**
   * Verifies the user's access token on the upgrade, before a socket exists.
   *
   * Throwing fails the upgrade with an HTTP status, so an unauthenticated
   * browser never reaches `open`. Same shape as the edge socket's resolve
   * next door; the difference is only where the token comes from and which
   * verifier checks it.
   *
   * The default `local` scope deliberately - see the note on the edge
   * route's resolve. `scoped` would push this check onto every route
   * mounted after this plugin, which would then demand a WebSocket
   * subprotocol from plain HTTP requests.
   */
  .resolve(async ({ headers }): Promise<{ viewer: { id: string } }> => {
    const token = tokenFromSubprotocol(headers["sec-websocket-protocol"]);

    try {
      const claims = await jwtVerifier.verify(token);
      if (!claims.sub) throw new UnauthorizedError("Access token has no subject claim");
      return { viewer: { id: claims.sub } };
    } catch (err) {
      // Never pass the reason through - "expired" versus "bad signature"
      // tells a forger which part to fix next. Logged instead, exactly as
      // `requireUser` does it.
      if (err instanceof JwtVerificationError) {
        console.warn(
          `[telemetry] token rejected: ${err.cause instanceof Error ? err.cause.message : err.message}`,
        );
        throw new UnauthorizedError("Invalid or expired access token");
      }
      throw err;
    }
  })
  .ws("/ws/telemetry", {
    async open(ws) {
      const { viewer } = ws.data as { viewer: { id: string } };
      const allowed = await visibleFarmIds(viewer.id);

      const send = (message: AppTelemetryMessage) => ws.send(message);

      const unsubscribe = subscribe((printer: AppPrinterState) => {
        // The line that keeps one organisation's telemetry out of another's
        // console. `maySee` is in `telemetry/broadcast.ts` so it can be
        // tested without standing up a socket.
        if (!maySee(allowed, printer)) return;
        send({ type: "printer.telemetry", printer });
      });

      // Cleanup is attached to the socket rather than kept in a module-level
      // map: there is exactly one unsubscribe per socket and the close
      // handler is the only thing that needs it. A leaked subscriber would
      // push into a dead socket forever.
      unsubscribers.set(ws.id, unsubscribe);

      // Sent before any telemetry, so a client knows the difference between
      // "nothing to show you" and "not connected yet". Counting farms the
      // user can see is not the same as counting printers, so this asks for
      // the printers.
      send({ type: "ready", visiblePrinters: await countVisiblePrinters(viewer.id) });
    },

    message(_ws, _message) {
      // Deliberately nothing. This socket is push-only; see the module
      // comment. Ignoring rather than closing, because a well-meaning client
      // sending a keepalive should not be punished for it.
    },

    close(ws) {
      unsubscribers.get(ws.id)?.();
      unsubscribers.delete(ws.id);
    },
  });