// A pool of the workers of worker.ts. It runs the eigenvalue jobs of
// modes.ts through a JobRunner, and splits the modes of a model between the
// workers for the synthesis: the main thread steps the contact (which
// couples all modes and is cheap), each worker forms the parts of its
// modes, and the main thread sums them.

import { type LinearSystem } from "../model/assemble.ts";
import { type Loss, type Plate, type SimParams } from "../params.ts";
import { type Job, type JobResult, type JobRunner, type ModeSet, computeModes } from "./modes.ts";
import {
  type ContactResult,
  type ModalModel,
  type Synthesis,
  attackPressure,
  listenerWeights,
  modalRates,
  simulateContact,
  sliceContact,
  sliceModel,
  strikeGains,
  tailStart,
} from "./synth.ts";
import { type WorkerRequest, type WorkerResponse } from "./worker.ts";

interface Pending {
  resolve: (r: WorkerResponse) => void;
  reject: (e: Error) => void;
}

class PoolWorker {
  worker: Worker;
  busy = false;
  pending = new Map<number, Pending>();
  ready: Pending | null = null;
  constructor() {
    this.worker = new Worker(new URL("./worker.ts", import.meta.url), { type: "module" });
    this.worker.onmessage = (ev: MessageEvent<WorkerResponse>) => {
      const r = ev.data;
      if (r.type === "ready") { this.ready?.resolve(r); this.ready = null; return; }
      const p = this.pending.get(r.id);
      if (!p) return;
      this.pending.delete(r.id);
      if (r.type === "error") p.reject(new Error(r.message)); else p.resolve(r);
    };
    this.worker.onerror = (e) => {
      for (const p of this.pending.values()) p.reject(new Error(e.message));
      this.pending.clear();
    };
  }
  send(msg: WorkerRequest, transfer: ArrayBuffer[] = []): Promise<WorkerResponse> {
    return new Promise((resolve, reject) => {
      if (msg.type === "system" || msg.type === "model") this.ready = { resolve, reject };
      else this.pending.set(msg.id, { resolve, reject });
      this.worker.postMessage(msg, transfer);
    });
  }
}

export class WorkerPool {
  private workers: PoolWorker[] = [];
  private queue: (() => void)[] = [];
  private modelSplit: { from: number; to: number }[] = [];
  private nextId = 1;

  constructor(public readonly size: number) {
    for (let i = 0; i < size; i++) this.workers.push(new PoolWorker());
  }

  static defaultSize(): number {
    const hw = typeof navigator !== "undefined" ? navigator.hardwareConcurrency || 4 : 4;
    return Math.max(2, Math.min(12, hw));
  }

  /** Waits for an idle worker. */
  private acquire(): Promise<PoolWorker> {
    return new Promise((resolve) => {
      const tryTake = (): boolean => {
        const w = this.workers.find((w) => !w.busy);
        if (!w) return false;
        w.busy = true;
        resolve(w);
        return true;
      };
      if (!tryTake()) this.queue.push(() => tryTake());
    });
  }
  private release(w: PoolWorker): void {
    w.busy = false;
    const next = this.queue.shift();
    if (next) next();
  }

  /** Sends the linear system to every worker. */
  async setSystem(sys: LinearSystem): Promise<void> {
    await Promise.all(this.workers.map((w) => w.send({ type: "system", n: sys.n, bw: sys.bw, band: sys.band, mass: sys.mass })));
  }

  runner(): JobRunner {
    return {
      run: async (job: Job): Promise<JobResult> => {
        const w = await this.acquire();
        try {
          const r = await w.send({ type: "job", id: this.nextId++, job });
          return (r as { result: JobResult }).result;
        } finally {
          this.release(w);
        }
      },
      close: () => {},
    };
  }

  /** The modes of the system, computed by the workers. */
  async modes(sys: LinearSystem, fMax: number, onProgress?: (done: number, total: number) => void): Promise<ModeSet> {
    await this.setSystem(sys);
    return computeModes(sys, fMax, this.runner(), { sliceSize: sys.fields === 1 ? 30 : 100, onProgress });
  }

  /** Splits the modes of the model between the workers. */
  async setModel(model: ModalModel): Promise<void> {
    const n = this.workers.length;
    this.modelSplit = [];
    const per = Math.ceil(model.count / n);
    await Promise.all(
      this.workers.map((w, i) => {
        const from = Math.min(model.count, i * per), to = Math.min(model.count, (i + 1) * per);
        this.modelSplit.push({ from, to });
        const slice = sliceModel(model, from, to);
        return w.send({ type: "model", model: slice });
      }),
    );
  }

  /**
   * Synthesizes a strike with the model last set. The audio comes first
   * (through `onAudio`), then the movie frames; the returned promise
   * resolves when both are done.
   */
  async synthesize(model: ModalModel, pl: Plate, loss: Loss, p: SimParams, onAudio?: (audio: Float32Array, contact: ContactResult) => void, withMovie = true): Promise<Synthesis> {
    const sigma = modalRates(model, loss);
    const g = strikeGains(model, pl, p);
    const contact = simulateContact(model, sigma, g, p, p.fps);
    const { weight, delay, maxDelay } = listenerWeights(pl, p);
    const tStart = tailStart(contact, maxDelay);
    const len = Math.min(contact.tableLen, Math.ceil(tStart * contact.tableRate) + 2);
    const nSamples = Math.round(p.duration * p.fs);
    const nFrames = withMovie ? Math.round(p.duration * p.fps) : 0;
    const parts = await Promise.all(
      this.workers.map((w, i) => {
        const { from, to } = this.modelSplit[i];
        if (to <= from) return Promise.resolve(null);
        return w.send({
          type: "synth", id: this.nextId++, sigma: sigma.slice(from, to), contact: sliceContact(contact, from, to),
          weight, delay, tStart, len, fs: p.fs, nSamples, audio: true, movie: null,
        });
      }),
    );
    const audio = new Float32Array(nSamples);
    const table = new Float32Array(model.nRows * len);
    for (const r of parts) {
      if (!r || r.type !== "synth") continue;
      if (r.audio) for (let i = 0; i < nSamples; i++) audio[i] += r.audio[i];
      if (r.table) for (let i = 0; i < table.length; i++) table[i] += r.table[i];
    }
    attackPressure(model, table, len, contact.tableRate, weight, delay, p.fs, tStart, audio);
    onAudio?.(audio, contact);
    const frames = new Float32Array(nFrames * model.nGrid);
    if (withMovie) {
      const mparts = await Promise.all(
        this.workers.map((w, i) => {
          const { from, to } = this.modelSplit[i];
          if (to <= from) return Promise.resolve(null);
          return w.send({
            type: "synth", id: this.nextId++, sigma: sigma.slice(from, to), contact: sliceContact(contact, from, to),
            weight, delay, tStart, len, fs: p.fs, nSamples, audio: false, movie: { fps: p.fps, nFrames },
          });
        }),
      );
      for (const r of mparts) if (r && r.type === "synth" && r.frames) for (let i = 0; i < frames.length; i++) frames[i] += r.frames[i];
    }
    return { audio, frames, nFrames, contact };
  }

  terminate(): void {
    for (const w of this.workers) w.worker.terminate();
    this.workers = [];
  }
}