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

import type { ToolcraftModelImportLimits } from "../../schema/types";
import { encodeToolcraftModelDocument } from "../canonical/model-document-codec";
import { digestToolcraftModelDocument } from "../canonical/model-document-digest";
import { createValidModelDocument } from "../canonical/model-document-test-support";
import { TOOLCRAFT_DEFAULT_MODEL_IMPORT_LIMITS } from "../model-import-limits";
import { encodeToolcraftModelRepairPlanEnvelope } from "../topology/model-repair-plan-codec";
import { analyzeToolcraftModel } from "../topology/model-topology-analysis";
import { createToolcraftModelWorkerClient } from "./model-import-worker-client";
import {
  FakeWorker,
  forgedDigest,
} from "./model-import-worker-client-test-support";
import { snapshotToolcraftModelWorkerAnalysisSummary } from "./model-import-worker-diagnostics";
import type { ToolcraftModelWorkerRequest } from "./model-import-worker-protocol";
import {
  createToolcraftModelWorkerReceiptDigest,
  digestToolcraftModelWorkerBytes,
} from "./model-import-worker-receipt";
import {
  fakeRepairPlan,
  repairRequest,
} from "./model-import-worker-runtime-test-support";

type RepairRequest = Extract<ToolcraftModelWorkerRequest, { kind: "repair" }>;

const modelImportLimitNames = Object.keys(
  TOOLCRAFT_DEFAULT_MODEL_IMPORT_LIMITS,
) as (keyof ToolcraftModelImportLimits)[];

function repairResultFor(request: RepairRequest) {
  const document = createValidModelDocument();
  const canonicalDocument = new Uint8Array(
    encodeToolcraftModelDocument(document),
  ).slice().buffer;
  const envelope = request.payload.repairPlanEnvelope;
  const resultWithoutReceipt = {
    analysis: snapshotToolcraftModelWorkerAnalysisSummary(analyzeToolcraftModel(
      document,
      "realtime-mesh",
      TOOLCRAFT_DEFAULT_MODEL_IMPORT_LIMITS,
    )),
    canonicalDocument,
    canonicalDocumentDigest: digestToolcraftModelDocument(document),
    operation: "repair" as const,
    repairPlanDigest: envelope.planDigest,
  };
  return {
    generation: request.generation,
    jobId: request.jobId,
    kind: "result" as const,
    result: {
      ...resultWithoutReceipt,
      receiptDigest: createToolcraftModelWorkerReceiptDigest({
        generation: request.generation,
        jobId: request.jobId,
        repairPlanEnvelope: {
          byteLength: envelope.bytes.byteLength,
          envelopeDigest: envelope.envelopeDigest,
          planDigest: envelope.planDigest,
          version: envelope.version,
        },
        result: resultWithoutReceipt,
      }),
    },
  };
}

function repairInput(jobId: string) {
  const document = createValidModelDocument();
  const plan = fakeRepairPlan();
  const planBytes = encodeToolcraftModelRepairPlanEnvelope(plan);
  return {
    canonicalDocument: encodeToolcraftModelDocument(document),
    canonicalDocumentDigest: digestToolcraftModelDocument(document),
    jobId,
    limits: TOOLCRAFT_DEFAULT_MODEL_IMPORT_LIMITS,
    repairPlanEnvelope: {
      bytes: planBytes.buffer,
      envelopeDigest: digestToolcraftModelWorkerBytes(planBytes),
      planDigest: plan.planDigest,
      version: 1 as const,
    },
  };
}

function emitRepairResult(
  worker: FakeWorker,
  request: RepairRequest,
  response: ReturnType<typeof repairResultFor>,
): void {
  worker.emitMessage({
    generation: request.generation,
    jobId: request.jobId,
    kind: "progress",
    phase: "repairing",
    progress: 0.9,
  });
  worker.emitMessage(response);
}

describe("model worker client repair receipts", () => {
  it.each(modelImportLimitNames)(
    "binds %s into the protected result receipt",
    (limitName) => {
      const request = repairRequest(1, fakeRepairPlan());
      if (request.kind !== "repair") throw new Error("Expected repair request.");
      const authentic = repairResultFor(request);
      const { receiptDigest: _receiptDigest, ...result } = authentic.result;
      const receiptInput = {
        generation: request.generation,
        jobId: request.jobId,
        repairPlanEnvelope: {
          byteLength: request.payload.repairPlanEnvelope.bytes.byteLength,
          envelopeDigest: request.payload.repairPlanEnvelope.envelopeDigest,
          planDigest: request.payload.repairPlanEnvelope.planDigest,
          version: request.payload.repairPlanEnvelope.version,
        },
        result,
      };
      const baseline = createToolcraftModelWorkerReceiptDigest(receiptInput);
      const current = result.analysis.limits[limitName];
      const mutatedLimit = limitName === "maxArchiveCompressionRatio"
        ? current - 0.5
        : current - 1;

      expect(
        createToolcraftModelWorkerReceiptDigest({
          ...receiptInput,
          result: {
            ...result,
            analysis: {
              ...result.analysis,
              limits: {
                ...result.analysis.limits,
                [limitName]: mutatedLimit,
              },
            },
          },
        }),
      ).not.toBe(baseline);
    },
  );

  it("accepts an authentic result bound to the expected input envelope", async () => {
    const worker = new FakeWorker();
    const sha256 = vi.fn(async (bytes: ArrayBuffer | Uint8Array) =>
      digestToolcraftModelWorkerBytes(bytes)
    );
    const onResult = vi.fn();
    const client = createToolcraftModelWorkerClient({
      sha256,
      workerFactory: () => worker,
    });
    const pending = client.repair(repairInput("authentic-repair-receipt"), {
      onResult,
    });
    await vi.waitFor(() => expect(worker.posts).toHaveLength(1));
    expect(sha256).toHaveBeenCalledTimes(2);
    const request = worker.posts[0]!.message;
    if (request.kind !== "repair") throw new Error("Expected repair request.");
    const authentic = repairResultFor(request);
    expect(authentic.result.repairPlanDigest).toBe(
      request.payload.repairPlanEnvelope.planDigest,
    );

    emitRepairResult(worker, request, authentic);

    await expect(pending).resolves.toMatchObject({
      kind: "result",
      result: {
        operation: "repair",
        repairPlanDigest: request.payload.repairPlanEnvelope.planDigest,
      },
    });
    expect(sha256).toHaveBeenCalledTimes(4);
    expect(onResult).toHaveBeenCalledTimes(1);
    expect(worker.terminate).not.toHaveBeenCalled();
  });

  it("rejects an otherwise-valid result with a forged receipt", async () => {
    const worker = new FakeWorker();
    const sha256 = vi.fn(async (bytes: ArrayBuffer | Uint8Array) =>
      digestToolcraftModelWorkerBytes(bytes)
    );
    const onResult = vi.fn();
    const client = createToolcraftModelWorkerClient({
      sha256,
      workerFactory: () => worker,
    });
    const pending = client.repair(repairInput("forged-repair-receipt"), {
      onResult,
    });
    await vi.waitFor(() => expect(worker.posts).toHaveLength(1));
    expect(sha256).toHaveBeenCalledTimes(2);
    const request = worker.posts[0]!.message;
    if (request.kind !== "repair") throw new Error("Expected repair request.");
    const authentic = repairResultFor(request);
    expect(authentic.result.repairPlanDigest).toBe(
      request.payload.repairPlanEnvelope.planDigest,
    );

    emitRepairResult(worker, request, {
      ...authentic,
      result: { ...authentic.result, receiptDigest: forgedDigest },
    });

    await expect(pending).resolves.toMatchObject({
      feedback: { code: "model-worker-invalid-message" },
      kind: "error",
    });
    expect(sha256).toHaveBeenCalledTimes(4);
    expect(onResult).not.toHaveBeenCalled();
    expect(worker.terminate).toHaveBeenCalledTimes(1);
  });
});