import { describe, expect, it } from "bun:test";
import type { UIMessageChunk } from "ai";
import { createAnswerStreamProjection } from "../../src/chat/answer-stream-projection.js";

describe("canonical answer stream projection", () => {
  it("splits text around UI blocks, updates them in place and finishes each segment exactly once", () => {
    const chunks: UIMessageChunk[] = [];
    const stream = createAnswerStreamProjection({ messageId: "a", sessionId: "s", write: (chunk) => chunks.push(chunk), onChange() {} });
    stream.appendText("查询资料。");
    stream.appendData("answer_ui", { block_id: "chart", fallback_text: "初始" });
    stream.appendData("answer_ui", { block_id: "chart", fallback_text: "更新" });
    stream.appendText("待现场核验。");
    stream.finish(); stream.finish(); stream.appendText("不应追加");
    expect(stream.parts.map((part) => part.type)).toEqual(["text", "data-answer_ui", "text"]);
    expect(stream.parts[1].data).toEqual({ block_id: "chart", fallback_text: "更新" });
    expect(stream.parts[0].text).toBe("查询资料。");
    expect(stream.parts[2].text).toBe("待现场核验。");
    expect(chunks.filter((chunk) => chunk.type === "text-start").map((chunk) => chunk.id)).toEqual(["answer", "answer-1"]);
    expect(chunks.filter((chunk) => chunk.type === "text-end").map((chunk) => chunk.id)).toEqual(["answer", "answer-1"]);
    expect(chunks.filter((chunk) => chunk.type === "finish")).toHaveLength(1);
  });
  it("continues canonical parts for persistence after the browser disconnects", () => {
    let connected = true;
    const written: UIMessageChunk[] = [];
    const snapshots: unknown[] = [];
    const stream = createAnswerStreamProjection({ messageId: "a", sessionId: "s", write: (chunk) => { if (connected) written.push(chunk); }, onChange: () => snapshots.push(structuredClone(stream.parts)) });
    stream.appendText("第一段");
    connected = false;
    const count = written.length;
    stream.appendData("answer_ui", { block_id: "evidence", fallback_text: "依据" });
    stream.appendText("第二段"); stream.finish();
    expect(written).toHaveLength(count);
    expect(snapshots.at(-1)).toEqual([
      { type: "text", text: "第一段", state: "done" },
      { type: "data-answer_ui", id: "evidence", data: { block_id: "evidence", fallback_text: "依据" } },
      { type: "text", text: "第二段", state: "done" },
    ]);
  });
});