import type { UIMessageChunk } from "ai";

/** Canonical parts are projected independently of an attached browser response. */
export function createAnswerStreamProjection(options: {
  messageId: string;
  sessionId: string;
  write: (chunk: UIMessageChunk) => void;
  onChange: () => void;
}) {
  const parts: Record<string, unknown>[] = [];
  let segment = 0;
  let textId = "answer";
  let currentText: Record<string, unknown> | undefined;
  let finished = false;
  const endText = () => {
    if (!currentText) return;
    currentText.state = "done";
    options.write({ type: "text-end", id: textId });
    currentText = undefined;
    textId = `answer-${++segment}`;
  };
  const appendText = (delta: string) => {
    if (!delta || finished) return;
    if (!currentText) {
      currentText = { type: "text", text: "", state: "streaming" };
      parts.push(currentText);
      options.write({ type: "text-start", id: textId });
    }
    currentText.text = `${String(currentText.text)}${delta}`;
    options.write({ type: "text-delta", id: textId, delta });
    options.write({ type: "data-stream_token", transient: true, data: {
      session_id: options.sessionId, message_id: options.messageId, content: delta,
    } } as UIMessageChunk);
    options.onChange();
  };
  const appendData = (event: string, data: Record<string, unknown>) => {
    if (finished) return;
    if (event === "answer_ui") endText();
    const type = `data-${event}`;
    const id = [data.id, data.request_id, data.envelope_id, data.block_id].find((value): value is string => typeof value === "string");
    const existing = id === undefined ? -1 : parts.findIndex((part) => part.type === type && part.id === id);
    const part = id ? { type, id, data } : { type, data };
    if (existing >= 0) parts[existing] = part;
    else parts.push(part);
    options.write(part as UIMessageChunk);
    options.onChange();
  };
  const finish = (finishReason: "stop" | "error" = "stop") => {
    if (finished) return;
    if (!parts.length) appendText("Agent 已完成处理。");
    endText();
    finished = true;
    options.onChange();
    options.write({ type: "finish", finishReason });
  };
  return { parts, appendText, appendData, finish, get finished() { return finished; } };
}