import { env, exports } from "cloudflare:workers";
import { describe, it, expect } from "vitest";
import { MessageType } from "../types";
import type { UIMessage as ChatMessage } from "ai";
import { connectChatWS } from "./test-utils";
import { getAgentByName } from "agents";

/**
 * Helper: send a CF_AGENT_USE_CHAT_REQUEST and wait for the done response.
 */
function sendChatRequest(
  ws: WebSocket,
  requestId: string,
  messages: ChatMessage[],
  extraBody?: Record<string, unknown>
): Promise<boolean> {
  return new Promise<boolean>((resolve) => {
    const timeout = setTimeout(() => resolve(false), 3000);

    const handler = (e: MessageEvent) => {
      const data = JSON.parse(e.data as string);
      if (
        data.type === MessageType.CF_AGENT_USE_CHAT_RESPONSE &&
        data.id === requestId &&
        data.done
      ) {
        clearTimeout(timeout);
        ws.removeEventListener("message", handler);
        resolve(true);
      }
    };
    ws.addEventListener("message", handler);

    ws.send(
      JSON.stringify({
        type: MessageType.CF_AGENT_USE_CHAT_REQUEST,
        id: requestId,
        init: {
          method: "POST",
          body: JSON.stringify({ messages, ...extraBody })
        }
      })
    );
  });
}

/**
 * Helper: fetch persisted messages via HTTP.
 */
async function fetchPersistedMessages(room: string): Promise<ChatMessage[]> {
  const res = await exports.default.fetch(
    `http://example.com/agents/test-chat-agent/${room}/get-messages`
  );
  expect(res.status).toBe(200);
  return (await res.json()) as ChatMessage[];
}

describe("Regenerate message (reconcile stale rows)", () => {
  it("deletes the stale assistant message when client sends truncated array", async () => {
    const room = crypto.randomUUID();
    const { ws } = await connectChatWS(`/agents/test-chat-agent/${room}`);
    await new Promise((r) => setTimeout(r, 50));

    const userMsg: ChatMessage = {
      id: "user-1",
      role: "user",
      parts: [{ type: "text", text: "Hello" }]
    };

    // Step 1: Send the first message — server will respond with an assistant message
    const done1 = await sendChatRequest(ws, "req-1", [userMsg]);
    expect(done1).toBe(true);

    // Verify we have user + assistant in the DB
    const afterFirst = await fetchPersistedMessages(room);
    expect(afterFirst.length).toBeGreaterThanOrEqual(2);
    const assistantMsg = afterFirst.find((m) => m.role === "assistant");
    expect(assistantMsg).toBeDefined();

    // Step 2: Simulate regenerate() — client removes the assistant message
    // and sends only the user message (truncated array).
    const done2 = await sendChatRequest(ws, "req-2", [userMsg]);
    expect(done2).toBe(true);

    // Step 3: Verify the old assistant message was deleted and replaced
    const afterRegenerate = await fetchPersistedMessages(room);
    const userMessages = afterRegenerate.filter((m) => m.role === "user");
    const assistantMessages = afterRegenerate.filter(
      (m) => m.role === "assistant"
    );

    // Should have exactly 1 user message and 1 new assistant message
    expect(userMessages.length).toBe(1);
    expect(userMessages[0].id).toBe("user-1");
    expect(assistantMessages.length).toBe(1);

    // The new assistant message should NOT be the same as the old one
    // (it was regenerated by the server)
    expect(assistantMessages[0].id).not.toBe(assistantMsg!.id);

    ws.close(1000);
  });

  it("does not delete messages when using CF_AGENT_CHAT_MESSAGES (no reconcile)", async () => {
    const room = crypto.randomUUID();
    const { ws } = await connectChatWS(`/agents/test-chat-agent/${room}`);
    await new Promise((r) => setTimeout(r, 50));

    const agentStub = await getAgentByName(env.TestChatAgent, room);

    // Seed the DB with two messages directly
    const userMsg: ChatMessage = {
      id: "user-seed",
      role: "user",
      parts: [{ type: "text", text: "Hello" }]
    };
    const assistantMsg: ChatMessage = {
      id: "assistant-seed",
      role: "assistant",
      parts: [{ type: "text", text: "Hi there!" }]
    };
    await agentStub.persistMessages([userMsg, assistantMsg]);

    // Send a CF_AGENT_CHAT_MESSAGES with only a new message (no reconcile)
    const newMsg: ChatMessage = {
      id: "new-msg",
      role: "user",
      parts: [{ type: "text", text: "Follow up" }]
    };

    ws.send(
      JSON.stringify({
        type: MessageType.CF_AGENT_CHAT_MESSAGES,
        messages: [newMsg]
      })
    );

    await new Promise((r) => setTimeout(r, 100));

    // All 3 messages should be preserved — CF_AGENT_CHAT_MESSAGES does not reconcile
    const persisted = await fetchPersistedMessages(room);
    expect(persisted.length).toBe(3);
    expect(persisted.map((m) => m.id)).toContain("user-seed");
    expect(persisted.map((m) => m.id)).toContain("assistant-seed");
    expect(persisted.map((m) => m.id)).toContain("new-msg");

    ws.close(1000);
  });

  it("reconcile deletes stale rows when subset is persisted", async () => {
    const room = crypto.randomUUID();
    const { ws } = await connectChatWS(`/agents/test-chat-agent/${room}`);
    await new Promise((r) => setTimeout(r, 50));

    const agentStub = await getAgentByName(env.TestChatAgent, room);

    // Seed the DB with 3 messages directly
    const msg1: ChatMessage = {
      id: "m1",
      role: "user",
      parts: [{ type: "text", text: "Hello" }]
    };
    const msg2: ChatMessage = {
      id: "m2",
      role: "assistant",
      parts: [{ type: "text", text: "Hi!" }]
    };
    const msg3: ChatMessage = {
      id: "m3",
      role: "user",
      parts: [{ type: "text", text: "How are you?" }]
    };
    await agentStub.persistMessages([msg1, msg2, msg3]);

    let persisted = (await agentStub.getPersistedMessages()) as ChatMessage[];
    expect(persisted.length).toBe(3);

    // Persist with reconcile, omitting msg2 (simulates regenerate removing assistant)
    await agentStub.persistMessages([msg1, msg3], [], {
      _deleteStaleRows: true
    });

    // msg2 should be deleted
    persisted = (await agentStub.getPersistedMessages()) as ChatMessage[];
    expect(persisted.length).toBe(2);
    expect(persisted.map((m) => m.id)).toEqual(["m1", "m3"]);

    ws.close(1000);
  });

  it("reconcile preserves server messages when incoming set contains new IDs", async () => {
    const room = crypto.randomUUID();
    const { ws } = await connectChatWS(`/agents/test-chat-agent/${room}`);
    await new Promise((r) => setTimeout(r, 50));

    const agentStub = await getAgentByName(env.TestChatAgent, room);

    // Seed the DB with a user message and an assistant response
    const user1: ChatMessage = {
      id: "u1",
      role: "user",
      parts: [{ type: "text", text: "First" }]
    };
    const asst1: ChatMessage = {
      id: "a1",
      role: "assistant",
      parts: [{ type: "text", text: "Reply to first" }]
    };
    await agentStub.persistMessages([user1, asst1]);

    // Simulate a client that reconnects with partial history and appends
    // a new user message (does not include the assistant message it never saw)
    const user2: ChatMessage = {
      id: "u2",
      role: "user",
      parts: [{ type: "text", text: "Second" }]
    };
    await agentStub.persistMessages([user1, user2], [], {
      _deleteStaleRows: true
    });

    // The assistant message should be preserved — the client sent a new
    // message ID ("u2") not in the server state, so stale deletion is skipped
    const persisted = (await agentStub.getPersistedMessages()) as ChatMessage[];
    expect(persisted.length).toBe(3);
    expect(persisted.map((m) => m.id)).toEqual(["u1", "a1", "u2"]);

    ws.close(1000);
  });

  it("reconcile is a no-op when incoming set matches DB", async () => {
    const room = crypto.randomUUID();
    const { ws } = await connectChatWS(`/agents/test-chat-agent/${room}`);
    await new Promise((r) => setTimeout(r, 50));

    const agentStub = await getAgentByName(env.TestChatAgent, room);

    const msg1: ChatMessage = {
      id: "keep-1",
      role: "user",
      parts: [{ type: "text", text: "Hello" }]
    };
    const msg2: ChatMessage = {
      id: "keep-2",
      role: "assistant",
      parts: [{ type: "text", text: "Hi!" }]
    };
    await agentStub.persistMessages([msg1, msg2]);

    // Persist the same set with reconcile — nothing should change
    await agentStub.persistMessages([msg1, msg2], [], {
      _deleteStaleRows: true
    });

    const persisted = (await agentStub.getPersistedMessages()) as ChatMessage[];
    expect(persisted.length).toBe(2);
    expect(persisted.map((m) => m.id)).toEqual(["keep-1", "keep-2"]);

    ws.close(1000);
  });

  it("trigger field is stripped from options.body (not leaked to user code)", async () => {
    const room = crypto.randomUUID();
    const { ws } = await connectChatWS(`/agents/test-chat-agent/${room}`);

    const agentStub = await getAgentByName(env.TestChatAgent, room);
    await agentStub.clearCapturedContext();

    const userMsg: ChatMessage = {
      id: "msg-trigger",
      role: "user",
      parts: [{ type: "text", text: "Hello" }]
    };

    // Send with trigger field in the body (as the transport now does)
    // but no other custom fields
    const done = await sendChatRequest(ws, "req-trigger", [userMsg], {
      trigger: "regenerate-message"
    });
    expect(done).toBe(true);

    await new Promise((r) => setTimeout(r, 100));

    // trigger is destructured separately on the server (like messages and
    // clientTools), so options.body should be undefined when no other
    // custom fields are present.
    const capturedBody = await agentStub.getCapturedBody();
    expect(capturedBody).toBeUndefined();

    ws.close(1000);
  });

  it("trigger field does not pollute options.body when custom fields are present", async () => {
    const room = crypto.randomUUID();
    const { ws } = await connectChatWS(`/agents/test-chat-agent/${room}`);

    const agentStub = await getAgentByName(env.TestChatAgent, room);
    await agentStub.clearCapturedContext();

    const userMsg: ChatMessage = {
      id: "msg-trigger-2",
      role: "user",
      parts: [{ type: "text", text: "Hello" }]
    };

    // Send with trigger AND a real custom field
    const done = await sendChatRequest(ws, "req-trigger-2", [userMsg], {
      trigger: "regenerate-message",
      model: "claude-sonnet"
    });
    expect(done).toBe(true);

    await new Promise((r) => setTimeout(r, 100));

    // options.body should contain only the custom field, not trigger
    const capturedBody = await agentStub.getCapturedBody();
    expect(capturedBody).toEqual({ model: "claude-sonnet" });
    expect(capturedBody).not.toHaveProperty("trigger");

    ws.close(1000);
  });
});