import { randomUUID } from "node:crypto";
import { WebSocketServer, type WebSocket } from "ws";
import { formatComputer, parseJsonFrame, parseLinkMessage, type ResponseMessage } from "./protocol.js";

export type ComputerConnection = {
  computerId: number;
  label: string | null;
  ws: WebSocket;
  connectedAt: number;
  lastSeenAt: number;
};

type PendingRequest = {
  computerId: number;
  resolve: (message: ResponseMessage) => void;
};

export class LinkRegistry {
  private readonly computers = new Map<number, ComputerConnection>();
  private readonly pending = new Map<string, PendingRequest>();

  register(connection: ComputerConnection): void {
    const existing = this.computers.get(connection.computerId);
    if (existing && existing.ws !== connection.ws) {
      existing.ws.close();
    }
    this.computers.set(connection.computerId, connection);
  }

  unregister(ws: WebSocket): void {
    for (const [computerId, connection] of this.computers) {
      if (connection.ws === ws) {
        this.computers.delete(computerId);
      }
    }
  }

  count(): number {
    return this.computers.size;
  }

  snapshot(): ComputerConnection[] {
    return [...this.computers.values()];
  }

  handleFrame(ws: WebSocket, data: unknown): void {
    const message = parseLinkMessage(parseJsonFrame(data));
    if (!message) {
      return;
    }

    if (message.type === "hello") {
      const now = Date.now();
      this.register({
        computerId: message.computerId,
        label: message.computerLabel,
        ws,
        connectedAt: now,
        lastSeenAt: now,
      });
      ws.send(JSON.stringify({ type: "hello-ok" }));
      return;
    }

    const pending = this.pending.get(message.id);
    if (!pending) {
      return;
    }

    const connection = this.computers.get(pending.computerId);
    if (connection) {
      connection.lastSeenAt = Date.now();
    }
    this.pending.delete(message.id);
    pending.resolve(message);
  }

  async probeComputers(timeoutMs: number): Promise<string> {
    const computers = this.snapshot();
    if (computers.length === 0) {
      return "No computers connected.";
    }

    const lines = await Promise.all(computers.map((computer) => this.probeComputer(computer, timeoutMs)));
    return lines.join("\n");
  }

  private async probeComputer(computer: ComputerConnection, timeoutMs: number): Promise<string> {
    const id = randomUUID();
    const response = new Promise<ResponseMessage>((resolve) => {
      this.pending.set(id, { computerId: computer.computerId, resolve });
      computer.ws.send(JSON.stringify({ type: "request", id, method: "ping" }));
    });

    const timeout = new Promise<null>((resolve) => {
      setTimeout(() => {
        this.pending.delete(id);
        resolve(null);
      }, timeoutMs);
    });

    const result = await Promise.race([response, timeout]);
    if (!result) {
      return `timeout from ${formatComputer(computer.computerId, computer.label)}`;
    }

    if (result.ok && typeof result.result === "string") {
      return result.result;
    }

    return `error from ${formatComputer(computer.computerId, computer.label)}: ${result.error ?? "unknown error"}`;
  }
}

export function startLinkServer(options: { host: string; port: number; registry: LinkRegistry }): WebSocketServer {
  const server = new WebSocketServer({ host: options.host, port: options.port });

  server.on("connection", (ws) => {
    ws.on("message", (data) => options.registry.handleFrame(ws, data));
    ws.on("close", () => options.registry.unregister(ws));
    ws.on("error", () => options.registry.unregister(ws));
  });

  return server;
}