import type { BackendInfo, RunRequest, RunResult, Runner } from "./types";
import type { FromWorker, ToWorker } from "./protocol";

interface Pending {
  resolve: (r: RunResult) => void;
  onOutput?: (text: string) => void;
  output: string[];
  t0: number;
}

/**
 * Runs numbl in a Web Worker the runner owns. Cancelling terminates the worker
 * and a fresh one is created for the next run, so a runaway eigensolve can
 * always be stopped, with or without cross-origin isolation.
 */
export class WorkerRunner implements Runner {
  private worker: Worker | null = null;
  private nextId = 1;
  private pending = new Map<number, Pending>();
  private infoResolvers = new Map<number, (b: BackendInfo) => void>();
  private cachedBackend: BackendInfo | null = null;
  private readonly createWorker: () => Worker;

  constructor(createWorker?: () => Worker) {
    this.createWorker =
      createWorker ??
      (() => new Worker(new URL("./worker.ts", import.meta.url), { type: "module" }));
  }

  get backend(): Promise<BackendInfo> {
    if (this.cachedBackend) return Promise.resolve(this.cachedBackend);
    return new Promise((resolve) => {
      const id = this.nextId++;
      this.infoResolvers.set(id, (b) => {
        this.cachedBackend = b;
        resolve(b);
      });
      this.ensureWorker().postMessage({ type: "info", id } satisfies ToWorker);
    });
  }

  private ensureWorker(): Worker {
    if (this.worker) return this.worker;
    const w = this.createWorker();
    w.onmessage = (e: MessageEvent<FromWorker>) => this.onMessage(e.data);
    w.onerror = (e) => this.failAll(`numbl worker crashed: ${e.message || "unknown error"}`);
    this.worker = w;
    return w;
  }

  private onMessage(msg: FromWorker) {
    if (msg.type === "info") {
      this.infoResolvers.get(msg.id)?.(msg.backend);
      this.infoResolvers.delete(msg.id);
      return;
    }
    const p = this.pending.get(msg.id);
    if (!p) return;
    if (msg.type === "output") {
      p.output.push(msg.text);
      p.onOutput?.(msg.text);
    } else if (msg.type === "done") {
      this.pending.delete(msg.id);
      p.resolve(msg.result);
    }
  }

  private failAll(error: string, cancelled = false) {
    const backend = this.cachedBackend ?? {
      engine: "numbl",
      version: "unknown",
      optimization: "1",
      host: "web-worker",
    };
    for (const p of this.pending.values()) {
      p.resolve({
        ok: false,
        error,
        output: p.output,
        wallMs: performance.now() - p.t0,
        backend,
        cancelled,
      });
    }
    this.pending.clear();
  }

  run(req: RunRequest, onOutput?: (text: string) => void): Promise<RunResult> {
    return new Promise((resolve) => {
      const id = this.nextId++;
      this.pending.set(id, { resolve, onOutput, output: [], t0: performance.now() });
      this.ensureWorker().postMessage({
        type: "run",
        id,
        bundle: req.bundle,
        inputs: req.inputs,
      } satisfies ToWorker);
    });
  }

  cancel(): void {
    if (!this.worker) return;
    this.worker.terminate();
    this.worker = null;
    this.failAll("run cancelled", true);
  }

  dispose(): void {
    this.cancel();
  }
}