import { describe, it, expect, vi } from "vitest";
import {
  adaptTraceToUiMessages,
  buildToolRenderOverridesFromSnapshots,
  type TraceEnvelope,
  type TraceMessage,
  type TraceWidgetSnapshot,
} from "../trace-viewer-adapter";

// Mock external dependencies used by the adapter
vi.mock("@mcp-ui/client", () => ({
  isUIResource: (value: Record<string, unknown>) =>
    !!value?.resource &&
    typeof (value.resource as Record<string, unknown>).uri === "string" &&
    (value.resource as Record<string, unknown>).uri
      ?.toString()
      .startsWith("ui://"),
}));

vi.mock("@modelcontextprotocol/ext-apps/app-bridge", () => ({
  getToolUiResourceUri: (tool: { _meta?: Record<string, unknown> }) => {
    const uri = (tool._meta?.ui as Record<string, unknown> | undefined)
      ?.resourceUri;
    return typeof uri === "string" ? uri : undefined;
  },
}));

// Helper to build a minimal widget snapshot
function makeWidgetSnapshot(
  overrides: Partial<TraceWidgetSnapshot> = {},
): TraceWidgetSnapshot {
  return {
    toolCallId: "call-1",
    toolName: "create_view",
    protocol: "mcp-apps",
    serverId: "server-1",
    resourceUri: "ui://widget/create-view.html",
    toolMetadata: { ui: { resourceUri: "ui://widget/create-view.html" } },
    widgetCsp: null,
    widgetPermissions: null,
    widgetPermissive: true,
    prefersBorder: true,
    ...overrides,
  };
}

describe("adaptTraceToUiMessages", () => {
  // --- Test 1: Text + tool-call + tool-result grouping ---
  it("groups assistant text + tool-call with matching tool-result into a single UIMessage", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "assistant",
          content: [
            { type: "text", text: "Let me do that." },
            {
              type: "tool-call",
              toolCallId: "call-1",
              toolName: "get_data",
              input: { id: 42 },
            },
          ],
        },
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolCallId: "call-1",
              toolName: "get_data",
              output: { type: "json", value: { text: "result text" } },
            },
          ],
        },
      ],
    };

    const result = adaptTraceToUiMessages({ trace });
    expect(result.messages).toHaveLength(1);

    const msg = result.messages[0];
    expect(msg.role).toBe("assistant");

    // text, dynamic-tool, text-fallback
    expect(msg.parts.length).toBeGreaterThanOrEqual(3);
    expect(msg.parts[0]).toEqual({ type: "text", text: "Let me do that." });
    expect(msg.parts[1]).toMatchObject({
      type: "dynamic-tool",
      toolCallId: "call-1",
      toolName: "get_data",
      state: "output-available",
    });
    // text fallback from tool result
    expect(msg.parts[2]).toMatchObject({ type: "text" });
  });

  it("attaches readable tool output to the tool part in attached-to-tool mode", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "assistant",
          content: [
            { type: "text", text: "Let me do that." },
            {
              type: "tool-call",
              toolCallId: "call-1",
              toolName: "read_me",
              input: { id: 42 },
            },
          ],
        },
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolCallId: "call-1",
              toolName: "read_me",
              output: { type: "json", value: { text: "result text" } },
            },
          ],
        },
      ],
    };

    const result = adaptTraceToUiMessages({
      trace,
      toolResultDisplay: "attached-to-tool",
    });
    expect(result.messages).toHaveLength(1);

    const msg = result.messages[0];
    expect(msg.parts).toHaveLength(2);
    expect(msg.parts[1]).toMatchObject({
      type: "dynamic-tool",
      toolCallId: "call-1",
      toolName: "read_me",
      traceDisplayText: "result text",
      traceDisplayMode: "markdown",
    });
    expect(
      msg.parts.find((part) => part.type === "text" && part !== msg.parts[0]),
    ).toBeUndefined();
  });

  it("attaches structured tool output as json markdown in attached-to-tool mode", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "assistant",
          content: [
            {
              type: "tool-call",
              toolCallId: "call-1",
              toolName: "read_me",
              input: {},
            },
          ],
        },
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolCallId: "call-1",
              toolName: "read_me",
              output: { hello: "world" },
            },
          ],
        },
      ],
    };

    const result = adaptTraceToUiMessages({
      trace,
      toolResultDisplay: "attached-to-tool",
    });

    expect(result.messages[0].parts).toHaveLength(1);
    expect(result.messages[0].parts[0]).toMatchObject({
      type: "dynamic-tool",
      traceDisplayMode: "json-markdown",
    });
    expect((result.messages[0].parts[0] as any).traceDisplayText).toContain(
      '"hello": "world"',
    );
  });

  it("keeps the raw result in the shared tool card without a sibling", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "assistant",
          content: [
            {
              type: "tool-call",
              toolCallId: "call-1",
              toolName: "read_me",
              input: { id: 42 },
            },
          ],
        },
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolCallId: "call-1",
              toolName: "read_me",
              output: { type: "json", value: { hello: "world" } },
            },
          ],
        },
      ],
    };

    const result = adaptTraceToUiMessages({
      trace,
      toolResultDisplay: "tool-card",
    });

    expect(result.messages[0].parts).toHaveLength(1);
    expect(result.messages[0].parts[0]).toMatchObject({
      type: "dynamic-tool",
      input: { id: 42 },
      output: { hello: "world" },
    });
    expect(result.messages[0].parts[0]).not.toHaveProperty("traceDisplayMode");
  });

  // --- Test 2: Multiple tool calls ---
  it("groups multiple tool-calls and results into a single assistant UIMessage", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "assistant",
          content: [
            {
              type: "tool-call",
              toolCallId: "call-1",
              toolName: "tool_a",
              input: {},
            },
            {
              type: "tool-call",
              toolCallId: "call-2",
              toolName: "tool_b",
              input: {},
            },
          ],
        },
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolCallId: "call-1",
              toolName: "tool_a",
              output: "result-a",
            },
            {
              type: "tool-result",
              toolCallId: "call-2",
              toolName: "tool_b",
              output: "result-b",
            },
          ],
        },
      ],
    };

    const result = adaptTraceToUiMessages({ trace });
    expect(result.messages).toHaveLength(1);

    const msg = result.messages[0];
    const dynamicTools = msg.parts.filter((p) => p.type === "dynamic-tool");
    expect(dynamicTools).toHaveLength(2);
    expect(dynamicTools[0]).toMatchObject({
      toolCallId: "call-1",
      toolName: "tool_a",
    });
    expect(dynamicTools[1]).toMatchObject({
      toolCallId: "call-2",
      toolName: "tool_b",
    });
  });

  // --- Test 3: Orphan tool-results ---
  it("creates synthetic assistant UIMessage for orphan tool-results", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolCallId: "orphan-1",
              toolName: "orphan_tool",
              output: "orphan output",
            },
          ],
        },
      ],
    };

    const result = adaptTraceToUiMessages({ trace });
    expect(result.messages).toHaveLength(1);

    const msg = result.messages[0];
    expect(msg.role).toBe("assistant");
    expect(msg.parts.some((p) => p.type === "dynamic-tool")).toBe(true);
  });

  // --- Test 4: String content ---
  it("normalizes string content into a text part", () => {
    const trace: TraceMessage[] = [{ role: "user", content: "Hello" }];

    const result = adaptTraceToUiMessages({ trace });
    expect(result.messages).toHaveLength(1);

    const msg = result.messages[0];
    expect(msg.role).toBe("user");
    expect(msg.parts).toEqual([{ type: "text", text: "Hello" }]);
  });

  // --- Test 5: Missing toolCallId ---
  it("assigns deterministic synthetic ID when toolCallId is missing", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "assistant",
          content: [
            {
              type: "tool-call",
              toolName: "my_tool",
              input: {},
            },
          ],
        },
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolName: "my_tool",
              output: "done",
            },
          ],
        },
      ],
    };

    const result = adaptTraceToUiMessages({ trace });
    const toolPart = result.messages[0].parts.find(
      (p) => p.type === "dynamic-tool",
    );
    expect(toolPart).toBeDefined();
    expect((toolPart as any).toolCallId).toBe("trace-tool-0-0-my_tool");
  });

  it("maps source indices to canonical focus message ids for user, assistant, and tool messages", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "user",
          content: "Draw me a diagram",
        },
        {
          role: "assistant",
          content: [
            {
              type: "tool-call",
              toolCallId: "call-1",
              toolName: "create_view",
              input: { title: "Architecture" },
            },
          ],
        },
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolCallId: "call-1",
              toolName: "create_view",
              output: { ok: true },
            },
          ],
        },
        {
          role: "assistant",
          content: [{ type: "text", text: "Saved." }],
        },
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolCallId: "orphan-1",
              toolName: "read_me",
              output: "orphan output",
            },
          ],
        },
      ],
    };

    const result = adaptTraceToUiMessages({ trace });

    expect(result.sourceMessageIndexToUiMessageIds[0]).toEqual([
      "trace-user-0",
    ]);
    expect(result.sourceMessageIndexToFocusUiMessageId[0]).toBe("trace-user-0");

    // Tool-bearing messages are re-keyed to `trace-tool-<toolCallId>` by the
    // inspector adapter; the source-index lookups are remapped to match so
    // "Reveal in Chat" resolves a live id. Non-tool messages keep their id.
    expect(result.sourceMessageIndexToUiMessageIds[1]).toEqual([
      "trace-tool-call-1",
    ]);
    expect(result.sourceMessageIndexToFocusUiMessageId[1]).toBe(
      "trace-tool-call-1",
    );
    expect(result.sourceMessageIndexToFocusUiMessageId[2]).toBe(
      "trace-tool-call-1",
    );
    expect(result.sourceMessageIndexToFocusUiMessageId[3]).toBe(
      "trace-assistant-3",
    );
    expect(result.sourceMessageIndexToFocusUiMessageId[4]).toBe(
      "trace-tool-orphan-1",
    );
  });

  // --- Test 6: Deterministic synthetic IDs ---
  it("produces identical output for identical input across calls and does not mutate input", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "assistant",
          content: [
            {
              type: "tool-call",
              toolName: "stable_tool",
              input: { x: 1 },
            },
          ],
        },
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolName: "stable_tool",
              output: "stable output",
            },
          ],
        },
      ],
    };

    const traceBefore = JSON.parse(JSON.stringify(trace));

    const result1 = adaptTraceToUiMessages({ trace });
    const result2 = adaptTraceToUiMessages({ trace });

    expect(result1).toEqual(result2);

    // The original trace must not have been mutated
    expect(trace).toEqual(traceBefore);
  });

  // --- Test 7: _meta / _serverId preservation ---
  it("preserves _serverId from tool-result output", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "assistant",
          content: [
            {
              type: "tool-call",
              toolCallId: "call-srv",
              toolName: "srv_tool",
              input: {},
            },
          ],
        },
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolCallId: "call-srv",
              toolName: "srv_tool",
              output: {
                type: "json",
                value: { _serverId: "server-original", data: 1 },
              },
            },
          ],
        },
      ],
    };

    const result = adaptTraceToUiMessages({ trace });
    const toolPart = result.messages[0].parts.find(
      (p) => p.type === "dynamic-tool",
    ) as any;
    expect(toolPart.output._serverId).toBe("server-original");
  });

  it("merges serverId field from tool-result part onto output", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "assistant",
          content: [
            {
              type: "tool-call",
              toolCallId: "call-merge",
              toolName: "merge_tool",
              input: {},
            },
          ],
        },
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolCallId: "call-merge",
              toolName: "merge_tool",
              serverId: "merged-server",
              result: { data: "hello" },
            },
          ],
        },
      ],
    };

    const result = adaptTraceToUiMessages({ trace });
    const toolPart = result.messages[0].parts.find(
      (p) => p.type === "dynamic-tool",
    ) as any;
    expect(toolPart.output._serverId).toBe("merged-server");
  });

  // --- Test 8: Text extraction vs fenced-JSON fallback ---
  it("appends extracted text when tool output has extractable text", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "assistant",
          content: [
            {
              type: "tool-call",
              toolCallId: "call-text",
              toolName: "text_tool",
              input: {},
            },
          ],
        },
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolCallId: "call-text",
              toolName: "text_tool",
              output: {
                type: "json",
                value: {
                  content: [{ type: "text", text: "extracted content" }],
                },
              },
            },
          ],
        },
      ],
    };

    const result = adaptTraceToUiMessages({ trace });
    const textParts = result.messages[0].parts.filter((p) => p.type === "text");
    expect(textParts.some((p) => (p as any).text === "extracted content")).toBe(
      true,
    );
  });

  it("renders JSON text-block output as a structured result", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "assistant",
          content: [
            {
              type: "tool-call",
              toolCallId: "call-json-text",
              toolName: "list_searches",
              input: {},
            },
          ],
        },
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolCallId: "call-json-text",
              toolName: "list_searches",
              output: {
                type: "json",
                value: {
                  content: [
                    {
                      type: "text",
                      text: '{"searches":[{"id":"admissions-community"}],"nextCursor":null}',
                    },
                  ],
                },
              },
            },
          ],
        },
      ],
    };

    const result = adaptTraceToUiMessages({ trace });
    expect(result.messages[0].parts).toContainEqual({
      type: "data-result",
      data: {
        searches: [{ id: "admissions-community" }],
        nextCursor: null,
      },
    });
    expect(
      result.messages[0].parts.some(
        (part) =>
          part.type === "text" &&
          (part as { text?: string }).text?.startsWith('{"searches"'),
      ),
    ).toBe(false);
  });

  it("renders structured tool output as a structured result", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "assistant",
          content: [
            {
              type: "tool-call",
              toolCallId: "call-json",
              toolName: "json_tool",
              input: {},
            },
          ],
        },
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolCallId: "call-json",
              toolName: "json_tool",
              output: {
                type: "json",
                value: { count: 5, items: ["a", "b"] },
              },
            },
          ],
        },
      ],
    };

    const result = adaptTraceToUiMessages({ trace });
    expect(result.messages[0].parts).toContainEqual({
      type: "data-result",
      data: { count: 5, items: ["a", "b"] },
    });
  });

  it("formats JSON text output alongside a replayed widget", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "assistant",
          content: [
            {
              type: "tool-call",
              toolCallId: "call-widget-json",
              toolName: "create_view",
              input: {},
            },
          ],
        },
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolCallId: "call-widget-json",
              toolName: "create_view",
              result: {
                content: [
                  {
                    type: "text",
                    text: '{"team":"Barcelona","teams":["Arsenal","Barcelona"]}',
                  },
                ],
              },
            },
          ],
        },
      ],
      widgetSnapshots: [
        makeWidgetSnapshot({
          toolCallId: "call-widget-json",
          widgetHtmlUrl: "https://storage.example.com/widget.html",
        }),
      ],
    };

    const result = adaptTraceToUiMessages({ trace });
    expect(result.toolRenderOverrides["call-widget-json"]).toBeDefined();
    expect(result.messages[0].parts).toContainEqual({
      type: "data-result",
      data: { team: "Barcelona", teams: ["Arsenal", "Barcelona"] },
    });
  });

  // --- Test 9: Reasoning parts with no state ---
  it("passes through reasoning parts unchanged", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "assistant",
          content: [{ type: "reasoning", text: "thinking..." } as any],
        },
      ],
    };

    const result = adaptTraceToUiMessages({ trace });
    expect(result.messages).toHaveLength(1);

    const reasoningPart = result.messages[0].parts.find(
      (p) => (p as any).type === "reasoning",
    );
    expect(reasoningPart).toBeDefined();
    expect((reasoningPart as any).text).toBe("thinking...");
  });

  // --- Test 10: Offline snapshot replay ---
  it("creates replay override with cachedWidgetHtmlUrl and isOffline when snapshot has widgetHtmlUrl", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "assistant",
          content: [
            {
              type: "tool-call",
              toolCallId: "call-1",
              toolName: "create_view",
              input: { title: "Flow" },
            },
          ],
        },
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolCallId: "call-1",
              toolName: "create_view",
              output: { type: "json", value: { ok: true } },
            },
          ],
        },
      ],
      widgetSnapshots: [
        makeWidgetSnapshot({
          widgetHtmlUrl: "https://storage.example.com/widget.html",
          toolOutput: {
            content: [{ type: "text", text: "stored result" }],
            structuredContent: { checkpointId: "checkpoint-1" },
          },
        }),
      ],
    };

    // No connected servers — offline replay should still work
    const result = adaptTraceToUiMessages({ trace });

    expect(result.toolRenderOverrides["call-1"]).toBeDefined();
    expect(result.toolRenderOverrides["call-1"].cachedWidgetHtmlUrl).toBe(
      "https://storage.example.com/widget.html",
    );
    expect(result.toolRenderOverrides["call-1"].isOffline).toBe(true);
    expect(result.toolRenderOverrides["call-1"].toolOutput).toEqual({
      content: [{ type: "text", text: "stored result" }],
      structuredContent: { checkpointId: "checkpoint-1" },
    });
  });

  it("preserves widget metadata from tool output when result is simplified", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "assistant",
          content: [
            {
              type: "tool-call",
              toolCallId: "call-widget-live",
              toolName: "create_view",
              input: { title: "Dog" },
            },
          ],
        },
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolCallId: "call-widget-live",
              toolName: "create_view",
              serverId: "server-1",
              output: {
                type: "json",
                value: {
                  _meta: {
                    ui: { resourceUri: "ui://widget/create-view.html" },
                  },
                  structuredContent: {
                    checkpointId: "checkpoint-1",
                  },
                },
              },
              result: {
                structuredContent: {
                  checkpointId: "checkpoint-1",
                },
              },
            },
          ],
        },
      ],
    };

    const result = adaptTraceToUiMessages({
      trace,
      connectedServerIds: ["server-1"],
    });
    const toolPart = result.messages[0].parts.find(
      (part) => part.type === "dynamic-tool",
    ) as any;

    expect(toolPart.output._meta.ui.resourceUri).toBe(
      "ui://widget/create-view.html",
    );
    expect(toolPart.output._serverId).toBe("server-1");
    expect(toolPart.output.structuredContent).toEqual({
      checkpointId: "checkpoint-1",
    });
  });

  it("does not invent empty tool output when snapshot output is absent", () => {
    const overrides = buildToolRenderOverridesFromSnapshots([
      makeWidgetSnapshot({
        widgetHtmlUrl: "https://storage.example.com/widget.html",
      }),
    ]);

    expect(overrides["call-1"]).toBeDefined();
    expect(overrides["call-1"].cachedWidgetHtmlUrl).toBe(
      "https://storage.example.com/widget.html",
    );
    expect(overrides["call-1"].toolOutput).toBeUndefined();
  });

  describe("buildToolRenderOverridesFromSnapshots / preferLiveWhenPossible", () => {
    it("sets liveFetchPreferred for mcp-apps snapshots with a resourceUri but keeps the cached URL as fallback", () => {
      const overrides = buildToolRenderOverridesFromSnapshots(
        [
          makeWidgetSnapshot({
            toolCallId: "call-mcp-live",
            protocol: "mcp-apps",
            resourceUri: "ui://widget/create-view.html",
            widgetHtmlUrl: "https://storage.example.com/mcp-widget.html",
          }),
        ],
        { preferLiveWhenPossible: true },
      );

      const override = overrides["call-mcp-live"];
      expect(override).toBeDefined();
      expect(override.liveFetchPreferred).toBe(true);
      expect(override.cachedWidgetHtmlUrl).toBe(
        "https://storage.example.com/mcp-widget.html",
      );
      expect(override.isOffline).toBe(true);
      expect(override.resourceUri).toBe("ui://widget/create-view.html");
    });

    it("leaves liveFetchPreferred unset for mcp-apps snapshots without a resourceUri", () => {
      const overrides = buildToolRenderOverridesFromSnapshots(
        [
          makeWidgetSnapshot({
            toolCallId: "call-mcp-no-uri",
            protocol: "mcp-apps",
            resourceUri: undefined,
            widgetHtmlUrl: "https://storage.example.com/no-uri.html",
          }),
        ],
        { preferLiveWhenPossible: true },
      );

      const override = overrides["call-mcp-no-uri"];
      expect(override.liveFetchPreferred).toBe(false);
      expect(override.cachedWidgetHtmlUrl).toBe(
        "https://storage.example.com/no-uri.html",
      );
      expect(override.isOffline).toBe(true);
    });

    it("leaves liveFetchPreferred unset for openai-apps snapshots", () => {
      const overrides = buildToolRenderOverridesFromSnapshots(
        [
          makeWidgetSnapshot({
            toolCallId: "call-openai",
            protocol: "openai-apps",
            resourceUri: undefined,
            toolMetadata: { "openai/outputTemplate": "__cached__" },
            widgetHtmlUrl: "https://storage.example.com/openai-widget.html",
          }),
        ],
        { preferLiveWhenPossible: true },
      );

      const override = overrides["call-openai"];
      expect(override.liveFetchPreferred).toBe(false);
      expect(override.cachedWidgetHtmlUrl).toBe(
        "https://storage.example.com/openai-widget.html",
      );
      expect(override.isOffline).toBe(true);
    });

    it("defaults to liveFetchPreferred=false when the option is not set", () => {
      const overrides = buildToolRenderOverridesFromSnapshots([
        makeWidgetSnapshot({
          toolCallId: "call-default",
          widgetHtmlUrl: "https://storage.example.com/default.html",
        }),
      ]);

      const override = overrides["call-default"];
      expect(override.liveFetchPreferred).toBe(false);
      expect(override.cachedWidgetHtmlUrl).toBe(
        "https://storage.example.com/default.html",
      );
      expect(override.isOffline).toBe(true);
    });

    it("sets liveFetchPreferred without a cached fallback when the snapshot has no widget HTML", () => {
      const overrides = buildToolRenderOverridesFromSnapshots(
        [
          makeWidgetSnapshot({
            toolCallId: "call-no-url",
            protocol: "mcp-apps",
            resourceUri: "ui://widget/create-view.html",
            widgetHtmlUrl: undefined,
          }),
        ],
        { preferLiveWhenPossible: true },
      );

      const override = overrides["call-no-url"];
      expect(override.liveFetchPreferred).toBe(true);
      expect(override.cachedWidgetHtmlUrl).toBeUndefined();
      expect(override.isOffline).toBe(false);
    });
  });

  it("treats nested result.isError tool-results as output errors", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "assistant",
          content: [
            {
              type: "tool-call",
              toolCallId: "call-error",
              toolName: "create_view",
              input: { elements: "[]" },
            },
          ],
        },
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolCallId: "call-error",
              toolName: "create_view",
              result: {
                isError: true,
                content: [
                  {
                    type: "text",
                    text: "Invalid JSON in elements",
                  },
                ],
              },
            },
          ],
        },
      ],
      widgetSnapshots: [
        makeWidgetSnapshot({
          toolCallId: "call-error",
          widgetHtmlUrl: "https://storage.example.com/error-widget.html",
        }),
      ],
    };

    const result = adaptTraceToUiMessages({ trace });
    const toolPart = result.messages[0].parts.find(
      (part) => part.type === "dynamic-tool",
    ) as any;
    const textParts = result.messages[0].parts.filter(
      (part) => part.type === "text",
    ) as Array<{ text?: string }>;

    expect(toolPart.state).toBe("output-error");
    expect(result.toolRenderOverrides["call-error"]).toBeUndefined();
    expect(
      textParts.some(
        (part) => part.text === "Tool error: Invalid JSON in elements",
      ),
    ).toBe(true);
  });

  it("skips replay for failed widget calls but still replays later successful snapshots", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "assistant",
          content: [
            {
              type: "tool-call",
              toolCallId: "call-error",
              toolName: "create_view",
              input: { elements: '[{"type":"rectangle"}]' },
            },
            {
              type: "tool-call",
              toolCallId: "call-success",
              toolName: "create_view",
              input: { elements: '[{"type":"ellipse"}]' },
            },
          ],
        },
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolCallId: "call-error",
              toolName: "create_view",
              output: {
                type: "json",
                value: {
                  isError: true,
                  content: [
                    {
                      type: "text",
                      text: "Invalid JSON in elements",
                    },
                  ],
                },
              },
            },
            {
              type: "tool-result",
              toolCallId: "call-success",
              toolName: "create_view",
              output: {
                type: "json",
                value: { ok: true },
              },
            },
          ],
        },
      ],
      widgetSnapshots: [
        makeWidgetSnapshot({
          toolCallId: "call-error",
          widgetHtmlUrl: "https://storage.example.com/error-widget.html",
        }),
        makeWidgetSnapshot({
          toolCallId: "call-success",
          widgetHtmlUrl: "https://storage.example.com/success-widget.html",
        }),
      ],
    };

    const result = adaptTraceToUiMessages({ trace });
    const dynamicTools = result.messages[0].parts.filter(
      (part) => part.type === "dynamic-tool",
    ) as any[];

    expect(dynamicTools[0].toolCallId).toBe("call-error");
    expect(dynamicTools[0].state).toBe("output-error");
    expect(dynamicTools[1].toolCallId).toBe("call-success");
    expect(dynamicTools[1].state).toBe("output-available");
    expect(result.toolRenderOverrides["call-error"]).toBeUndefined();
    expect(result.toolRenderOverrides["call-success"]).toBeDefined();
    expect(result.toolRenderOverrides["call-success"].cachedWidgetHtmlUrl).toBe(
      "https://storage.example.com/success-widget.html",
    );
  });

  // --- Test 11: Widget scrub when no replay possible ---
  it("scrubs widget output and sets empty toolMetadata when no replay is possible", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "assistant",
          content: [
            {
              type: "tool-call",
              toolCallId: "call-widget",
              toolName: "widget_tool",
              input: {},
            },
          ],
        },
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolCallId: "call-widget",
              toolName: "widget_tool",
              output: {
                type: "json",
                value: {
                  _meta: {
                    ui: { resourceUri: "ui://test/widget.html" },
                  },
                  content: [
                    {
                      type: "resource",
                      resource: { uri: "ui://test/widget.html", text: "html" },
                    },
                  ],
                },
              },
            },
          ],
        },
      ],
    };

    // No snapshot, no connected servers → scrub
    const result = adaptTraceToUiMessages({
      trace,
      toolsMetadata: {
        widget_tool: { ui: { resourceUri: "ui://test/widget.html" } },
      },
    });

    expect(result.toolRenderOverrides["call-widget"]).toBeDefined();
    expect(result.toolRenderOverrides["call-widget"].toolMetadata).toEqual({});

    // The dynamic-tool output should be scrubbed (no resource content)
    const toolPart = result.messages[0].parts.find(
      (p) => p.type === "dynamic-tool",
    ) as any;
    expect(toolPart.output).toBeDefined();

    // Resource with ui:// URI should be removed
    const outputContent = (toolPart.output as Record<string, unknown>)?.content;
    if (Array.isArray(outputContent)) {
      const uiResources = outputContent.filter((item: any) =>
        item?.resource?.uri?.startsWith("ui://"),
      );
      expect(uiResources).toHaveLength(0);
    }
  });

  it("keeps streamed widget resources when embedded metadata and server id are enough for live replay", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "assistant",
          content: [
            {
              type: "tool-call",
              toolCallId: "call-stream-widget",
              toolName: "create_view",
              input: {},
            },
          ],
        },
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolCallId: "call-stream-widget",
              toolName: "create_view",
              output: {
                type: "json",
                value: {
                  _serverId: "server-1",
                  _meta: {
                    ui: { resourceUri: "ui://test/widget.html" },
                  },
                  content: [
                    {
                      type: "resource",
                      resource: { uri: "ui://test/widget.html", text: "html" },
                    },
                  ],
                },
              },
            },
          ],
        },
      ],
    };

    const result = adaptTraceToUiMessages({
      trace,
      connectedServerIds: ["server-1"],
    });

    expect(result.toolRenderOverrides["call-stream-widget"]).toBeUndefined();
    const toolPart = result.messages[0].parts.find(
      (part) => part.type === "dynamic-tool",
    ) as any;
    const outputContent = toolPart.output?.content;
    expect(Array.isArray(outputContent)).toBe(true);
    expect(
      outputContent?.some(
        (item: any) => item?.resource?.uri === "ui://test/widget.html",
      ),
    ).toBe(true);
  });

  it("scrubs streamed widget resources when the server id exists but is not connected", () => {
    const trace: TraceEnvelope = {
      messages: [
        {
          role: "assistant",
          content: [
            {
              type: "tool-call",
              toolCallId: "call-disconnected-widget",
              toolName: "create_view",
              input: {},
            },
          ],
        },
        {
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolCallId: "call-disconnected-widget",
              toolName: "create_view",
              output: {
                type: "json",
                value: {
                  _serverId: "server-1",
                  _meta: {
                    ui: { resourceUri: "ui://test/widget.html" },
                  },
                  content: [
                    {
                      type: "resource",
                      resource: { uri: "ui://test/widget.html", text: "html" },
                    },
                  ],
                },
              },
            },
          ],
        },
      ],
    };

    const result = adaptTraceToUiMessages({
      trace,
      connectedServerIds: [],
    });

    expect(result.toolRenderOverrides["call-disconnected-widget"]).toBeDefined();
    expect(
      result.toolRenderOverrides["call-disconnected-widget"].toolMetadata,
    ).toEqual({});

    const toolPart = result.messages[0].parts.find(
      (part) => part.type === "dynamic-tool",
    ) as any;
    const outputContent = toolPart.output?.content;
    expect(Array.isArray(outputContent)).toBe(true);
    expect(
      outputContent?.some(
        (item: any) => item?.resource?.uri === "ui://test/widget.html",
      ),
    ).toBe(false);
  });
});