import { afterEach, describe, expect, it, vi } from "vitest";
import { WorkerSupervisor, type WorkerLike } from "./WorkerSupervisor";
import { WORKER_PROTOCOL_VERSION, type WorkerRequest } from "./worker-protocol";

class FakeWorker implements WorkerLike {
  onmessage: ((event: MessageEvent<unknown>) => void) | null = null;
  onerror: ((event: ErrorEvent) => void) | null = null;
  onmessageerror: ((event: MessageEvent<unknown>) => void) | null = null;
  readonly messages: unknown[] = [];
  terminateCalls = 0;

  postMessage(message: unknown): void {
    this.messages.push(message);
  }

  terminate(): void {
    this.terminateCalls += 1;
  }

  succeed(payload: string): void {
    const request = this.messages.at(-1) as WorkerRequest<string>;
    this.onmessage?.(
      new MessageEvent("message", {
        data: {
          protocolVersion: WORKER_PROTOCOL_VERSION,
          requestId: request.requestId,
          generation: request.generation,
          ok: true,
          payload,
        },
      }),
    );
  }
}

afterEach(() => {
  vi.useRealTimers();
});

describe("WorkerSupervisor", () => {
  it("terminates a timed-out worker and succeeds with a fresh generation", async () => {
    vi.useFakeTimers();
    const workers: FakeWorker[] = [];
    const supervisor = new WorkerSupervisor<string, string>(
      "test worker",
      () => {
        const worker = new FakeWorker();
        workers.push(worker);
        return worker;
      },
    );

    const first = supervisor.run("hang", 25);
    const firstExpectation = expect(first).rejects.toMatchObject({
      kind: "timeout",
    });
    await vi.advanceTimersByTimeAsync(25);
    await firstExpectation;
    expect(workers[0]?.terminateCalls).toBe(1);

    const second = supervisor.run("ok", 25);
    expect(workers).toHaveLength(2);
    workers[1]?.succeed("recovered");
    await expect(second).resolves.toBe("recovered");
    expect(supervisor.currentGeneration).toBe(2);
    supervisor.dispose();
  });

  it("supersedes an active request and ignores the stale worker", async () => {
    const workers: FakeWorker[] = [];
    const supervisor = new WorkerSupervisor<string, string>(
      "test worker",
      () => {
        const worker = new FakeWorker();
        workers.push(worker);
        return worker;
      },
    );
    const firstWorkerMessage: {
      current?: ((event: MessageEvent<unknown>) => void) | null;
    } = {};
    const first = supervisor.run("first", 1_000);
    firstWorkerMessage.current = workers[0]?.onmessage;
    const second = supervisor.run("second", 1_000, { supersede: true });
    await expect(first).rejects.toMatchObject({ kind: "cancelled" });
    expect(workers[0]?.terminateCalls).toBe(1);

    const oldRequest = workers[0]?.messages[0] as WorkerRequest<string>;
    firstWorkerMessage.current?.(
      new MessageEvent("message", {
        data: {
          protocolVersion: WORKER_PROTOCOL_VERSION,
          requestId: oldRequest.requestId,
          generation: oldRequest.generation,
          ok: true,
          payload: "stale",
        },
      }),
    );
    workers[1]?.succeed("fresh");
    await expect(second).resolves.toBe("fresh");
    supervisor.dispose();
  });

  it("restarts after a crash", async () => {
    const workers: FakeWorker[] = [];
    const supervisor = new WorkerSupervisor<string, string>(
      "test worker",
      () => {
        const worker = new FakeWorker();
        workers.push(worker);
        return worker;
      },
    );
    const first = supervisor.run("crash", 1_000);
    workers[0]?.onerror?.(
      new ErrorEvent("error", { message: "simulated crash" }),
    );
    await expect(first).rejects.toMatchObject({ kind: "crash" });
    const second = supervisor.run("next", 1_000);
    workers[1]?.succeed("ok");
    await expect(second).resolves.toBe("ok");
    supervisor.dispose();
  });
});
