import { afterEach, describe, expect, it, vi } from "vitest";
import type { DiffPreview, Message } from "../src/dimsdk/core/model.ts";

const { computeDiffPreviewWithShiki } = vi.hoisted(() => ({
  computeDiffPreviewWithShiki: vi.fn(),
}));

vi.mock("../src/dimsdk/ui/chat/diff-highlight.ts", () => ({
  computeDiffPreviewWithShiki,
}));

const {
  clearDiffHighlightInFlight,
  scheduleToolCallDiffHighlightByCallId,
} = await import("../src/dimsdk/core/store/diff-highlight.ts");

function createAssistantMessage(parts: Extract<Message, { role: "assistant" }>["parts"]): Extract<Message, { role: "assistant" }> {
  return {
    id: "assistant-1",
    role: "assistant",
    parts,
  };
}

async function waitFor(predicate: () => boolean, timeoutMs = 1000): Promise<void> {
  const deadline = Date.now() + timeoutMs;
  while (Date.now() < deadline) {
    if (predicate()) return;
    await new Promise((resolve) => setTimeout(resolve, 0));
  }
  throw new Error("timed out waiting for highlight update");
}

afterEach(() => {
  clearDiffHighlightInFlight();
  computeDiffPreviewWithShiki.mockReset();
});

describe("diff highlight store wrapper", () => {
  it("reads edit args from camelCase fields and updates the tool call preview", async () => {
    const preview: DiffPreview = {
      kind: "shiki",
      language: "ts",
      theme: "monokai",
      oldString: "before",
      newString: "after",
      lines: [
        { type: "removed", content: "before", tokens: [{ content: "before" }] },
        { type: "added", content: "after", tokens: [{ content: "after" }] },
      ],
    };
    computeDiffPreviewWithShiki.mockResolvedValue(preview);

    let messages: Message[] = [
      createAssistantMessage([
        {
          type: "tool_call",
          call: {
            id: "edit-1",
            name: "Edit",
            arguments: {
              oldString: "before",
              newString: "after",
              filePath: "/tmp/demo.ts",
            },
          },
        },
      ]),
    ];

    const state = {
      get messages() {
        return messages;
      },
      set messages(value: Message[]) {
        messages = value;
      },
      ui: {
        theme: {
          mode: "monokai-dark",
        },
      },
    };

    scheduleToolCallDiffHighlightByCallId({
      state,
      callId: "edit-1",
      updateAssistant(assistantId, updater) {
        state.messages = state.messages.map((message) =>
          message.role === "assistant" && message.id === assistantId
            ? updater(message)
            : message,
        );
      },
    });

    await waitFor(() => {
      const assistant = state.messages[0];
      return (
        assistant?.role === "assistant"
        && assistant.parts[0]?.type === "tool_call"
        && assistant.parts[0].call.diffPreview?.kind === "shiki"
      );
    });

    expect(computeDiffPreviewWithShiki).toHaveBeenCalledWith(
      "before",
      "after",
      expect.objectContaining({
        filePath: "/tmp/demo.ts",
        theme: "monokai",
      }),
    );

    const assistant = state.messages[0];
    expect(assistant).toBeDefined();
    if (!assistant || assistant.role !== "assistant") {
      throw new Error("assistant message missing");
    }
    expect(assistant.parts[0]).toMatchObject({
      type: "tool_call",
      call: {
        id: "edit-1",
        diffPreview: preview,
      },
    });
  });
});
