import { expect, it } from "bun:test";
import type { UIMessageChunk } from "ai";
import { createChatRunAccumulator, publishRunEvent } from "../../src/routes/chat-run-accumulator.js";
import { createInitialStreamingMessages, type ActiveRun } from "../../src/routes/chatUiState.js";

it.each([false, true])("projects one event stream into reconnect and persisted messages (collaboration=%s)", (collaboration) => {
  const run: ActiveRun = {
    clientSessionId: "s", controller: new AbortController(), status: "running",
    messages: createInitialStreamingMessages([], "查询"), pendingPermissions: new Map(), pendingQuestions: new Map(), subscribers: new Set(),
  };
  const chunks: UIMessageChunk[] = [];
  const stream = createChatRunAccumulator({ run, messages: [{ id: "user", role: "user", parts: [{ type: "text", text: "查询" }] }],
    messageId: "answer", sessionId: "s", collaboration, write: (chunk) => chunks.push(chunk), onStatus() {} });
  stream.apply("progress", { id: "query", status: "running", title: "查询资料" });
  stream.apply("token", { content: "已有数据" });
  stream.apply("permission_request", { request_id: "permission", session_id: "s", permission: "shell", patterns: [], always: [], created_at: 1 });
  stream.apply("permission_response", { request_id: "permission", reply: "once" });
  stream.appendInput({ id: "extra", role: "user", parts: [{ type: "text", text: "补充" }] }, "补充");
  stream.appendInput({ id: "extra", role: "user", parts: [] }, "重复");
  stream.apply("final_answer", { content: "结论" });
  stream.apply("done", {});
  expect(run.status).toBe("completed");
  expect(run.pendingPermissions.size).toBe(0);
  expect(run.messages.at(-1)).toMatchObject({ id: "answer", content: "结论", permissions: [{ status: "approved_once" }], progress: [{ status: "completed" }, { status: "completed" }] });
  const saved = stream.uiMessages();
  expect(saved.map((message) => message.id)).toEqual(["user", "extra", "answer"]);
  expect(saved.at(-1)?.parts).toEqual((run.messages.at(-1) as { parts: typeof saved[number]["parts"] }).parts);
  expect(saved.at(-1)?.parts.some((part) => part.type === "data-progress")).toBe(!collaboration);
  expect(chunks.filter((chunk) => chunk.type === "finish")).toHaveLength(1);
});

it.each(["done", "error", "aborted"])("publishes the same %s terminal state to reconnect subscribers and storage", (event) => {
  const run: ActiveRun = {
    clientSessionId: "s", controller: new AbortController(), status: "running",
    messages: createInitialStreamingMessages([], "查询"), pendingPermissions: new Map(), pendingQuestions: new Map(), subscribers: new Set(),
  };
  const received: Array<{ event: string; data: Record<string, unknown> }> = [];
  run.subscribers.add({ write: (event, data) => received.push({ event, data }), close() {} });
  const stream = createChatRunAccumulator({ run, messages: [], messageId: "answer", sessionId: "s",
    collaboration: false, write() {}, onStatus() {} });
  publishRunEvent(run, "progress", { id: "tool", title: "正在处理", status: "running", phase: "tool" });
  publishRunEvent(run, "token", { content: "部分结果" });
  publishRunEvent(run, event, { session_id: "s", message: "执行已停止" });
  const canonical = stream.uiMessages().at(-1)?.parts.find((part) => part.type === "data-run_status");
  expect(received.at(-2)).toMatchObject({ event: "run_status", data: { status: event === "done" ? "completed" : event } });
  expect(canonical).toMatchObject({ data: { status: received.at(-2)?.data.status } });
});