import { describe, expect, it } from "vitest";
import type { Message } from "../src/dimsdk/core/model.ts";
import { createLocalThinkingPartId } from "../src/dimsdk/core/transcript-part-id.ts";
import { createToolUiActions } from "../src/dimsdk/core/store/tool-ui.ts";

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

describe("tool ui actions", () => {
  it("collapses a tool-call group together with consecutive sibling calls", () => {
    const state = {
      messages: [
        createAssistantMessage("a1", [
          {
            type: "tool_call",
            call: { id: "group", name: "Edit", arguments: {} },
          },
          {
            type: "tool_output",
            toolCallId: "group",
            toolName: "Edit",
            output: "running",
          },
          {
            type: "tool_call",
            call: { id: "child", name: "Edit", arguments: {} },
          },
          {
            type: "tool_result",
            result: { toolCallId: "child", output: "ok", status: "success" },
          },
          {
            type: "tool_call",
            call: { id: "stop", name: "Read", arguments: {} },
          },
        ]),
      ] satisfies Message[],
      selectedToolCallId: null as string | null,
      expandedBashOutputs: new Set<string>(),
      expandedToolOutputLineCounts: new Map<string, number>(),
      unifiedEditDiffToolCallIds: new Set<string>(),
      collapsedToolCallIds: new Set<string>(),
      expandedToolCallGroups: new Set<string>(["group"]),
    };

    const actions = createToolUiActions({
      state,
      updateAssistant() {},
      scheduleToolCallReadHighlightByCallId() {},
      scheduleToolCallDiffHighlightByCallId() {},
    });

    actions.toggleToolCallGroupCollapsed("group");

    expect([...state.collapsedToolCallIds].sort()).toEqual(["child", "group"]);
    expect(state.expandedToolCallGroups.has("group")).toBe(false);
    expect(state.collapsedToolCallIds.has("stop")).toBe(false);
  });

  it("updates all assistant thinking parts through updateAssistant", () => {
    let messages: Message[] = [
      createAssistantMessage("a1", [
        {
          type: "thinking",
          partId: createLocalThinkingPartId("a1", 0),
          title: "Thinking",
          collapsed: false,
          content: "step 1",
        },
        {
          type: "markdown",
          markdown: "done",
        },
      ]),
      {
        id: "u1",
        role: "user",
        content: "hi",
      },
    ];

    const state = {
      get messages() {
        return messages;
      },
      set messages(value: Message[]) {
        messages = value;
      },
      selectedToolCallId: null as string | null,
      expandedBashOutputs: new Set<string>(),
      expandedToolOutputLineCounts: new Map<string, number>(),
      unifiedEditDiffToolCallIds: new Set<string>(),
      collapsedToolCallIds: new Set<string>(),
      expandedToolCallGroups: new Set<string>(),
    };

    const actions = createToolUiActions({
      state,
      updateAssistant(assistantId, updater) {
        state.messages = state.messages.map((message) =>
          message.role === "assistant" && message.id === assistantId
            ? updater(message)
            : message,
        );
      },
      scheduleToolCallReadHighlightByCallId() {},
      scheduleToolCallDiffHighlightByCallId() {},
    });

    actions.setAllThinkingCollapsed(true);

    const assistant = state.messages[0];
    expect(assistant?.role).toBe("assistant");
    if (assistant?.role !== "assistant") {
      throw new Error("assistant message missing");
    }
    expect(assistant.parts[0]).toMatchObject({
      type: "thinking",
      partId: createLocalThinkingPartId("a1", 0),
      collapsed: true,
    });
    expect(assistant.parts[1]).toMatchObject({
      type: "markdown",
      markdown: "done",
    });
    expect(state.messages[1]).toMatchObject({
      role: "user",
      content: "hi",
    });
  });
});
