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

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

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

const {
  clearReadHighlightInFlight,
  scheduleToolCallReadHighlightByCallId,
} = await import("../src/dimsdk/core/store/read-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 read highlight update");
}

afterEach(() => {
  clearReadHighlightInFlight();
  computeReadPreviewWithShiki.mockReset();
});

describe("read highlight store wrapper", () => {
  it("uses the path alias and applies the generated preview to the read tool call", async () => {
    const preview: ReadPreview = {
      kind: "shiki",
      mode: "collapsed",
      language: "ts",
      theme: "monokai",
      hash: "preview-hash",
      totalLines: 2,
      lines: [
        {
          gutter: "1",
          content: "const x = 1",
          tokens: [{ content: "const" }],
        },
        {
          gutter: "2",
          content: "console.log(x)",
          tokens: [{ content: "console" }],
        },
      ],
    };
    computeReadPreviewWithShiki.mockResolvedValue(preview);

    const output = JSON.stringify([
      {
        type: "text",
        text: "1  const x = 1\n2  console.log(x)",
      },
    ]);

    let messages: Message[] = [
      createAssistantMessage([
        {
          type: "tool_call",
          call: {
            id: "read-1",
            name: "Read",
            arguments: {
              path: "/tmp/demo.ts",
            },
          },
        },
        {
          type: "tool_result",
          result: {
            toolCallId: "read-1",
            output,
            status: "success",
          },
        },
      ]),
    ];

    const state = {
      get messages() {
        return messages;
      },
      set messages(value: Message[]) {
        messages = value;
      },
      selectedToolCallId: "read-1",
      expandedBashOutputs: new Set<string>(),
      ui: {
        theme: {
          mode: "monokai-dark",
        },
      },
    };

    scheduleToolCallReadHighlightByCallId({
      state,
      callId: "read-1",
      mode: "collapsed",
      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.readPreview?.kind === "shiki"
      );
    });

    expect(computeReadPreviewWithShiki).toHaveBeenCalledWith(
      output,
      expect.objectContaining({
        filePath: "/tmp/demo.ts",
        mode: "collapsed",
        maxLines: 10,
        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: "read-1",
        readPreview: preview,
      },
    });
  });
});
