import React from "react";
import { fireEvent, screen, waitFor } from "@testing-library/react";
import {
  MockStepwiseAgent,
  renderWithCopilotKit,
  runStartedEvent,
  runFinishedEvent,
  stateSnapshotEvent,
  testId,
  textMessageStartEvent,
  textMessageContentEvent,
  textMessageEndEvent,
} from "../../__tests__/utils/test-helpers";
import { ReactCustomMessageRenderer } from "../../types/react-custom-message-renderer";
import { useCopilotKit } from "../../providers/CopilotKitProvider";
import { useCopilotChatConfiguration } from "../../providers/CopilotChatConfigurationProvider";
import { useAgent } from "../../hooks/use-agent";
import { CopilotKitCoreReact } from "../../lib/react-core";
import { Message } from "@ag-ui/core";

// Test shim: some environments lack setCredentials on CopilotKitCoreReact.
if (!(CopilotKitCoreReact.prototype as any).setCredentials) {
  (CopilotKitCoreReact.prototype as any).setCredentials = () => {};
}

type SnapshotRendererProps = {
  message: Message;
  position: "before" | "after";
  runId: string;
  messageIndex: number;
  messageIndexInRun: number;
  numberOfMessagesInRun: number;
  agentId: string;
  stateSnapshot: any;
};

const SnapshotRenderer: React.FC<SnapshotRendererProps> = ({
  position,
  message,
  runId,
  stateSnapshot,
}) => {
  if (position !== "after" || message.role !== "assistant") {
    return null;
  }

  const { copilotkit } = useCopilotKit();
  const config = useCopilotChatConfiguration();
  const typedSnapshot = stateSnapshot as { history?: number[] } | undefined;
  const runHistory = typedSnapshot?.history ?? [];

  let count: number | undefined;
  if (config) {
    const runIds = copilotkit.getRunIdsForThread(
      config.agentId,
      config.threadId,
    );
    const runIndex = runIds.indexOf(runId);
    if (runIndex >= 0 && runIndex < runHistory.length) {
      count = runHistory[runIndex];
    }
  }

  if (count === undefined) {
    count = runHistory[runHistory.length - 1];
  }

  return (
    <div data-testid={`state-${message.id}`} data-run-id={runId}>
      State: {count ?? "null"}
    </div>
  );
};

const LiveStateRenderer: React.FC<SnapshotRendererProps> = ({
  messageIndexInRun,
  position,
}) => {
  const { agent } = useAgent();
  const currentStep = (agent?.state as { current_step?: string } | undefined)
    ?.current_step;

  if (position !== "after") {
    return null;
  }
  if (messageIndexInRun !== 0) {
    return null;
  }
  if (!currentStep) {
    return null;
  }

  return <div data-testid="live-step">{currentStep}</div>;
};

describe("CopilotKitProvider custom message renderers E2E", () => {
  it("renders state snapshots before assistant text starts", async () => {
    const agent = new MockStepwiseAgent();

    const customRenderer: ReactCustomMessageRenderer = {
      render: LiveStateRenderer,
    };

    renderWithCopilotKit({
      agent,
      renderCustomMessages: [customRenderer],
    });

    const input = await screen.findByRole("textbox");

    fireEvent.change(input, { target: { value: "Who am I?" } });
    fireEvent.keyDown(input, { key: "Enter", code: "Enter" });

    await waitFor(() => {
      expect(screen.getByText("Who am I?")).toBeDefined();
    });

    agent.emit(runStartedEvent());
    agent.emit(stateSnapshotEvent({ current_step: "Processing..." }));

    await waitFor(() => {
      expect(screen.getByTestId("live-step").textContent).toContain(
        "Processing...",
      );
    });
  });

  it("renders stored state snapshots for sequential runs", async () => {
    const agent = new MockStepwiseAgent();
    const history: number[] = [];

    const emitSnapshot = (count: number) => {
      history.push(count);
      agent.emit(stateSnapshotEvent({ history: [...history] }));
    };

    const customRenderer: ReactCustomMessageRenderer = {
      render: SnapshotRenderer,
    };

    renderWithCopilotKit({
      agent,
      renderCustomMessages: [customRenderer],
    });

    const input = await screen.findByRole("textbox");

    const firstAssistantId = testId("assistant-message");
    fireEvent.change(input, { target: { value: "First question" } });
    fireEvent.keyDown(input, { key: "Enter", code: "Enter" });

    await waitFor(() => {
      expect(screen.getByText("First question")).toBeDefined();
    });

    agent.emit(runStartedEvent());
    emitSnapshot(1);
    agent.emit(textMessageStartEvent(firstAssistantId));
    agent.emit(textMessageContentEvent(firstAssistantId, "First answer"));
    agent.emit(textMessageEndEvent(firstAssistantId));
    agent.emit(runFinishedEvent());

    await waitFor(() => {
      expect(
        screen.getByTestId(`state-${firstAssistantId}`).textContent,
      ).toContain("State: 1");
    });
    const firstRunId = screen
      .getByTestId(`state-${firstAssistantId}`)
      .getAttribute("data-run-id");
    expect(firstRunId).toBeTruthy();

    const secondAssistantId = testId("assistant-message");
    fireEvent.change(input, { target: { value: "Second question" } });
    fireEvent.keyDown(input, { key: "Enter", code: "Enter" });

    await waitFor(() => {
      expect(screen.getByText("Second question")).toBeDefined();
    });

    agent.emit(runStartedEvent());
    emitSnapshot(2);
    agent.emit(textMessageStartEvent(secondAssistantId));
    agent.emit(textMessageContentEvent(secondAssistantId, "Second answer"));
    agent.emit(textMessageEndEvent(secondAssistantId));
    agent.emit(runFinishedEvent());
    agent.complete();

    await waitFor(() => {
      expect(
        screen.getByTestId(`state-${secondAssistantId}`).textContent,
      ).toContain("State: 2");
    });
    const secondRunId = screen
      .getByTestId(`state-${secondAssistantId}`)
      .getAttribute("data-run-id");

    expect(secondRunId).not.toBe(firstRunId);

    const firstRunIdAfterSecond = screen
      .getByTestId(`state-${firstAssistantId}`)
      .getAttribute("data-run-id");
    expect(firstRunIdAfterSecond).toBe(firstRunId);

    expect(
      screen.getByTestId(`state-${firstAssistantId}`).textContent,
    ).toContain("State: 1");
  });

  it("renders only at specified position (before vs after)", async () => {
    const agent = new MockStepwiseAgent();
    const positions: string[] = [];

    const PositionRenderer: React.FC<SnapshotRendererProps> = ({
      position,
      message,
    }) => {
      positions.push(position);
      return (
        <div data-testid={`${position}-${message.id}`}>
          {position}: {message.role}
        </div>
      );
    };

    renderWithCopilotKit({
      agent,
      renderCustomMessages: [{ render: PositionRenderer }],
    });

    const input = await screen.findByRole("textbox");
    const messageId = testId("message");

    fireEvent.change(input, { target: { value: "Test" } });
    fireEvent.keyDown(input, { key: "Enter", code: "Enter" });

    await waitFor(() => {
      expect(screen.getByText("Test")).toBeDefined();
    });

    agent.emit(runStartedEvent());
    agent.emit(textMessageStartEvent(messageId));
    agent.emit(textMessageContentEvent(messageId, "Response"));
    agent.emit(textMessageEndEvent(messageId));
    agent.emit(runFinishedEvent());

    await waitFor(() => {
      expect(screen.getByTestId(`before-${messageId}`)).toBeDefined();
      expect(screen.getByTestId(`after-${messageId}`)).toBeDefined();
    });

    expect(screen.getByTestId(`before-${messageId}`).textContent).toBe(
      "before: assistant",
    );
    expect(screen.getByTestId(`after-${messageId}`).textContent).toBe(
      "after: assistant",
    );

    // Verify renderer was called for both positions
    expect(positions).toContain("before");
    expect(positions).toContain("after");
  });

  it("filters by message role correctly", async () => {
    const agent = new MockStepwiseAgent();

    const AssistantOnlyRenderer: React.FC<SnapshotRendererProps> = ({
      message,
      position,
    }) => {
      if (message.role !== "assistant" || position !== "after") return null;
      return (
        <div data-testid={`assistant-badge-${message.id}`}>AI Response</div>
      );
    };

    renderWithCopilotKit({
      agent,
      renderCustomMessages: [{ render: AssistantOnlyRenderer }],
    });

    const input = await screen.findByRole("textbox");
    const assistantId = testId("assistant");

    fireEvent.change(input, { target: { value: "User message" } });
    fireEvent.keyDown(input, { key: "Enter", code: "Enter" });

    await waitFor(() => {
      expect(screen.getByText("User message")).toBeDefined();
    });

    // User message should not have the badge
    const userMessages = screen.queryAllByTestId(/^assistant-badge-/);
    expect(userMessages.length).toBe(0);

    agent.emit(runStartedEvent());
    agent.emit(textMessageStartEvent(assistantId));
    agent.emit(textMessageContentEvent(assistantId, "AI response"));
    agent.emit(textMessageEndEvent(assistantId));
    agent.emit(runFinishedEvent());

    // Assistant message should have the badge
    await waitFor(() => {
      expect(
        screen.getByTestId(`assistant-badge-${assistantId}`),
      ).toBeDefined();
    });
  });

  it("executes multiple renderers in order", async () => {
    const agent = new MockStepwiseAgent();
    const executionOrder: string[] = [];

    const FirstRenderer: React.FC<SnapshotRendererProps> = ({
      message,
      position,
    }) => {
      if (position !== "after" || message.role !== "assistant") return null;
      executionOrder.push("first");
      return <div data-testid={`first-${message.id}`}>First</div>;
    };

    const SecondRenderer: React.FC<SnapshotRendererProps> = ({
      message,
      position,
    }) => {
      if (position !== "after" || message.role !== "assistant") return null;
      executionOrder.push("second");
      return <div data-testid={`second-${message.id}`}>Second</div>;
    };

    renderWithCopilotKit({
      agent,
      renderCustomMessages: [
        { render: FirstRenderer },
        { render: SecondRenderer },
      ],
    });

    const input = await screen.findByRole("textbox");
    const messageId = testId("message");

    fireEvent.change(input, { target: { value: "Test" } });
    fireEvent.keyDown(input, { key: "Enter", code: "Enter" });

    await waitFor(() => {
      expect(screen.getByText("Test")).toBeDefined();
    });

    agent.emit(runStartedEvent());
    agent.emit(textMessageStartEvent(messageId));
    agent.emit(textMessageContentEvent(messageId, "Response"));
    agent.emit(textMessageEndEvent(messageId));
    agent.emit(runFinishedEvent());

    await waitFor(() => {
      expect(screen.getByTestId(`first-${messageId}`)).toBeDefined();
    });

    // Only first renderer should execute since it returns a result
    expect(executionOrder).toEqual(["first"]);
    expect(screen.queryByTestId(`second-${messageId}`)).toBeNull();
  });

  it("respects agent-scoped renderers", async () => {
    const agent1 = new MockStepwiseAgent();
    const agent2 = new MockStepwiseAgent();

    const Agent1Renderer: React.FC<SnapshotRendererProps> = ({
      message,
      position,
    }) => {
      if (position !== "after" || message.role !== "assistant") return null;
      return <div data-testid={`agent1-badge-${message.id}`}>Agent 1</div>;
    };

    const Agent2Renderer: React.FC<SnapshotRendererProps> = ({
      message,
      position,
    }) => {
      if (position !== "after" || message.role !== "assistant") return null;
      return <div data-testid={`agent2-badge-${message.id}`}>Agent 2</div>;
    };

    renderWithCopilotKit({
      agents: { agent1: agent1, agent2: agent2 },
      agentId: "agent1",
      renderCustomMessages: [
        { agentId: "agent1", render: Agent1Renderer },
        { agentId: "agent2", render: Agent2Renderer },
      ],
    });

    const input = await screen.findByRole("textbox");
    const messageId = testId("message");

    fireEvent.change(input, { target: { value: "Test" } });
    fireEvent.keyDown(input, { key: "Enter", code: "Enter" });

    await waitFor(() => {
      expect(screen.getByText("Test")).toBeDefined();
    });

    agent1.emit(runStartedEvent());
    agent1.emit(textMessageStartEvent(messageId));
    agent1.emit(textMessageContentEvent(messageId, "Response"));
    agent1.emit(textMessageEndEvent(messageId));
    agent1.emit(runFinishedEvent());

    await waitFor(() => {
      expect(screen.getByTestId(`agent1-badge-${messageId}`)).toBeDefined();
    });

    // Agent2's renderer should not execute for agent1's messages
    expect(screen.queryByTestId(`agent2-badge-${messageId}`)).toBeNull();
  });

  it("prioritizes agent-specific renderers over global renderers", async () => {
    const agent = new MockStepwiseAgent();

    const GlobalRenderer: React.FC<SnapshotRendererProps> = ({
      message,
      position,
    }) => {
      if (position !== "after" || message.role !== "assistant") return null;
      return <div data-testid={`global-${message.id}`}>Global</div>;
    };

    const SpecificRenderer: React.FC<SnapshotRendererProps> = ({
      message,
      position,
    }) => {
      if (position !== "after" || message.role !== "assistant") return null;
      return <div data-testid={`specific-${message.id}`}>Specific</div>;
    };

    const agentId = "specific-agent";

    renderWithCopilotKit({
      agents: { [agentId]: agent },
      agentId,
      renderCustomMessages: [
        { render: GlobalRenderer },
        { agentId, render: SpecificRenderer },
      ],
    });

    const input = await screen.findByRole("textbox");
    const messageId = testId("message");

    fireEvent.change(input, { target: { value: "Test" } });
    fireEvent.keyDown(input, { key: "Enter", code: "Enter" });

    await waitFor(() => {
      expect(screen.getByText("Test")).toBeDefined();
    });

    agent.emit(runStartedEvent());
    agent.emit(textMessageStartEvent(messageId));
    agent.emit(textMessageContentEvent(messageId, "Response"));
    agent.emit(textMessageEndEvent(messageId));
    agent.emit(runFinishedEvent());

    await waitFor(() => {
      expect(screen.getByTestId(`specific-${messageId}`)).toBeDefined();
    });

    // Global renderer should not execute since specific one takes precedence
    expect(screen.queryByTestId(`global-${messageId}`)).toBeNull();
  });

  it("handles missing state snapshots gracefully", async () => {
    const agent = new MockStepwiseAgent();

    const StateRenderer: React.FC<SnapshotRendererProps> = ({
      message,
      position,
      stateSnapshot,
    }) => {
      if (position !== "after" || message.role !== "assistant") return null;
      return (
        <div data-testid={`state-${message.id}`}>
          {stateSnapshot ? JSON.stringify(stateSnapshot) : "no-state"}
        </div>
      );
    };

    renderWithCopilotKit({
      agent,
      renderCustomMessages: [{ render: StateRenderer }],
    });

    const input = await screen.findByRole("textbox");
    const messageId = testId("message");

    fireEvent.change(input, { target: { value: "Test" } });
    fireEvent.keyDown(input, { key: "Enter", code: "Enter" });

    await waitFor(() => {
      expect(screen.getByText("Test")).toBeDefined();
    });

    // Don't emit state snapshot
    agent.emit(runStartedEvent());
    agent.emit(textMessageStartEvent(messageId));
    agent.emit(textMessageContentEvent(messageId, "Response"));
    agent.emit(textMessageEndEvent(messageId));
    agent.emit(runFinishedEvent());

    await waitFor(() => {
      expect(screen.getByTestId(`state-${messageId}`)).toBeDefined();
    });

    expect(screen.getByTestId(`state-${messageId}`).textContent).toBe(
      "no-state",
    );
  });

  it("provides correct message index properties", async () => {
    const agent = new MockStepwiseAgent();
    let capturedProps: {
      messageIndex?: number;
      messageIndexInRun?: number;
      numberOfMessagesInRun?: number;
    } | null = null;

    const IndexRenderer: React.FC<SnapshotRendererProps> = (props) => {
      const {
        message,
        position,
        messageIndex,
        messageIndexInRun,
        numberOfMessagesInRun,
      } = props;
      if (position !== "after" || message.role !== "assistant") return null;

      capturedProps = {
        messageIndex,
        messageIndexInRun,
        numberOfMessagesInRun,
      };

      return (
        <div data-testid={`index-${message.id}`}>
          {messageIndex}/{messageIndexInRun}/{numberOfMessagesInRun}
        </div>
      );
    };

    renderWithCopilotKit({
      agent,
      renderCustomMessages: [{ render: IndexRenderer }],
    });

    const input = await screen.findByRole("textbox");
    const msg1 = testId("msg1");
    const msg2 = testId("msg2");

    // First exchange
    fireEvent.change(input, { target: { value: "First" } });
    fireEvent.keyDown(input, { key: "Enter", code: "Enter" });

    await waitFor(() => {
      expect(screen.getByText("First")).toBeDefined();
    });

    agent.emit(runStartedEvent());
    agent.emit(textMessageStartEvent(msg1));
    agent.emit(textMessageContentEvent(msg1, "Response 1"));
    agent.emit(textMessageEndEvent(msg1));
    agent.emit(runFinishedEvent());

    await waitFor(() => {
      expect(screen.getByTestId(`index-${msg1}`)).toBeDefined();
    });

    // Second exchange
    fireEvent.change(input, { target: { value: "Second" } });
    fireEvent.keyDown(input, { key: "Enter", code: "Enter" });

    await waitFor(() => {
      expect(screen.getByText("Second")).toBeDefined();
    });

    agent.emit(runStartedEvent());
    agent.emit(textMessageStartEvent(msg2));
    agent.emit(textMessageContentEvent(msg2, "Response 2"));
    agent.emit(textMessageEndEvent(msg2));
    agent.emit(runFinishedEvent());

    await waitFor(() => {
      expect(screen.getByTestId(`index-${msg2}`)).toBeDefined();
    });

    // Verify the captured props are meaningful
    expect(capturedProps).toBeTruthy();
    if (!capturedProps) {
      throw new Error("Expected capturedProps to be set");
    }
    const props = capturedProps as {
      messageIndex: number;
      messageIndexInRun: number;
      numberOfMessagesInRun: number;
    };
    expect(typeof props.messageIndex).toBe("number");
    expect(typeof props.messageIndexInRun).toBe("number");
    expect(typeof props.numberOfMessagesInRun).toBe("number");
  });

  it("works across multi-turn conversations", async () => {
    const agent = new MockStepwiseAgent();

    const TurnCounter: React.FC<SnapshotRendererProps> = ({
      message,
      position,
      stateSnapshot,
    }) => {
      if (position !== "after" || message.role !== "assistant") return null;
      const snapshot = stateSnapshot as { turn?: number } | undefined;
      const turn = snapshot?.turn ?? 0;
      return <div data-testid={`turn-${message.id}`}>Turn: {turn}</div>;
    };

    renderWithCopilotKit({
      agent,
      renderCustomMessages: [{ render: TurnCounter }],
    });

    const input = await screen.findByRole("textbox");

    // Turn 1
    const msg1 = testId("msg1");
    fireEvent.change(input, { target: { value: "Turn 1" } });
    fireEvent.keyDown(input, { key: "Enter", code: "Enter" });

    await waitFor(() => {
      expect(screen.getByText("Turn 1")).toBeDefined();
    });

    agent.emit(runStartedEvent());
    agent.emit(stateSnapshotEvent({ turn: 1 }));
    agent.emit(textMessageStartEvent(msg1));
    agent.emit(textMessageContentEvent(msg1, "Response 1"));
    agent.emit(textMessageEndEvent(msg1));
    agent.emit(runFinishedEvent());

    await waitFor(() => {
      expect(screen.getByTestId(`turn-${msg1}`).textContent).toBe("Turn: 1");
    });

    // Turn 2
    const msg2 = testId("msg2");
    fireEvent.change(input, { target: { value: "Turn 2" } });
    fireEvent.keyDown(input, { key: "Enter", code: "Enter" });

    await waitFor(() => {
      expect(screen.getByText("Turn 2")).toBeDefined();
    });

    agent.emit(runStartedEvent());
    agent.emit(stateSnapshotEvent({ turn: 2 }));
    agent.emit(textMessageStartEvent(msg2));
    agent.emit(textMessageContentEvent(msg2, "Response 2"));
    agent.emit(textMessageEndEvent(msg2));
    agent.emit(runFinishedEvent());

    await waitFor(() => {
      expect(screen.getByTestId(`turn-${msg2}`).textContent).toBe("Turn: 2");
    });

    // Turn 3
    const msg3 = testId("msg3");
    fireEvent.change(input, { target: { value: "Turn 3" } });
    fireEvent.keyDown(input, { key: "Enter", code: "Enter" });

    await waitFor(() => {
      expect(screen.getByText("Turn 3")).toBeDefined();
    });

    agent.emit(runStartedEvent());
    agent.emit(stateSnapshotEvent({ turn: 3 }));
    agent.emit(textMessageStartEvent(msg3));
    agent.emit(textMessageContentEvent(msg3, "Response 3"));
    agent.emit(textMessageEndEvent(msg3));
    agent.emit(runFinishedEvent());

    await waitFor(() => {
      expect(screen.getByTestId(`turn-${msg3}`).textContent).toBe("Turn: 3");
    });

    // Verify the renderer works across multiple turns
    // All messages should have turn counters rendered
    expect(screen.getByTestId(`turn-${msg1}`)).toBeDefined();
    expect(screen.getByTestId(`turn-${msg2}`)).toBeDefined();
    expect(screen.getByTestId(`turn-${msg3}`)).toBeDefined();
  });

  it("handles renderers returning null without breaking", async () => {
    const agent = new MockStepwiseAgent();

    const NullRenderer: React.FC<SnapshotRendererProps> = () => {
      return null;
    };

    const FallbackRenderer: React.FC<SnapshotRendererProps> = ({
      message,
      position,
    }) => {
      if (position !== "after" || message.role !== "assistant") return null;
      return <div data-testid={`fallback-${message.id}`}>Fallback</div>;
    };

    renderWithCopilotKit({
      agent,
      renderCustomMessages: [
        { render: NullRenderer },
        { render: FallbackRenderer },
      ],
    });

    const input = await screen.findByRole("textbox");
    const messageId = testId("message");

    fireEvent.change(input, { target: { value: "Test" } });
    fireEvent.keyDown(input, { key: "Enter", code: "Enter" });

    await waitFor(() => {
      expect(screen.getByText("Test")).toBeDefined();
    });

    agent.emit(runStartedEvent());
    agent.emit(textMessageStartEvent(messageId));
    agent.emit(textMessageContentEvent(messageId, "Response"));
    agent.emit(textMessageEndEvent(messageId));
    agent.emit(runFinishedEvent());

    // First renderer returns null, so second should not execute (first one doesn't break the chain)
    await waitFor(() => {
      const fallback = screen.queryByTestId(`fallback-${messageId}`);
      expect(fallback).toBeNull();
    });
  });

  it("re-renders custom message when state updates within the same run", async () => {
    const agent = new MockStepwiseAgent();

    const StateCountRenderer: React.FC<SnapshotRendererProps> = ({
      message,
      position,
      stateSnapshot,
    }) => {
      if (position !== "after" || message.role !== "assistant") return null;
      const snapshot = stateSnapshot as { count?: number } | undefined;
      return (
        <div data-testid={`count-${message.id}`}>
          Count: {snapshot?.count ?? "none"}
        </div>
      );
    };

    renderWithCopilotKit({
      agent,
      renderCustomMessages: [{ render: StateCountRenderer }],
    });

    const input = await screen.findByRole("textbox");
    const messageId = testId("message");

    fireEvent.change(input, { target: { value: "Test" } });
    fireEvent.keyDown(input, { key: "Enter", code: "Enter" });

    await waitFor(() => {
      expect(screen.getByText("Test")).toBeDefined();
    });

    // Start run
    agent.emit(runStartedEvent());

    // Emit message
    agent.emit(textMessageStartEvent(messageId));
    agent.emit(textMessageContentEvent(messageId, "Response"));
    agent.emit(textMessageEndEvent(messageId));

    // First state update
    agent.emit(stateSnapshotEvent({ count: 1 }));

    await waitFor(() => {
      expect(screen.getByTestId(`count-${messageId}`).textContent).toBe(
        "Count: 1",
      );
    });

    // Second state update WITHIN THE SAME RUN
    agent.emit(stateSnapshotEvent({ count: 2 }));

    // This assertion will FAIL due to the bug - custom message won't re-render
    await waitFor(() => {
      expect(screen.getByTestId(`count-${messageId}`).textContent).toBe(
        "Count: 2",
      );
    });

    agent.emit(runFinishedEvent());
  });

  it("receives state snapshots from different runs", async () => {
    const agent = new MockStepwiseAgent();
    const receivedSnapshots: Array<{ messageId: string; count: number }> = [];

    const CounterRenderer: React.FC<SnapshotRendererProps> = ({
      message,
      position,
      stateSnapshot,
    }) => {
      if (position !== "after" || message.role !== "assistant") return null;
      const snapshot = stateSnapshot as { count?: number } | undefined;
      const count = snapshot?.count ?? 0;

      // Track what snapshot this message received
      const existing = receivedSnapshots.find(
        (s) => s.messageId === message.id,
      );
      if (!existing) {
        receivedSnapshots.push({ messageId: message.id, count });
      }

      return <div data-testid={`count-${message.id}`}>Count: {count}</div>;
    };

    renderWithCopilotKit({
      agent,
      renderCustomMessages: [{ render: CounterRenderer }],
    });

    const input = await screen.findByRole("textbox");
    const msg1 = testId("msg1");

    fireEvent.change(input, { target: { value: "First" } });
    fireEvent.keyDown(input, { key: "Enter", code: "Enter" });

    await waitFor(() => {
      expect(screen.getByText("First")).toBeDefined();
    });

    agent.emit(runStartedEvent());
    agent.emit(stateSnapshotEvent({ count: 5 }));
    agent.emit(textMessageStartEvent(msg1));
    agent.emit(textMessageContentEvent(msg1, "Response"));
    agent.emit(textMessageEndEvent(msg1));
    agent.emit(runFinishedEvent());

    await waitFor(() => {
      expect(screen.getByTestId(`count-${msg1}`)).toBeDefined();
    });

    // Update state in new run
    const msg2 = testId("msg2");
    fireEvent.change(input, { target: { value: "Second" } });
    fireEvent.keyDown(input, { key: "Enter", code: "Enter" });

    await waitFor(() => {
      expect(screen.getByText("Second")).toBeDefined();
    });

    agent.emit(runStartedEvent());
    agent.emit(stateSnapshotEvent({ count: 10 }));
    agent.emit(textMessageStartEvent(msg2));
    agent.emit(textMessageContentEvent(msg2, "Response 2"));
    agent.emit(textMessageEndEvent(msg2));
    agent.emit(runFinishedEvent());

    await waitFor(() => {
      expect(screen.getByTestId(`count-${msg2}`)).toBeDefined();
    });

    // Verify both messages received state snapshots
    expect(receivedSnapshots.length).toBe(2);
    expect(receivedSnapshots.some((s) => s.messageId === msg1)).toBe(true);
    expect(receivedSnapshots.some((s) => s.messageId === msg2)).toBe(true);
  });
});