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

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

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

const {
  clearMarkdownCodeHighlightState,
  scheduleAssistantMarkdownCodeHighlight,
} = await import("../src/dimsdk/core/store/markdown-code-highlight.ts");

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

afterEach(() => {
  clearMarkdownCodeHighlightState();
  computeMarkdownCodeBlockPreviewWithShiki.mockReset();
  vi.useRealTimers();
});

describe("markdown code highlight store wrapper", () => {
  it("applies preview blocks for fenced markdown after the debounce window", async () => {
    vi.useFakeTimers();

    const preview: MarkdownCodeBlockPreview = {
      kind: "shiki",
      language: "ts",
      theme: "monokai",
      hash: "preview-hash",
      totalLines: 1,
      lines: [
        {
          content: "const answer = 42",
          tokens: [{ content: "const" }],
        },
      ],
    };
    computeMarkdownCodeBlockPreviewWithShiki.mockResolvedValue(preview);

    let messages: Message[] = [
      createAssistantMessage("```ts\nconst answer = 42\n```"),
    ];

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

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

    await vi.advanceTimersByTimeAsync(150);

    expect(computeMarkdownCodeBlockPreviewWithShiki).toHaveBeenCalledWith(
      "const answer = 42",
      expect.objectContaining({
        language: "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: "markdown",
      codeBlocks: [expect.objectContaining({ kind: "shiki" })],
    });
  });
});
