import { queryClient } from "@/contexts";
import type { Message } from "@/types/messages";
import {
  addUserMessage,
  findLastBotMessage,
  removePlaceholderUserMessage,
  updateMessage,
  updateMessageProperties,
} from "../message-utils";

const QUERY_KEY = ["useGetMessagesQuery", { id: "flow-1", session_id: "s1" }];

function buildMessage(overrides: Partial<Message> = {}): Message {
  return {
    id: "msg-1",
    text: "hello",
    sender: "Machine",
    sender_name: "AI",
    session_id: "s1",
    flow_id: "flow-1",
    timestamp: new Date().toISOString(),
    files: [],
    ...overrides,
  } as Message;
}

beforeEach(() => {
  queryClient.clear();
});

describe("findLastBotMessage", () => {
  it("should_return_last_bot_message_when_cache_has_messages", () => {
    const msg1 = buildMessage({ id: "m1", sender: "User" });
    const msg2 = buildMessage({ id: "m2", sender: "Machine" });
    const msg3 = buildMessage({ id: "m3", sender: "Machine" });

    queryClient.setQueryData(QUERY_KEY, [msg1, msg2, msg3]);

    const result = findLastBotMessage();

    expect(result).not.toBeNull();
    expect(result!.message.id).toBe("m3");
    expect(result!.queryKey).toEqual(QUERY_KEY);
  });

  it("should_return_null_when_no_bot_messages_exist", () => {
    const userMsg = buildMessage({ id: "m1", sender: "User" });
    queryClient.setQueryData(QUERY_KEY, [userMsg]);

    expect(findLastBotMessage()).toBeNull();
  });

  it("should_return_null_when_cache_is_empty", () => {
    expect(findLastBotMessage()).toBeNull();
  });

  it("should_skip_bot_messages_without_id", () => {
    const msgNoId = buildMessage({
      id: null as unknown as string,
      sender: "Machine",
    });
    const msgWithId = buildMessage({ id: "m2", sender: "Machine" });

    queryClient.setQueryData(QUERY_KEY, [msgNoId, msgWithId]);

    const result = findLastBotMessage();
    expect(result!.message.id).toBe("m2");
  });

  it("should_return_last_bot_message_for_given_flow_and_session_only", () => {
    const s1Key = ["useGetMessagesQuery", { id: "flow-1", session_id: "s1" }];
    const s2Key = ["useGetMessagesQuery", { id: "flow-1", session_id: "s2" }];
    const msgS1 = buildMessage({
      id: "m-s1",
      sender: "Machine",
      session_id: "s1",
    });
    const msgS2 = buildMessage({
      id: "m-s2",
      sender: "Machine",
      session_id: "s2",
    });

    queryClient.setQueryData(s1Key, [msgS1]);
    queryClient.setQueryData(s2Key, [msgS2]);

    const resultS1 = findLastBotMessage("flow-1", "s1");
    const resultS2 = findLastBotMessage("flow-1", "s2");

    expect(resultS1).not.toBeNull();
    expect(resultS1!.message.id).toBe("m-s1");
    expect(resultS2).not.toBeNull();
    expect(resultS2!.message.id).toBe("m-s2");
  });
});

describe("updateMessageProperties", () => {
  it("should_merge_new_properties_into_target_message", () => {
    const msg = buildMessage({
      id: "m1",
      properties: { state: "complete", source: { id: "s" } },
    });
    queryClient.setQueryData(QUERY_KEY, [msg]);

    updateMessageProperties("m1", QUERY_KEY, { build_duration: 1500 });

    const updated = queryClient.getQueryData<Message[]>(QUERY_KEY);
    expect(updated![0].properties).toEqual(
      expect.objectContaining({
        state: "complete",
        build_duration: 1500,
      }),
    );
  });

  it("should_not_modify_other_messages", () => {
    const msg1 = buildMessage({ id: "m1" });
    const msg2 = buildMessage({ id: "m2" });
    queryClient.setQueryData(QUERY_KEY, [msg1, msg2]);

    updateMessageProperties("m1", QUERY_KEY, { build_duration: 999 });

    const updated = queryClient.getQueryData<Message[]>(QUERY_KEY);
    expect(updated![1].properties).toEqual(msg2.properties);
  });

  it("should_handle_message_with_empty_properties", () => {
    const msg = buildMessage({ id: "m1", properties: {} });
    queryClient.setQueryData(QUERY_KEY, [msg]);

    updateMessageProperties("m1", QUERY_KEY, { build_duration: 500 });

    const updated = queryClient.getQueryData<Message[]>(QUERY_KEY);
    expect(updated![0].properties).toEqual(
      expect.objectContaining({ build_duration: 500 }),
    );
  });
});

describe("removePlaceholderUserMessage", () => {
  it("should_remove_placeholder_user_message_when_buildFlow_fails", () => {
    // Arrange — simulate addUserMessage creating a placeholder (id: null)
    const placeholder = buildMessage({
      id: null as unknown as string,
      sender: "User",
      sender_name: "User",
      text: "dsada",
      flow_id: "flow-1",
      session_id: "s1",
    });
    addUserMessage(placeholder);

    // Verify placeholder exists in cache before removal
    const before = queryClient.getQueryData<Message[]>(QUERY_KEY);
    expect(before).toHaveLength(1);
    expect(before![0].id).toBeNull();

    // Act — remove placeholder (simulates error path cleanup)
    removePlaceholderUserMessage("s1", "flow-1");

    // Assert — placeholder should be gone
    const after = queryClient.getQueryData<Message[]>(QUERY_KEY);
    expect(after).toHaveLength(0);
  });

  it("should_not_remove_real_messages_when_removing_placeholder", () => {
    // Arrange — cache has a real message and a placeholder
    const realMessage = buildMessage({
      id: "msg-real",
      sender: "User",
      sender_name: "User",
      text: "real message",
      flow_id: "flow-1",
      session_id: "s1",
    });
    const placeholder = buildMessage({
      id: null as unknown as string,
      sender: "User",
      sender_name: "User",
      text: "placeholder",
      flow_id: "flow-1",
      session_id: "s1",
    });
    queryClient.setQueryData(QUERY_KEY, [realMessage, placeholder]);

    // Act
    removePlaceholderUserMessage("s1", "flow-1");

    // Assert — only real message remains
    const after = queryClient.getQueryData<Message[]>(QUERY_KEY);
    expect(after).toHaveLength(1);
    expect(after![0].id).toBe("msg-real");
  });

  it("should_not_modify_cache_when_no_placeholder_exists", () => {
    // Arrange — only real messages
    const realMessage = buildMessage({
      id: "msg-real",
      sender: "Machine",
      text: "bot response",
      flow_id: "flow-1",
      session_id: "s1",
    });
    queryClient.setQueryData(QUERY_KEY, [realMessage]);

    // Act
    removePlaceholderUserMessage("s1", "flow-1");

    // Assert — nothing changed
    const after = queryClient.getQueryData<Message[]>(QUERY_KEY);
    expect(after).toHaveLength(1);
    expect(after![0].id).toBe("msg-real");
  });
});

describe("updateMessage — streaming token preservation (UI-014)", () => {
  // QA UI-014 (PR #12992): Pokédex Agent runs, but tool invocations do not
  // appear in the Playground "agent-steps" panel. The backend API stream emits
  // them correctly (API-011 PASS) — the loss happens in the FE store: the
  // `token` event handler in `message-event-handler.ts` builds a token Message
  // that omits `content_blocks`/`properties`/`category`, and `updateMessage`
  // replaces the cached message with it. Any tool steps added by prior
  // `add_message` events get clobbered when the first token chunk arrives.

  it("should_preserve_content_blocks_when_streaming_token_overwrites_message", () => {
    // Arrange — an existing bot message with tool steps (delivered by an
    // earlier `add_message` event with `on_tool_start` / `on_tool_end`
    // content blocks). State is "partial" because streaming is still in
    // progress, matching the live `add_message` payload the agent emits.
    const existing = buildMessage({
      id: "msg-1",
      text: "",
      sender: "Machine",
      flow_id: "flow-1",
      session_id: "s1",
      properties: { state: "partial" } as Message["properties"],
      content_blocks: [
        {
          title: "Agent Steps",
          contents: [
            {
              type: "tool_use",
              name: "fetch_pokemon",
              tool_input: { name: "pikachu" },
              output: '{"name":"pikachu","type":"electric"}',
              error: null,
              header: { title: "Executed **fetch_pokemon**", icon: "Hammer" },
              duration: 412,
            },
          ],
        },
      ] as unknown as Message["content_blocks"],
    });
    queryClient.setQueryData(QUERY_KEY, [existing]);

    // Act — a streaming token chunk arrives. The message-event-handler builds
    // a Message that has the partial-state property but does NOT carry the
    // content_blocks (tokens only carry chunk + id + sender/timestamp).
    updateMessage({
      id: "msg-1",
      flow_id: "flow-1",
      session_id: "s1",
      text: "Pika",
      sender: "Machine",
      sender_name: "AI",
      timestamp: existing.timestamp,
      files: [],
      properties: { state: "partial" } as Message["properties"],
    } as Message);

    // Assert — the tool-step content_blocks must survive the token update.
    const updated = queryClient.getQueryData<Message[]>(QUERY_KEY);
    expect(updated).toHaveLength(1);
    expect(updated![0].content_blocks).toBeDefined();
    expect(updated![0].content_blocks).toHaveLength(1);
    expect(updated![0].content_blocks![0].contents[0]).toMatchObject({
      type: "tool_use",
      name: "fetch_pokemon",
    });
  });

  it("should_still_accumulate_text_chunks_after_preserving_content_blocks", () => {
    // Regression guard: the fix for the content_blocks loss must not break
    // the existing token-accumulation behavior asserted elsewhere in the
    // updateMessage path.
    const existing = buildMessage({
      id: "msg-1",
      text: "Hello ",
      sender: "Machine",
      flow_id: "flow-1",
      session_id: "s1",
      properties: { state: "partial" } as Message["properties"],
      content_blocks: [
        { title: "Agent Steps", contents: [] },
      ] as unknown as Message["content_blocks"],
    });
    queryClient.setQueryData(QUERY_KEY, [existing]);

    updateMessage({
      id: "msg-1",
      flow_id: "flow-1",
      session_id: "s1",
      text: "world",
      sender: "Machine",
      sender_name: "AI",
      timestamp: existing.timestamp,
      files: [],
      properties: { state: "partial" } as Message["properties"],
    } as Message);

    const updated = queryClient.getQueryData<Message[]>(QUERY_KEY);
    expect(updated![0].text).toBe("Hello world");
    expect(updated![0].content_blocks).toHaveLength(1);
  });
});
