import { describe, it, expect, beforeEach, vi, afterEach } from "vitest";
import {
  MessageRepository,
  ExportedMessageRepository,
} from "../runtime/utils/message-repository";
import type { ThreadMessage } from "../types/message";
import type { TextMessagePart } from "../types/message";
import type { ThreadMessageLike } from "../runtime/utils/thread-message-like";

// Mock generateId and generateOptimisticId to make tests deterministic
const mockGenerateId = vi.fn();
const mockGenerateOptimisticId = vi.fn();
const mockIsOptimisticId = vi.fn((id: string) =>
  id.startsWith("__optimistic__"),
);

vi.mock("../utils/id", async (importOriginal) => {
  const original = await importOriginal<typeof import("../utils/id")>();
  return {
    ...original,
    generateId: () => mockGenerateId(),
    generateOptimisticId: () => mockGenerateOptimisticId(),
    isOptimisticId: (id: string) => mockIsOptimisticId(id),
  };
});

/**
 * Tests for the MessageRepository class, which manages message threads with branching capabilities.
 *
 * This suite verifies that the repository:
 * - Correctly manages message additions, updates, and deletions
 * - Properly maintains parent-child relationships between messages
 * - Handles branch creation and switching between branches
 * - Successfully imports and exports repository state
 * - Correctly manages optimistic messages in the thread
 * - Handles edge cases and error conditions gracefully
 */
describe("MessageRepository", () => {
  let repository: MessageRepository;
  let nextMockId = 1;

  /**
   * Creates a test ThreadMessage with the given overrides.
   */
  const createTestMessage = (overrides = {}): ThreadMessage => ({
    id: "test-id",
    role: "assistant",
    createdAt: new Date(),
    content: [{ type: "text", text: "Test message" }],
    status: { type: "complete", reason: "stop" },
    metadata: {
      unstable_state: null,
      unstable_annotations: [],
      unstable_data: [],
      steps: [],
      custom: {},
    },
    ...overrides,
  });

  /**
   * Creates a test CoreMessage with the given overrides.
   */
  const createThreadMessageLike = (overrides = {}): ThreadMessageLike => ({
    role: "assistant",
    content: [{ type: "text", text: "Test message" }],
    ...overrides,
  });

  beforeEach(() => {
    repository = new MessageRepository();
    // Reset mocks with predictable counter-based values
    nextMockId = 1;
    mockGenerateId.mockImplementation(() => `mock-id-${nextMockId++}`);
    mockGenerateOptimisticId.mockImplementation(
      () => `__optimistic__mock-id-${nextMockId++}`,
    );
  });

  afterEach(() => {
    vi.clearAllMocks();
  });

  // Core functionality tests - these test the public contract
  describe("Basic CRUD operations", () => {
    it("should add a new message to the repository", () => {
      const message = createTestMessage({ id: "message-id" });
      repository.addOrUpdateMessage(null, message);

      const messages = repository.getMessages();
      expect(messages).toContain(message);
    });

    it("should update an existing message", () => {
      const message = createTestMessage({ id: "message-id" });
      repository.addOrUpdateMessage(null, message);

      const updatedContent = [
        { type: "text", text: "Updated message" },
      ] as const;
      const updatedMessage = createTestMessage({
        id: "message-id",
        content: updatedContent,
      });

      repository.addOrUpdateMessage(null, updatedMessage);

      const retrievedMessage = repository.getMessage("message-id").message;
      expect(retrievedMessage.content).toEqual(updatedContent);
    });

    it("should establish parent-child relationships between messages", () => {
      const parent = createTestMessage({ id: "parent-id" });
      const child = createTestMessage({ id: "child-id" });

      repository.addOrUpdateMessage(null, parent);
      repository.addOrUpdateMessage("parent-id", child);

      const childWithParent = repository.getMessage("child-id");
      expect(childWithParent.parentId).toBe("parent-id");
    });

    it("should throw an error when parent message is not found", () => {
      const message = createTestMessage();

      expect(() => {
        repository.addOrUpdateMessage("non-existent-id", message);
      }).toThrow(/Parent message not found/);
    });

    it("should retrieve all messages in the current branch", () => {
      const parent = createTestMessage({ id: "parent-id" });
      const child = createTestMessage({ id: "child-id" });
      const grandchild = createTestMessage({ id: "grandchild-id" });

      repository.addOrUpdateMessage(null, parent);
      repository.addOrUpdateMessage("parent-id", child);
      repository.addOrUpdateMessage("child-id", grandchild);

      const messages = repository.getMessages();

      expect(messages.map((m) => m.id)).toEqual([
        "parent-id",
        "child-id",
        "grandchild-id",
      ]);
    });

    it("should track the head message", () => {
      const parent = createTestMessage({ id: "parent-id" });
      const child = createTestMessage({ id: "child-id" });

      repository.addOrUpdateMessage(null, parent);
      expect(repository.headId).toBe("parent-id");

      repository.addOrUpdateMessage("parent-id", child);
      expect(repository.headId).toBe("child-id");
    });

    it("should delete a message and adjust the head", () => {
      const parent = createTestMessage({ id: "parent-id" });
      const child = createTestMessage({ id: "child-id" });

      repository.addOrUpdateMessage(null, parent);
      repository.addOrUpdateMessage("parent-id", child);

      expect(repository.headId).toBe("child-id");

      repository.deleteMessage("child-id");

      expect(repository.headId).toBe("parent-id");

      const messages = repository.getMessages();
      expect(messages.map((m) => m.id)).toEqual(["parent-id"]);
    });

    it("should clear all messages", () => {
      const message = createTestMessage();
      repository.addOrUpdateMessage(null, message);

      repository.clear();

      expect(repository.getMessages()).toHaveLength(0);
      expect(repository.headId).toBeNull();
    });
  });

  describe("Branch management", () => {
    it("should create multiple branches from a parent message", () => {
      const parent = createTestMessage({ id: "parent-id" });
      const branch1 = createTestMessage({ id: "branch1-id" });
      const branch2 = createTestMessage({ id: "branch2-id" });

      repository.addOrUpdateMessage(null, parent);
      repository.addOrUpdateMessage("parent-id", branch1);
      repository.addOrUpdateMessage("parent-id", branch2);

      repository.switchToBranch("branch1-id");
      expect(repository.headId).toBe("branch1-id");

      repository.switchToBranch("branch2-id");
      expect(repository.headId).toBe("branch2-id");

      const branches = repository.getBranches("branch1-id");
      expect(branches).toContain("branch1-id");
      expect(branches).toContain("branch2-id");
    });

    it("should switch between branches and maintain branch state", () => {
      const parent = createTestMessage({ id: "parent-id" });
      const branch1 = createTestMessage({ id: "branch1-id" });
      const branch2 = createTestMessage({ id: "branch2-id" });

      repository.addOrUpdateMessage(null, parent);
      repository.addOrUpdateMessage("parent-id", branch1);
      repository.addOrUpdateMessage("parent-id", branch2);

      repository.switchToBranch("branch1-id");
      expect(repository.headId).toBe("branch1-id");

      const messages1 = repository.getMessages();
      expect(messages1.map((m) => m.id)).toEqual(["parent-id", "branch1-id"]);

      repository.switchToBranch("branch2-id");
      expect(repository.headId).toBe("branch2-id");

      const messages2 = repository.getMessages();
      expect(messages2.map((m) => m.id)).toEqual(["parent-id", "branch2-id"]);
    });

    it("should throw error when switching to a non-existent branch", () => {
      expect(() => {
        repository.switchToBranch("non-existent-id");
      }).toThrow(/Branch not found/);
    });

    it("should reset head to an earlier message in the tree", () => {
      const parent = createTestMessage({ id: "parent-id" });
      const child = createTestMessage({ id: "child-id" });
      const grandchild = createTestMessage({ id: "grandchild-id" });

      repository.addOrUpdateMessage(null, parent);
      repository.addOrUpdateMessage("parent-id", child);
      repository.addOrUpdateMessage("child-id", grandchild);

      repository.resetHead("parent-id");

      expect(repository.headId).toBe("parent-id");

      const messages = repository.getMessages();
      expect(messages.map((m) => m.id)).toEqual(["parent-id"]);
    });

    it("should remove children when resetting head to a message with children", () => {
      const parent = createTestMessage({ id: "parent-id" });
      const child = createTestMessage({ id: "child-id" });
      const grandchild1 = createTestMessage({ id: "grandchild1-id" });
      const grandchild2 = createTestMessage({ id: "grandchild2-id" });
      const greatGrandchild = createTestMessage({ id: "greatgrandchild-id" });

      repository.addOrUpdateMessage(null, parent);
      repository.addOrUpdateMessage("parent-id", child);
      repository.addOrUpdateMessage("child-id", grandchild1);
      repository.addOrUpdateMessage("child-id", grandchild2);
      repository.addOrUpdateMessage("grandchild2-id", greatGrandchild);

      repository.resetHead("child-id");

      expect(repository.headId).toBe("child-id");

      const messages = repository.getMessages();
      expect(messages.map((m) => m.id)).toEqual(["parent-id", "child-id"]);

      expect(() => repository.getMessage("grandchild1-id")).toThrow(
        /Message not found/,
      );
      expect(() => repository.getMessage("grandchild2-id")).toThrow(
        /Message not found/,
      );
      expect(() => repository.getMessage("greatgrandchild-id")).toThrow(
        /Message not found/,
      );

      const branches = repository.getBranches("child-id");
      expect(branches).toEqual(["child-id"]);
    });

    it("should reset head to null when null is passed", () => {
      const message = createTestMessage();
      repository.addOrUpdateMessage(null, message);

      repository.resetHead(null);

      expect(repository.headId).toBeNull();
      expect(repository.getMessages()).toHaveLength(0);
    });
  });

  describe("Optimistic messages", () => {
    it("should create an optimistic message with a unique ID", () => {
      mockGenerateOptimisticId.mockReturnValue("__optimistic__generated-id");

      const coreMessage = createThreadMessageLike();
      const optimisticId = repository.appendOptimisticMessage(
        null,
        coreMessage,
      );

      expect(optimisticId).toBe("__optimistic__generated-id");
      expect(repository.getMessage(optimisticId).message.status?.type).toBe(
        "running",
      );
    });

    it("should create an optimistic message as a child of a specified parent", () => {
      const parent = createTestMessage({ id: "parent-id" });
      repository.addOrUpdateMessage(null, parent);

      const coreMessage = createThreadMessageLike();
      const optimisticId = repository.appendOptimisticMessage(
        "parent-id",
        coreMessage,
      );

      const result = repository.getMessage(optimisticId);
      expect(result.parentId).toBe("parent-id");
    });

    it("should retry generating unique optimistic IDs if initial one exists", () => {
      mockGenerateOptimisticId.mockReturnValueOnce("__optimistic__existing-id");

      const existingMessage = createTestMessage({
        id: "__optimistic__existing-id",
      });
      repository.addOrUpdateMessage(null, existingMessage);

      mockGenerateOptimisticId.mockReturnValueOnce("__optimistic__unique-id");

      const coreMessage = createThreadMessageLike();
      const optimisticId = repository.appendOptimisticMessage(
        null,
        coreMessage,
      );

      expect(optimisticId).toBe("__optimistic__unique-id");
      expect(mockGenerateOptimisticId).toHaveBeenCalledTimes(2);
    });
  });

  describe("Export and import", () => {
    it("should export the repository state", () => {
      const parent = createTestMessage({ id: "parent-id" });
      const child = createTestMessage({ id: "child-id" });

      repository.addOrUpdateMessage(null, parent);
      repository.addOrUpdateMessage("parent-id", child);

      const exported = repository.export();

      expect(exported.headId).toBe("child-id");
      expect(exported.messages).toHaveLength(2);
      expect(
        exported.messages.find((m) => m.message.id === "parent-id")?.parentId,
      ).toBeNull();
      expect(
        exported.messages.find((m) => m.message.id === "child-id")?.parentId,
      ).toBe("parent-id");
    });

    it("should import repository state", () => {
      const parent = createTestMessage({ id: "parent-id" });
      const child = createTestMessage({ id: "child-id" });

      const exported = {
        headId: "child-id",
        messages: [
          { message: parent, parentId: null },
          { message: child, parentId: "parent-id" },
        ],
      };

      repository.import(exported);

      expect(repository.headId).toBe("child-id");
      const messages = repository.getMessages();
      expect(messages.map((m) => m.id)).toEqual(["parent-id", "child-id"]);
    });

    it("should import with a specified head that is not the most recent message", () => {
      const parent = createTestMessage({ id: "parent-id" });
      const child1 = createTestMessage({ id: "child1-id" });
      const child2 = createTestMessage({ id: "child2-id" });

      const exported = {
        headId: "child1-id",
        messages: [
          { message: parent, parentId: null },
          { message: child1, parentId: "parent-id" },
          { message: child2, parentId: "parent-id" },
        ],
      };

      repository.import(exported);

      expect(repository.headId).toBe("child1-id");

      const messages = repository.getMessages();
      expect(messages.map((m) => m.id)).toEqual(["parent-id", "child1-id"]);

      repository.switchToBranch("child2-id");
      expect(repository.headId).toBe("child2-id");
    });

    it("should throw an error when importing with invalid parent references", () => {
      const child = createTestMessage({ id: "child-id" });

      const exported = {
        headId: "child-id",
        messages: [{ message: child, parentId: "non-existent-id" }],
      };

      expect(() => {
        repository.import(exported);
      }).toThrow(/Parent message not found/);
    });
  });

  describe("ExportedMessageRepository utility", () => {
    it("should convert an array of messages to repository format", () => {
      mockGenerateId.mockReturnValue("generated-id");

      const messages: ThreadMessageLike[] = [
        {
          role: "user" as const,
          content: [
            { type: "text" as const, text: "Hello" },
          ] as TextMessagePart[],
        },
        {
          role: "assistant" as const,
          content: [
            { type: "text" as const, text: "Hi there" },
          ] as TextMessagePart[],
        },
      ];

      const result = ExportedMessageRepository.fromArray(messages);

      expect(result.messages).toHaveLength(2);
      expect(result.messages[0]!.parentId).toBeNull();
      expect(result.messages[1]!.parentId).toBe("generated-id");
    });

    it("should handle empty message arrays", () => {
      const result = ExportedMessageRepository.fromArray([]);
      expect(result.messages).toHaveLength(0);
    });
  });

  describe("ExportedMessageRepository.fromBranchableArray", () => {
    it("should create a branching tree structure", () => {
      const items = [
        {
          message: {
            id: "user-1",
            role: "user" as const,
            content: [{ type: "text" as const, text: "Hello" }],
          },
          parentId: null,
        },
        {
          message: {
            id: "assistant-1",
            role: "assistant" as const,
            content: [{ type: "text" as const, text: "Hi there" }],
          },
          parentId: "user-1",
        },
        {
          message: {
            id: "assistant-2",
            role: "assistant" as const,
            content: [{ type: "text" as const, text: "Hey!" }],
          },
          parentId: "user-1",
        },
      ];

      const result = ExportedMessageRepository.fromBranchableArray(items);

      expect(result.messages).toHaveLength(3);
      expect(result.messages[0]!.parentId).toBeNull();
      expect(result.messages[0]!.message.id).toBe("user-1");
      expect(result.messages[1]!.parentId).toBe("user-1");
      expect(result.messages[1]!.message.id).toBe("assistant-1");
      expect(result.messages[2]!.parentId).toBe("user-1");
      expect(result.messages[2]!.message.id).toBe("assistant-2");
    });

    it("should support headId option", () => {
      const items = [
        {
          message: {
            id: "msg-1",
            role: "user" as const,
            content: [{ type: "text" as const, text: "Hello" }],
          },
          parentId: null,
        },
        {
          message: {
            id: "msg-2a",
            role: "assistant" as const,
            content: [{ type: "text" as const, text: "Branch A" }],
          },
          parentId: "msg-1",
        },
        {
          message: {
            id: "msg-2b",
            role: "assistant" as const,
            content: [{ type: "text" as const, text: "Branch B" }],
          },
          parentId: "msg-1",
        },
      ];

      const result = ExportedMessageRepository.fromBranchableArray(items, {
        headId: "msg-2b",
      });

      expect(result.headId).toBe("msg-2b");
    });

    it("should throw if a message has no id", () => {
      const items = [
        {
          message: {
            role: "user" as const,
            content: [{ type: "text" as const, text: "Hello" }],
          },
          parentId: null,
        },
      ];

      expect(() =>
        ExportedMessageRepository.fromBranchableArray(items),
      ).toThrow(/Each message must have an 'id' field set/);
    });

    it("should handle empty arrays", () => {
      const result = ExportedMessageRepository.fromBranchableArray([]);
      expect(result.messages).toHaveLength(0);
    });

    it("should work with MessageRepository.import for branch switching", () => {
      const items = [
        {
          message: {
            id: "user-1",
            role: "user" as const,
            content: [{ type: "text" as const, text: "Hello" }],
          },
          parentId: null,
        },
        {
          message: {
            id: "assistant-a",
            role: "assistant" as const,
            content: [{ type: "text" as const, text: "Response A" }],
          },
          parentId: "user-1",
        },
        {
          message: {
            id: "assistant-b",
            role: "assistant" as const,
            content: [{ type: "text" as const, text: "Response B" }],
          },
          parentId: "user-1",
        },
      ];

      const repo = ExportedMessageRepository.fromBranchableArray(items, {
        headId: "assistant-a",
      });

      repository.import(repo);

      // Should show branch A
      let messages = repository.getMessages();
      expect(messages.map((m) => m.id)).toEqual(["user-1", "assistant-a"]);

      // Switch to branch B
      repository.switchToBranch("assistant-b");
      messages = repository.getMessages();
      expect(messages.map((m) => m.id)).toEqual(["user-1", "assistant-b"]);
    });
  });

  describe("Complex scenarios", () => {
    it("should maintain tree structure after deletions", () => {
      const root = createTestMessage({ id: "root-id" });
      const nodeA = createTestMessage({ id: "A-id" });
      const nodeB = createTestMessage({ id: "B-id" });
      const nodeC = createTestMessage({ id: "C-id" });

      repository.addOrUpdateMessage(null, root);
      repository.addOrUpdateMessage("root-id", nodeA);
      repository.addOrUpdateMessage("A-id", nodeB);
      repository.addOrUpdateMessage("A-id", nodeC);

      repository.deleteMessage("B-id");

      repository.switchToBranch("C-id");
      expect(repository.headId).toBe("C-id");

      const messages = repository.getMessages();
      expect(messages.map((m) => m.id)).toEqual(["root-id", "A-id", "C-id"]);
    });

    it("should relink children when deleting a middle node", () => {
      const root = createTestMessage({ id: "root-id" });
      const nodeA = createTestMessage({ id: "A-id" });
      const nodeB = createTestMessage({ id: "B-id" });
      const nodeC = createTestMessage({ id: "C-id" });

      repository.addOrUpdateMessage(null, root);
      repository.addOrUpdateMessage("root-id", nodeA);
      repository.addOrUpdateMessage("A-id", nodeB);
      repository.addOrUpdateMessage("B-id", nodeC);

      repository.deleteMessage("B-id", "A-id");

      const c = repository.getMessage("C-id");
      expect(c.parentId).toBe("A-id");

      repository.switchToBranch("C-id");
      const messages = repository.getMessages();

      expect(messages.some((m) => m.id === "root-id")).toBe(true);
      expect(messages.some((m) => m.id === "A-id")).toBe(true);
      expect(messages.some((m) => m.id === "C-id")).toBe(true);
      expect(messages.some((m) => m.id === "B-id")).toBe(false);
    });

    it("should relink multiple children when deleting a parent node", () => {
      const root = createTestMessage({ id: "root-id" });
      const nodeA = createTestMessage({ id: "A-id" });
      const nodeB = createTestMessage({ id: "B-id" });
      const nodeC = createTestMessage({ id: "C-id" });
      const nodeD = createTestMessage({ id: "D-id" });

      repository.addOrUpdateMessage(null, root);
      repository.addOrUpdateMessage("root-id", nodeA);
      repository.addOrUpdateMessage("A-id", nodeB);
      repository.addOrUpdateMessage("A-id", nodeC);
      repository.addOrUpdateMessage("A-id", nodeD);

      repository.deleteMessage("A-id", "root-id");

      expect(repository.getMessage("B-id").parentId).toBe("root-id");
      expect(repository.getMessage("C-id").parentId).toBe("root-id");
      expect(repository.getMessage("D-id").parentId).toBe("root-id");

      repository.switchToBranch("B-id");
      const bMessages = repository.getMessages();
      expect(bMessages.some((m) => m.id === "root-id")).toBe(true);
      expect(bMessages.some((m) => m.id === "B-id")).toBe(true);
      expect(bMessages.some((m) => m.id === "A-id")).toBe(false);

      repository.switchToBranch("C-id");
      const cMessages = repository.getMessages();
      expect(cMessages.some((m) => m.id === "root-id")).toBe(true);
      expect(cMessages.some((m) => m.id === "C-id")).toBe(true);
      expect(cMessages.some((m) => m.id === "A-id")).toBe(false);

      repository.switchToBranch("D-id");
      const dMessages = repository.getMessages();
      expect(dMessages.some((m) => m.id === "root-id")).toBe(true);
      expect(dMessages.some((m) => m.id === "D-id")).toBe(true);
      expect(dMessages.some((m) => m.id === "A-id")).toBe(false);
    });

    it("should preserve message position when updating content", () => {
      const parent = createTestMessage({ id: "parent-id" });
      const child1 = createTestMessage({ id: "child1-id" });
      const child2 = createTestMessage({ id: "child2-id" });

      repository.addOrUpdateMessage(null, parent);
      repository.addOrUpdateMessage("parent-id", child1);
      repository.addOrUpdateMessage("child1-id", child2);

      const updatedChild1 = createTestMessage({
        id: "child1-id",
        content: [{ type: "text", text: "Updated content" }],
      });

      repository.addOrUpdateMessage("parent-id", updatedChild1);

      const messages = repository.getMessages();
      expect(messages.map((m) => m.id)).toEqual([
        "parent-id",
        "child1-id",
        "child2-id",
      ]);

      const MessagePart = messages[1]!.content[0];
      expect(MessagePart.type).toBe("text");
      expect((MessagePart as TextMessagePart).text).toBe("Updated content");
    });

    it("should handle re-parenting when messages are inserted at the start", () => {
      const messageA = createTestMessage({ id: "A" });
      const messageB = createTestMessage({ id: "B" });
      const messageC = createTestMessage({ id: "C" });

      repository.addOrUpdateMessage(null, messageA);
      repository.addOrUpdateMessage("A", messageB);
      repository.addOrUpdateMessage("B", messageC);

      expect(repository.getMessages().map((m) => m.id)).toEqual([
        "A",
        "B",
        "C",
      ]);
      expect(repository.headId).toBe("C");

      const messageX = createTestMessage({ id: "X" });
      const messageY = createTestMessage({ id: "Y" });

      repository.addOrUpdateMessage(null, messageX);
      repository.addOrUpdateMessage("X", messageY);

      repository.addOrUpdateMessage("Y", messageA);

      const messages = repository.getMessages();
      expect(messages.map((m) => m.id)).toEqual(["X", "Y", "A", "B", "C"]);

      expect(repository.getMessage("X").parentId).toBeNull();
      expect(repository.getMessage("Y").parentId).toBe("X");
      expect(repository.getMessage("A").parentId).toBe("Y");
      expect(repository.getMessage("B").parentId).toBe("A");
      expect(repository.getMessage("C").parentId).toBe("B");

      expect(repository.headId).toBe("C");
    });
  });
});