import { describe, expect, it, vi } from "vitest";

import { TOOLCRAFT_DEFAULT_MODEL_IMPORT_LIMITS } from "../model-import-limits";
import { createToolcraftModelWorkerClient } from "./model-import-worker-client";
import {
  completeDecode,
  deferred,
  FakeWorker,
  forgedDigest,
  importRequest,
  planDigest,
} from "./model-import-worker-client-test-support";
import { digestToolcraftModelWorkerBytes } from "./model-import-worker-receipt";

const LARGE_COPY_BYTE_LENGTH = 1024 * 1024 + 17;
const envelopeDigest = `sha256:${"3".repeat(64)}`;

function patternedBuffer(byteLength: number): ArrayBuffer {
  const bytes = new Uint8Array(byteLength);
  for (let index = 0; index < bytes.length; index += 1) {
    bytes[index] = (index * 31 + 7) & 0xff;
  }
  return bytes.buffer;
}

function detach(buffer: ArrayBuffer): void {
  structuredClone(buffer, { transfer: [buffer] });
}

function repairRequest(
  jobId: string,
  canonicalDocument: ArrayBuffer,
  envelopeBytes: ArrayBuffer,
) {
  return {
    canonicalDocument,
    canonicalDocumentDigest: forgedDigest,
    jobId,
    limits: TOOLCRAFT_DEFAULT_MODEL_IMPORT_LIMITS,
    repairPlanEnvelope: {
      bytes: envelopeBytes,
      envelopeDigest,
      planDigest,
      version: 1 as const,
    },
  };
}

describe("model worker client cooperative snapshots", () => {
  it("yields before posting a large source snapshot and transfers exact owned bytes", async () => {
    const worker = new FakeWorker();
    const firstYield = deferred<void>();
    const yieldToHost = vi.fn()
      .mockImplementationOnce(() => firstYield.promise)
      .mockResolvedValue(undefined);
    const source = patternedBuffer(LARGE_COPY_BYTE_LENGTH);
    const expectedDigest = digestToolcraftModelWorkerBytes(source);
    const client = createToolcraftModelWorkerClient({
      workerFactory: () => worker,
      yieldToHost,
    });

    const pending = client.importModel(importRequest("large-source", source));

    expect(pending).toBeInstanceOf(Promise);
    await vi.waitFor(() => expect(yieldToHost).toHaveBeenCalled());
    expect(worker.posts).toHaveLength(0);
    expect(digestToolcraftModelWorkerBytes(source)).toBe(expectedDigest);

    firstYield.resolve();
    await vi.waitFor(() => expect(worker.posts).toHaveLength(1));
    const posted = worker.posts[0]!;
    if (posted.message.kind !== "decode-and-analyze") {
      throw new Error("Expected decode request.");
    }
    const transferred = posted.message.payload.bundle.sourceFiles[0]!.bytes;
    expect(transferred).not.toBe(source);
    expect(digestToolcraftModelWorkerBytes(transferred)).toBe(expectedDigest);
    expect(posted.transfer).toEqual([transferred]);

    client.cancel("large-source");
    await expect(pending).resolves.toMatchObject({ kind: "cancelled" });
  });

  it("cannot late-post or publish when superseded during a source snapshot", async () => {
    const worker = new FakeWorker();
    const firstYield = deferred<void>();
    const yieldToHost = vi.fn()
      .mockImplementationOnce(() => firstYield.promise)
      .mockResolvedValue(undefined);
    const firstOnResult = vi.fn();
    const client = createToolcraftModelWorkerClient({
      workerFactory: () => worker,
      yieldToHost,
    });
    const first = client.importModel(
      importRequest("copy-first", patternedBuffer(LARGE_COPY_BYTE_LENGTH)),
      { onResult: firstOnResult },
    );
    await vi.waitFor(() => expect(yieldToHost).toHaveBeenCalled());

    const second = client.importModel(importRequest("copy-second"));

    await expect(first).resolves.toMatchObject({
      jobId: "copy-first",
      kind: "cancelled",
    });
    await vi.waitFor(() => expect(worker.posts).toHaveLength(1));
    expect(worker.posts[0]!.message.jobId).toBe("copy-second");

    firstYield.resolve();
    await vi.waitFor(() => expect(worker.posts).toHaveLength(1));
    expect(firstOnResult).not.toHaveBeenCalled();

    completeDecode(worker, worker.posts[0]!.message);
    await expect(second).resolves.toMatchObject({ kind: "result" });
    expect(firstOnResult).not.toHaveBeenCalled();
  });

  it("resolves a typed error when source bytes detach after activation", async () => {
    const worker = new FakeWorker();
    const firstYield = deferred<void>();
    const yieldToHost = vi.fn()
      .mockImplementationOnce(() => firstYield.promise)
      .mockResolvedValue(undefined);
    const source = patternedBuffer(LARGE_COPY_BYTE_LENGTH);
    const client = createToolcraftModelWorkerClient({
      workerFactory: () => worker,
      yieldToHost,
    });

    const pending = client.importModel(importRequest("detach-during-copy", source));
    await vi.waitFor(() => expect(yieldToHost).toHaveBeenCalled());

    detach(source);
    firstYield.resolve();

    await expect(pending).resolves.toMatchObject({
      feedback: { code: "model-worker-source-snapshot-failed" },
      jobId: "detach-during-copy",
      kind: "error",
    });
    expect(worker.posts).toHaveLength(0);
  });

  it("resolves a typed error when repair bytes detach after activation", async () => {
    const worker = new FakeWorker();
    const firstYield = deferred<void>();
    const yieldToHost = vi.fn()
      .mockImplementationOnce(() => firstYield.promise)
      .mockResolvedValue(undefined);
    const sha256 = vi.fn();
    const canonicalDocument = patternedBuffer(LARGE_COPY_BYTE_LENGTH);
    const client = createToolcraftModelWorkerClient({
      sha256,
      workerFactory: () => worker,
      yieldToHost,
    });

    const pending = client.repair(repairRequest(
      "detach-repair-during-copy",
      canonicalDocument,
      patternedBuffer(257),
    ));
    await vi.waitFor(() => expect(yieldToHost).toHaveBeenCalled());

    detach(canonicalDocument);
    firstYield.resolve();

    await expect(pending).resolves.toMatchObject({
      feedback: { code: "model-worker-repair-snapshot-failed" },
      jobId: "detach-repair-during-copy",
      kind: "error",
    });
    expect(sha256).not.toHaveBeenCalled();
    expect(worker.posts).toHaveLength(0);
  });

  it.each(["canonical-document", "plan-envelope"] as const)(
    "copies a large %s cooperatively and hashes both snapshots before posting",
    async (largeInput) => {
      const worker = new FakeWorker();
      const firstYield = deferred<void>();
      const yieldToHost = vi.fn()
        .mockImplementationOnce(() => firstYield.promise)
        .mockResolvedValue(undefined);
      const canonicalHash = deferred<string>();
      const envelopeHash = deferred<string>();
      const sha256 = vi.fn()
        .mockImplementationOnce(() => canonicalHash.promise)
        .mockImplementationOnce(() => envelopeHash.promise);
      const canonicalDocument = patternedBuffer(
        largeInput === "canonical-document" ? LARGE_COPY_BYTE_LENGTH : 257,
      );
      const envelopeBytes = patternedBuffer(
        largeInput === "plan-envelope" ? LARGE_COPY_BYTE_LENGTH : 257,
      );
      const expectedDocumentDigest = digestToolcraftModelWorkerBytes(
        canonicalDocument,
      );
      const expectedEnvelopeDigest = digestToolcraftModelWorkerBytes(envelopeBytes);
      const client = createToolcraftModelWorkerClient({
        sha256,
        workerFactory: () => worker,
        yieldToHost,
      });

      const pending = client.repair(
        repairRequest(`large-${largeInput}`, canonicalDocument, envelopeBytes),
      );

      expect(pending).toBeInstanceOf(Promise);
      await vi.waitFor(() => expect(yieldToHost).toHaveBeenCalled());
      expect(sha256).not.toHaveBeenCalled();
      expect(worker.posts).toHaveLength(0);

      firstYield.resolve();
      await vi.waitFor(() => expect(sha256).toHaveBeenCalledTimes(1));
      expect(worker.posts).toHaveLength(0);
      canonicalHash.resolve(forgedDigest);
      await vi.waitFor(() => expect(sha256).toHaveBeenCalledTimes(2));
      expect(worker.posts).toHaveLength(0);
      envelopeHash.resolve(envelopeDigest);

      await vi.waitFor(() => expect(worker.posts).toHaveLength(1));
      const posted = worker.posts[0]!;
      if (posted.message.kind !== "repair") {
        throw new Error("Expected repair request.");
      }
      expect(digestToolcraftModelWorkerBytes(
        posted.message.payload.canonicalDocument,
      )).toBe(expectedDocumentDigest);
      expect(digestToolcraftModelWorkerBytes(
        posted.message.payload.repairPlanEnvelope.bytes,
      )).toBe(expectedEnvelopeDigest);
      expect(posted.transfer).toEqual([
        posted.message.payload.canonicalDocument,
        posted.message.payload.repairPlanEnvelope.bytes,
      ]);

      client.cancel(`large-${largeInput}`);
      await expect(pending).resolves.toMatchObject({ kind: "cancelled" });
    },
  );
});

describe("model worker client input preflight", () => {
  it("returns a typed detached-source error without disturbing an active job", async () => {
    const worker = new FakeWorker();
    const client = createToolcraftModelWorkerClient({ workerFactory: () => worker });
    const active = client.importModel(importRequest("active-import"));
    const detachedSource = new Uint8Array([9, 8, 7]).buffer;
    const attemptedRequest = importRequest("detached-import", detachedSource);
    detach(detachedSource);
    let attempted!: ReturnType<typeof client.importModel>;

    expect(() => {
      attempted = client.importModel(attemptedRequest);
    }).not.toThrow();

    await expect(attempted).resolves.toMatchObject({
      feedback: { code: "source-file-buffer-invalid" },
      jobId: "detached-import",
      kind: "error",
    });
    expect(worker.terminate).not.toHaveBeenCalled();
    expect(worker.posts).toHaveLength(1);
    expect(worker.posts[0]!.message.jobId).toBe("active-import");

    completeDecode(worker, worker.posts[0]!.message);
    await expect(active).resolves.toMatchObject({ kind: "result" });
  });

  it("rejects aggregate source limits before superseding an active job", async () => {
    const worker = new FakeWorker();
    const client = createToolcraftModelWorkerClient({ workerFactory: () => worker });
    const active = client.importModel(importRequest("active-limit"));
    const attemptedRequest = importRequest(
      "oversized-import",
      new Uint8Array([1, 2, 3, 4]).buffer,
    );
    attemptedRequest.limits = {
      ...TOOLCRAFT_DEFAULT_MODEL_IMPORT_LIMITS,
      maxSourceBytes: 3,
    };

    const attempted = client.importModel(attemptedRequest);

    await expect(attempted).resolves.toMatchObject({
      feedback: { code: "source-byte-limit-exceeded" },
      jobId: "oversized-import",
      kind: "error",
    });
    expect(worker.terminate).not.toHaveBeenCalled();
    expect(worker.posts).toHaveLength(1);

    completeDecode(worker, worker.posts[0]!.message);
    await expect(active).resolves.toMatchObject({ kind: "result" });
  });

  it.each(["canonical-document", "plan-envelope"] as const)(
    "returns a typed detached %s error without disturbing an active job",
    async (detachedInput) => {
      const worker = new FakeWorker();
      const client = createToolcraftModelWorkerClient({ workerFactory: () => worker });
      const active = client.importModel(importRequest(`active-${detachedInput}`));
      const canonicalDocument = new Uint8Array([1, 2, 3]).buffer;
      const envelopeBytes = new Uint8Array([4, 5, 6]).buffer;
      const attemptedRequest = repairRequest(
        `detached-${detachedInput}`,
        canonicalDocument,
        envelopeBytes,
      );
      detach(detachedInput === "canonical-document" ? canonicalDocument : envelopeBytes);
      let attempted!: ReturnType<typeof client.repair>;

      expect(() => {
        attempted = client.repair(attemptedRequest);
      }).not.toThrow();

      await expect(attempted).resolves.toMatchObject({
        feedback: {
          code: detachedInput === "canonical-document"
            ? "repair-document-buffer-invalid"
            : "repair-plan-envelope-invalid",
        },
        jobId: `detached-${detachedInput}`,
        kind: "error",
      });
      expect(worker.terminate).not.toHaveBeenCalled();
      expect(worker.posts).toHaveLength(1);

      completeDecode(worker, worker.posts[0]!.message);
      await expect(active).resolves.toMatchObject({ kind: "result" });
    },
  );

  it("rejects malformed repair-envelope metadata before supersession", async () => {
    const worker = new FakeWorker();
    const client = createToolcraftModelWorkerClient({ workerFactory: () => worker });
    const active = client.importModel(importRequest("active-envelope"));
    const attemptedRequest = repairRequest(
      "malformed-envelope",
      new Uint8Array([1]).buffer,
      new Uint8Array([2]).buffer,
    );
    const malformedRequest = {
      ...attemptedRequest,
      repairPlanEnvelope: {
        ...attemptedRequest.repairPlanEnvelope,
        envelopeDigest: "not-a-digest",
      },
    };

    const attempted = client.repair(malformedRequest);

    await expect(attempted).resolves.toMatchObject({
      feedback: { code: "repair-plan-envelope-invalid" },
      jobId: "malformed-envelope",
      kind: "error",
    });
    expect(worker.terminate).not.toHaveBeenCalled();
    expect(worker.posts).toHaveLength(1);

    completeDecode(worker, worker.posts[0]!.message);
    await expect(active).resolves.toMatchObject({ kind: "result" });
  });
});