import { spawn } from "node:child_process";
import { randomUUID } from "node:crypto";
import type { IncomingMessage } from "node:http";
import { WebSocket, WebSocketServer } from "ws";
import { Config, VERSION } from "./config";
import { lastLaunchAt, profileDir, profileInUse, recordLaunch } from "./profile";
import { ReplayGuard, seal, unseal } from "./protocol";

export class BrowserUnavailable extends Error {
  // True when the browser was just started and is still loading: worth one more try.
  constructor(message: string, readonly starting = false) {
    super(message);
  }
}

interface Pending {
  resolve: (value: unknown) => void;
  reject: (err: Error) => void;
  timer: NodeJS.Timeout;
}

const AUTH_ERROR_MEMORY_MS = 30_000;
const RELAUNCH_COOLDOWN_MS = 90_000;
// Kept under LM Studio's limit for a single tool call, which counts the search after it too.
const CONNECT_WAIT_MS = 40_000;
const REPEAT_FAIL_MS = 60_000;

// Owns the local WebSocket that the browser extension connects to, and starts the
// browser profile when a request arrives and no extension is connected.
export class Bridge {
  private wss: WebSocketServer | null = null;
  private sock: WebSocket | null = null;
  private guard = new ReplayGuard();
  private pending = new Map<string, Pending>();
  private waiters: Array<() => void> = [];
  private lastLaunch = 0;
  private authErrorAt = 0;
  private lastAuthLog = 0;
  // The last time a call gave up waiting for the extension, and why. A call soon after that
  // fails within seconds instead of waiting the full time again for the same reason.
  private lastFailAt = 0;
  private lastFail = "";

  bindError: string | null = null;
  launchError: string | null = null;
  authError: string | null = null;
  extensionVersion: string | null = null;

  constructor(private cfg: Config, private log: (msg: string) => void) {}

  start(): Promise<void> {
    return new Promise((resolve) => {
      const wss = new WebSocketServer({ host: "127.0.0.1", port: this.cfg.port, maxPayload: 32 * 1024 * 1024 });
      wss.on("listening", () => {
        this.wss = wss;
        this.log(`listening on 127.0.0.1:${this.cfg.port}`);
        resolve();
      });
      wss.on("error", (err: NodeJS.ErrnoException) => {
        if (this.wss) {
          this.log(`server error: ${err.message}`);
          return;
        }
        this.bindError =
          err.code === "EADDRINUSE"
            ? `Port ${this.cfg.port} is already in use, ${this.cfg.host.portInUse}.`
            : `Could not open port ${this.cfg.port}: ${err.message}`;
        this.log(this.bindError);
        resolve();
      });
      wss.on("connection", (socket, req) => this.onConnection(socket, req));
    });
  }

  // Resolves once the port is free again, so a new bridge can open it straight away.
  close(): Promise<void> {
    for (const p of this.pending.values()) {
      clearTimeout(p.timer);
      p.reject(new Error("The server is shutting down."));
    }
    this.pending.clear();
    this.sock?.terminate();
    this.sock = null;
    const wss = this.wss;
    this.wss = null;
    if (!wss) return Promise.resolve();
    return new Promise((resolve) => {
      const done = setTimeout(resolve, 2000);
      wss.close(() => {
        clearTimeout(done);
        resolve();
      });
    });
  }

  isListening(): boolean {
    return this.wss !== null;
  }

  isConnected(): boolean {
    return this.sock !== null && this.sock.readyState === WebSocket.OPEN;
  }

  private onConnection(socket: WebSocket, req: IncomingMessage): void {
    // Web pages can open WebSockets to localhost too. Only the extension is welcome.
    const origin = req.headers.origin ?? "";
    if (/^https?:/i.test(origin)) {
      socket.close(1008, "forbidden");
      return;
    }
    let authed = false;
    const helloTimer = setTimeout(() => {
      if (!authed) socket.close(1008, "no hello");
    }, 5000);

    socket.on("message", (data) => {
      const msg = unseal(this.cfg.secret, data.toString(), this.guard);
      if (!msg) {
        if (!authed) {
          this.authError =
            `The browser extension connected, but its secret does not match ${this.cfg.host.secretSetting}. Copy the secret from the extension's options page into ${this.cfg.host.secretSetting}${this.cfg.host.afterChange}.`;
          this.authErrorAt = Date.now();
          if (Date.now() - this.lastAuthLog > 60_000) {
            this.lastAuthLog = Date.now();
            this.log("rejected an extension connection: secret mismatch");
          }
          socket.close(1008, "bad signature");
        }
        return;
      }
      if (!authed) {
        if (msg.type !== "hello") return;
        authed = true;
        clearTimeout(helloTimer);
        if (this.sock && this.sock !== socket) {
          try {
            this.sock.close(1000, "replaced");
          } catch {
            // already closed
          }
        }
        this.sock = socket;
        this.authError = null;
        this.lastFailAt = 0;
        this.extensionVersion = typeof msg.version === "string" ? msg.version : "unknown";
        socket.send(seal(this.cfg.secret, { type: "welcome", version: VERSION }));
        this.log(`extension connected (version ${this.extensionVersion})`);
        for (const wake of this.waiters.splice(0)) wake();
        return;
      }
      if (msg.type === "result" && typeof msg.reqId === "string") {
        const p = this.pending.get(msg.reqId);
        if (!p) {
          // A page that answered after its call gave up: close the tab it left open.
          const tabId = (msg.data as { tabId?: unknown } | undefined)?.tabId;
          if (msg.ok && typeof tabId === "number") this.request("close", { tabId }, 5000).catch(() => {});
          return;
        }
        this.pending.delete(msg.reqId);
        clearTimeout(p.timer);
        if (msg.ok) p.resolve(msg.data);
        else p.reject(new Error(typeof msg.error === "string" ? msg.error : "The browser reported an error."));
      }
    });

    socket.on("close", () => {
      clearTimeout(helloTimer);
      if (this.sock !== socket) return;
      this.sock = null;
      for (const p of this.pending.values()) {
        clearTimeout(p.timer);
        p.reject(new Error("The browser disconnected while the request was running."));
      }
      this.pending.clear();
      this.log("extension disconnected");
    });
    socket.on("error", () => {
      // close follows
    });
  }

  private waitForConnection(ms: number): Promise<boolean> {
    if (this.isConnected()) return Promise.resolve(true);
    return new Promise((resolve) => {
      let done = false;
      const wake = () => {
        if (done) return;
        done = true;
        clearTimeout(timer);
        resolve(true);
      };
      const timer = setTimeout(() => {
        if (done) return;
        done = true;
        this.waiters = this.waiters.filter((w) => w !== wake);
        resolve(this.isConnected());
      }, ms);
      this.waiters.push(wake);
    });
  }

  private recentAuthError(): string | null {
    return this.authError && Date.now() - this.authErrorAt < AUTH_ERROR_MEMORY_MS ? this.authError : null;
  }

  // Whether the configured profile is open in a running browser: true, false, or null if unknown.
  browserRunning(): boolean | null {
    const dir = profileDir(this.cfg.profile);
    return dir ? profileInUse(dir) : null;
  }

  private launchBrowser(): void {
    const now = Date.now();
    // The cooldown covers every copy of the server, since LM Studio may restart it between calls.
    if (now - Math.max(this.lastLaunch, lastLaunchAt(this.cfg.dataDir)) < RELAUNCH_COOLDOWN_MS) return;
    this.lastLaunch = now;
    if (!this.cfg.browserPath) {
      this.launchError =
        `Zen was not found, so it could not be started automatically. Open the browser profile yourself, or set ${this.cfg.host.browserSetting}.`;
      return;
    }
    try {
      const child = spawn(this.cfg.browserPath, ["-P", this.cfg.profile, "--no-remote"], {
        detached: true,
        stdio: "ignore",
        windowsHide: false,
      });
      child.on("error", (e) => {
        this.launchError = `Starting the browser failed: ${e.message}`;
      });
      child.unref();
      recordLaunch(this.cfg.dataDir, now);
      this.launchError = null;
      this.log(`started browser profile "${this.cfg.profile}"`);
    } catch (e) {
      this.launchError = `Starting the browser failed: ${(e as Error).message}`;
    }
  }

  async ensureConnected(waitMs = CONNECT_WAIT_MS): Promise<void> {
    if (this.isConnected()) return;
    if (!this.cfg.secret) throw new BrowserUnavailable(`The secret is not set. Copy it from the extension's options page into ${this.cfg.host.secretSetting}.`);
    if (!this.wss) throw new BrowserUnavailable(this.bindError ?? "The search server is not listening.");
    // A running browser reconnects within a couple of seconds; give it that chance first.
    if (await this.waitForConnection(2500)) return;
    const authProblem = this.recentAuthError();
    if (authProblem) throw new BrowserUnavailable(authProblem);
    if (Date.now() - this.lastFailAt < REPEAT_FAIL_MS) throw new BrowserUnavailable(`Still not connected. ${this.lastFail}`);
    // Starting the browser while it is already open only adds a window or tab, so wait instead.
    const running = this.browserRunning();
    if (running === true) this.log(`browser profile "${this.cfg.profile}" is already open; waiting for its extension`);
    else this.launchBrowser();
    if (await this.waitForConnection(waitMs)) return;
    const startedRecently = Date.now() - Math.max(this.lastLaunch, lastLaunchAt(this.cfg.dataDir)) < RELAUNCH_COOLDOWN_MS;
    const problem = this.recentAuthError() ?? this.launchError;
    if (!problem && startedRecently) {
      throw new BrowserUnavailable(`The "${this.cfg.profile}" browser profile was started and is still loading; it will not be started twice.`, true);
    }
    this.lastFailAt = Date.now();
    this.lastFail =
      this.recentAuthError() ??
      (this.launchError ? `${this.launchError} The extension did not connect within ${Math.round(waitMs / 1000)} seconds either.` : null) ??
      (running === true
        ? `The "${this.cfg.profile}" profile is open, but the simpleSearch Bridge extension in it did not connect within ${Math.round(waitMs / 1000)} seconds. Check that the extension is enabled in that profile.`
        : `The browser extension did not connect within ${Math.round(waitMs / 1000)} seconds. Check that the "${this.cfg.profile}" profile is open and the simpleSearch Bridge extension is installed and enabled in it.`);
    throw new BrowserUnavailable(this.lastFail);
  }

  request<T = unknown>(type: string, args: Record<string, unknown>, timeoutMs: number): Promise<T> {
    const socket = this.sock;
    if (!socket || socket.readyState !== WebSocket.OPEN) {
      return Promise.reject(new BrowserUnavailable("The browser is not connected."));
    }
    const id = randomUUID();
    return new Promise<T>((resolve, reject) => {
      const timer = setTimeout(() => {
        this.pending.delete(id);
        reject(new Error(`The browser did not answer within ${Math.round(timeoutMs / 1000)} seconds.`));
      }, timeoutMs);
      this.pending.set(id, { resolve: resolve as (v: unknown) => void, reject, timer });
      socket.send(seal(this.cfg.secret, { ...args, type, id }));
    });
  }
}

import { spawn } from "node:child_process";
import { randomUUID } from "node:crypto";
import type { IncomingMessage } from "node:http";
import { WebSocket, WebSocketServer } from "ws";
import { Config, VERSION } from "./config";
import { lastLaunchAt, profileDir, profileInUse, recordLaunch } from "./profile";
import { ReplayGuard, seal, unseal } from "./protocol";

export class BrowserUnavailable extends Error {
  // True when the browser was just started and is still loading: worth one more try.
  constructor(message: string, readonly starting = false) {
    super(message);
  }
}

interface Pending {
  resolve: (value: unknown) => void;
  reject: (err: Error) => void;
  timer: NodeJS.Timeout;
}

const AUTH_ERROR_MEMORY_MS = 30_000;
const RELAUNCH_COOLDOWN_MS = 90_000;
// Kept under LM Studio's limit for a single tool call, which counts the search after it too.
const CONNECT_WAIT_MS = 40_000;
const REPEAT_FAIL_MS = 60_000;

// Owns the local WebSocket that the browser extension connects to, and starts the
// browser profile when a request arrives and no extension is connected.
export class Bridge {
  private wss: WebSocketServer | null = null;
  private sock: WebSocket | null = null;
  private guard = new ReplayGuard();
  private pending = new Map<string, Pending>();
  private waiters: Array<() => void> = [];
  private lastLaunch = 0;
  private authErrorAt = 0;
  private lastAuthLog = 0;
  // The last time a call gave up waiting for the extension, and why. A call soon after that
  // fails within seconds instead of waiting the full time again for the same reason.
  private lastFailAt = 0;
  private lastFail = "";

  bindError: string | null = null;
  launchError: string | null = null;
  authError: string | null = null;
  extensionVersion: string | null = null;

  constructor(private cfg: Config, private log: (msg: string) => void) {}

  start(): Promise<void> {
    return new Promise((resolve) => {
      const wss = new WebSocketServer({ host: "127.0.0.1", port: this.cfg.port, maxPayload: 32 * 1024 * 1024 });
      wss.on("listening", () => {
        this.wss = wss;
        this.log(`listening on 127.0.0.1:${this.cfg.port}`);
        resolve();
      });
      wss.on("error", (err: NodeJS.ErrnoException) => {
        if (this.wss) {
          this.log(`server error: ${err.message}`);
          return;
        }
        this.bindError =
          err.code === "EADDRINUSE"
            ? `Port ${this.cfg.port} is already in use, ${this.cfg.host.portInUse}.`
            : `Could not open port ${this.cfg.port}: ${err.message}`;
        this.log(this.bindError);
        resolve();
      });
      wss.on("connection", (socket, req) => this.onConnection(socket, req));
    });
  }

  // Resolves once the port is free again, so a new bridge can open it straight away.
  close(): Promise<void> {
    for (const p of this.pending.values()) {
      clearTimeout(p.timer);
      p.reject(new Error("The server is shutting down."));
    }
    this.pending.clear();
    this.sock?.terminate();
    this.sock = null;
    const wss = this.wss;
    this.wss = null;
    if (!wss) return Promise.resolve();
    return new Promise((resolve) => {
      const done = setTimeout(resolve, 2000);
      wss.close(() => {
        clearTimeout(done);
        resolve();
      });
    });
  }

  isListening(): boolean {
    return this.wss !== null;
  }

  isConnected(): boolean {
    return this.sock !== null && this.sock.readyState === WebSocket.OPEN;
  }

  private onConnection(socket: WebSocket, req: IncomingMessage): void {
    // Web pages can open WebSockets to localhost too. Only the extension is welcome.
    const origin = req.headers.origin ?? "";
    if (/^https?:/i.test(origin)) {
      socket.close(1008, "forbidden");
      return;
    }
    let authed = false;
    const helloTimer = setTimeout(() => {
      if (!authed) socket.close(1008, "no hello");
    }, 5000);

    socket.on("message", (data) => {
      const msg = unseal(this.cfg.secret, data.toString(), this.guard);
      if (!msg) {
        if (!authed) {
          this.authError =
            `The browser extension connected, but its secret does not match ${this.cfg.host.secretSetting}. Copy the secret from the extension's options page into ${this.cfg.host.secretSetting}${this.cfg.host.afterChange}.`;
          this.authErrorAt = Date.now();
          if (Date.now() - this.lastAuthLog > 60_000) {
            this.lastAuthLog = Date.now();
            this.log("rejected an extension connection: secret mismatch");
          }
          socket.close(1008, "bad signature");
        }
        return;
      }
      if (!authed) {
        if (msg.type !== "hello") return;
        authed = true;
        clearTimeout(helloTimer);
        if (this.sock && this.sock !== socket) {
          try {
            this.sock.close(1000, "replaced");
          } catch {
            // already closed
          }
        }
        this.sock = socket;
        this.authError = null;
        this.lastFailAt = 0;
        this.extensionVersion = typeof msg.version === "string" ? msg.version : "unknown";
        socket.send(seal(this.cfg.secret, { type: "welcome", version: VERSION }));
        this.log(`extension connected (version ${this.extensionVersion})`);
        for (const wake of this.waiters.splice(0)) wake();
        return;
      }
      if (msg.type === "result" && typeof msg.reqId === "string") {
        const p = this.pending.get(msg.reqId);
        if (!p) {
          // A page that answered after its call gave up: close the tab it left open.
          const tabId = (msg.data as { tabId?: unknown } | undefined)?.tabId;
          if (msg.ok && typeof tabId === "number") this.request("close", { tabId }, 5000).catch(() => {});
          return;
        }
        this.pending.delete(msg.reqId);
        clearTimeout(p.timer);
        if (msg.ok) p.resolve(msg.data);
        else p.reject(new Error(typeof msg.error === "string" ? msg.error : "The browser reported an error."));
      }
    });

    socket.on("close", () => {
      clearTimeout(helloTimer);
      if (this.sock !== socket) return;
      this.sock = null;
      for (const p of this.pending.values()) {
        clearTimeout(p.timer);
        p.reject(new Error("The browser disconnected while the request was running."));
      }
      this.pending.clear();
      this.log("extension disconnected");
    });
    socket.on("error", () => {
      // close follows
    });
  }

  private waitForConnection(ms: number): Promise<boolean> {
    if (this.isConnected()) return Promise.resolve(true);
    return new Promise((resolve) => {
      let done = false;
      const wake = () => {
        if (done) return;
        done = true;
        clearTimeout(timer);
        resolve(true);
      };
      const timer = setTimeout(() => {
        if (done) return;
        done = true;
        this.waiters = this.waiters.filter((w) => w !== wake);
        resolve(this.isConnected());
      }, ms);
      this.waiters.push(wake);
    });
  }

  private recentAuthError(): string | null {
    return this.authError && Date.now() - this.authErrorAt < AUTH_ERROR_MEMORY_MS ? this.authError : null;
  }

  // Whether the configured profile is open in a running browser: true, false, or null if unknown.
  browserRunning(): boolean | null {
    const dir = profileDir(this.cfg.profile);
    return dir ? profileInUse(dir) : null;
  }

  private launchBrowser(): void {
    const now = Date.now();
    // The cooldown covers every copy of the server, since LM Studio may restart it between calls.
    if (now - Math.max(this.lastLaunch, lastLaunchAt(this.cfg.dataDir)) < RELAUNCH_COOLDOWN_MS) return;
    this.lastLaunch = now;
    if (!this.cfg.browserPath) {
      this.launchError =
        `Zen was not found, so it could not be started automatically. Open the browser profile yourself, or set ${this.cfg.host.browserSetting}.`;
      return;
    }
    try {
      const child = spawn(this.cfg.browserPath, ["-P", this.cfg.profile, "--no-remote"], {
        detached: true,
        stdio: "ignore",
        windowsHide: false,
      });
      child.on("error", (e) => {
        this.launchError = `Starting the browser failed: ${e.message}`;
      });
      child.unref();
      recordLaunch(this.cfg.dataDir, now);
      this.launchError = null;
      this.log(`started browser profile "${this.cfg.profile}"`);
    } catch (e) {
      this.launchError = `Starting the browser failed: ${(e as Error).message}`;
    }
  }

  async ensureConnected(waitMs = CONNECT_WAIT_MS): Promise<void> {
    if (this.isConnected()) return;
    if (!this.cfg.secret) throw new BrowserUnavailable(`The secret is not set. Copy it from the extension's options page into ${this.cfg.host.secretSetting}.`);
    if (!this.wss) throw new BrowserUnavailable(this.bindError ?? "The search server is not listening.");
    // A running browser reconnects within a couple of seconds; give it that chance first.
    if (await this.waitForConnection(2500)) return;
    const authProblem = this.recentAuthError();
    if (authProblem) throw new BrowserUnavailable(authProblem);
    if (Date.now() - this.lastFailAt < REPEAT_FAIL_MS) throw new BrowserUnavailable(`Still not connected. ${this.lastFail}`);
    // Starting the browser while it is already open only adds a window or tab, so wait instead.
    const running = this.browserRunning();
    if (running === true) this.log(`browser profile "${this.cfg.profile}" is already open; waiting for its extension`);
    else this.launchBrowser();
    if (await this.waitForConnection(waitMs)) return;
    const startedRecently = Date.now() - Math.max(this.lastLaunch, lastLaunchAt(this.cfg.dataDir)) < RELAUNCH_COOLDOWN_MS;
    const problem = this.recentAuthError() ?? this.launchError;
    if (!problem && startedRecently) {
      throw new BrowserUnavailable(`The "${this.cfg.profile}" browser profile was started and is still loading; it will not be started twice.`, true);
    }
    this.lastFailAt = Date.now();
    this.lastFail =
      this.recentAuthError() ??
      (this.launchError ? `${this.launchError} The extension did not connect within ${Math.round(waitMs / 1000)} seconds either.` : null) ??
      (running === true
        ? `The "${this.cfg.profile}" profile is open, but the simpleSearch Bridge extension in it did not connect within ${Math.round(waitMs / 1000)} seconds. Check that the extension is enabled in that profile.`
        : `The browser extension did not connect within ${Math.round(waitMs / 1000)} seconds. Check that the "${this.cfg.profile}" profile is open and the simpleSearch Bridge extension is installed and enabled in it.`);
    throw new BrowserUnavailable(this.lastFail);
  }

  request<T = unknown>(type: string, args: Record<string, unknown>, timeoutMs: number): Promise<T> {
    const socket = this.sock;
    if (!socket || socket.readyState !== WebSocket.OPEN) {
      return Promise.reject(new BrowserUnavailable("The browser is not connected."));
    }
    const id = randomUUID();
    return new Promise<T>((resolve, reject) => {
      const timer = setTimeout(() => {
        this.pending.delete(id);
        reject(new Error(`The browser did not answer within ${Math.round(timeoutMs / 1000)} seconds.`));
      }, timeoutMs);
      this.pending.set(id, { resolve: resolve as (v: unknown) => void, reject, timer });
      socket.send(seal(this.cfg.secret, { ...args, type, id }));
    });
  }
}