import type { InferenceRequest, InferenceResult } from "../shared/types";

export interface InferenceWorker {
  addEventListener(type: "message" | "error", listener: EventListener): void;
  removeEventListener(type: "message" | "error", listener: EventListener): void;
  postMessage(message: { requests: InferenceRequest[]; modelBaseUrl: string }): void;
  terminate(): void;
}

type WorkerFactory = () => InferenceWorker;
type WorkerReply = { results?: InferenceResult[]; error?: string };

/** Keeps ONNX Runtime configuration immutable by replacing the worker on change. */
export class InferenceWorkerHost {
  private worker: InferenceWorker | undefined;
  constructor(private readonly createWorker: WorkerFactory, private readonly modelBaseUrl: string) {}

  run(requests: InferenceRequest[]): Promise<InferenceResult[]> {
    if (!this.worker) this.worker = this.createWorker();
    const worker = this.worker;
    return new Promise((resolve, reject) => {
      const complete = (): void => {
        worker.removeEventListener("message", onMessage);
        worker.removeEventListener("error", onError);
      };
      const onMessage: EventListener = (event) => {
        const reply = (event as MessageEvent<WorkerReply>).data;
        complete();
        if (reply.error) reject(new Error(reply.error));
        else resolve(reply.results ?? []);
      };
      const onError: EventListener = (event) => {
        complete();
        const error = event as ErrorEvent;
        reject(new Error(error.message || "Inference worker failed"));
      };
      worker.addEventListener("message", onMessage);
      worker.addEventListener("error", onError);
      worker.postMessage({ requests, modelBaseUrl: this.modelBaseUrl });
    });
  }

  dispose(): void {
    this.worker?.terminate();
    this.worker = undefined;
  }
}
