import { describe, expect, it } from "bun:test";
import { streamPromptResponse } from "../../src/routes/chatStream.js";
import type { OpencodeRuntimeAdapter } from "../../src/runtime/opencode.js";

type Output = { event: string; data: Record<string, unknown> };
const textStart = { type: "message.part.updated", properties: { sessionID: "s", part: { id: "text", type: "text", text: "" } } };
const textDelta = (delta: string) => ({ type: "message.part.delta", properties: { sessionID: "s", partID: "text", field: "text", delta } });
const textEnd = { type: "message.part.updated", properties: { sessionID: "s", part: { id: "text", type: "text", text: "短说明。", time: { start: 1, end: 2 } } } };
const idle = { type: "session.idle", properties: { sessionID: "s" } };
const tool = (id: string, name: string, status = "completed", input: unknown = {}, output = "{}") => ({
  type: "message.part.updated", properties: { sessionID: "s", part: {
    id, type: "tool", tool: name, callID: id, state: { status, input, output, time: { start: 1, end: status === "completed" ? 2 : undefined } },
  } },
});
const final = tool("final", "final_answer", "completed", {}, JSON.stringify({ answer: "核查完成。", ui_block_ids: [] }));
const uiInput = {
  block_id: "note", phase: "process", fallback_text: "资料说明。",
  spec: { root: "note", elements: { note: { type: "Text", props: { text: "资料说明。", tone: "neutral" }, children: [] } } },
};
const ui = tool("ui", "present_answer", "completed", uiInput);

function createRuntime(events: AsyncIterable<unknown>, extra: Record<string, unknown> = {}) {
  return { prompt: async () => {}, messages: async () => [], abortSession: async () => {}, waitForSessionIdle: async () => {},
    subscribeEvents: async () => events, ...extra } as unknown as OpencodeRuntimeAdapter;
}
function run(runtime: OpencodeRuntimeAdapter, output: Output[], extra: Partial<Parameters<typeof streamPromptResponse>[0]> = {}) {
  return streamPromptResponse({ runtime, sessionId: "s", clientSessionId: "c", message: "核查", write: (event, data) => { output.push({ event, data }); }, ...extra });
}

describe("stream stage boundaries", () => {
  it.each(["text-end", "tool", "permission"])("releases narration before waiting at %s", async (boundary) => {
    const output: Output[] = [];
    const runtime = createRuntime({ async *[Symbol.asyncIterator]() {
      yield textStart;
      yield textDelta("短说明。");
      yield boundary === "text-end" ? textEnd : boundary === "tool" ? tool("query", "mcp_fixture_query", "running") : {
        type: "permission.asked", properties: { id: "permission", sessionID: "s", permission: "shell", patterns: ["probe"], always: [] },
      };
      // The next event is withheld until stage handling has completed.
      expect(output.filter((item) => item.event === "token").map((item) => item.data.content).join("")).toBe("短说明。");
      yield idle;
    } });
    await run(runtime, output);
  });

  it("keeps text before and after an answer component in wire order", async () => {
    const output: Output[] = [];
    const runtime = createRuntime({ async *[Symbol.asyncIterator]() {
      yield textStart; yield textDelta("组件前。"); yield ui;
      yield textDelta("组件后。"); yield textEnd; yield idle;
    } });
    await run(runtime, output);
    expect(output.filter((item) => ["token", "answer_ui"].includes(item.event)).map((item) => item.event === "token" ? item.data.content : "UI"))
      .toEqual(["组件前。", "UI", "组件后。"]);
  });

  it("publishes a plain final answer before native cleanup and never consumes later rewrites", async () => {
    const output: Output[] = [];
    let release!: () => void;
    const cleanup = new Promise<void>((resolve) => { release = resolve; });
    const runtime = createRuntime({ async *[Symbol.asyncIterator]() {
      yield final;
      throw new Error("must not start an automatic rewrite after final_answer");
    } }, { waitForSessionIdle: () => cleanup });
    const running = run(runtime, output);
    void running.catch(() => {});
    try {
      await Bun.sleep(15);
      expect(output.filter((item) => item.event === "final_answer").map((item) => item.data.content)).toEqual(["核查完成。"]);
      expect(output.some((item) => item.event === "done")).toBe(false);
    } finally { release(); }
    expect(await running).toMatchObject({ aborted: false, failed: false });
    expect(output.filter((item) => item.event === "final_answer")).toHaveLength(1);
  });

  it("holds final publication while a parallel tool is still running", async () => {
    const output: Output[] = [];
    const runtime = createRuntime({ async *[Symbol.asyncIterator]() {
      yield tool("query", "mcp_fixture_query", "running"); yield final;
      expect(output.some((item) => item.event === "final_answer")).toBe(false);
      yield tool("query", "mcp_fixture_query");
      throw new Error("must close once all tools finish");
    } });
    await run(runtime, output);
    expect(output.filter((item) => item.event === "final_answer")).toHaveLength(1);
  });

  it("waits for a pending answer component before resolving final references", async () => {
    const output: Output[] = [];
    const component = tool("ui", "present_answer", "completed", { ...uiInput, phase: "final" });
    const answer = tool("final", "final_answer", "completed", {}, JSON.stringify({ answer: "结论\n[[ui:note]]", ui_block_ids: ["note"] }));
    const runtime = createRuntime({ async *[Symbol.asyncIterator]() {
      yield tool("ui", "present_answer", "running"); yield answer;
      expect(output.some((item) => item.event === "final_answer")).toBe(false);
      yield component;
      throw new Error("must close after the pending component is published");
    } });
    await run(runtime, output);
    expect(output.filter((item) => item.event === "answer_ui")).toHaveLength(1);
    expect(output.find((item) => item.event === "final_answer")?.data.ui_block_ids).toEqual(["note"]);
  });

  it("does not interrupt an older answer when additional user input needs processing", async () => {
    const output: Output[] = [];
    let interrupted = false;
    const runtime = createRuntime({ async *[Symbol.asyncIterator]() {
      yield final;
      expect(output.some((item) => item.event === "final_answer")).toBe(false);
      expect(interrupted).toBe(false);
      yield idle;
    } }, { abortSession: async () => { interrupted = true; } });
    await run(runtime, output, { beforeAnswerClose: async () => false });
    expect(interrupted).toBe(false);
    expect(output.filter((item) => item.event === "final_answer")).toHaveLength(1);
  });

  it("drains received text before awaiting cancellation without inventing a final answer", async () => {
    const output: Output[] = [];
    let receivedAtAbort = "";
    const controller = new AbortController();
    const runtime = createRuntime({ async *[Symbol.asyncIterator]() {
      yield textStart; yield textDelta("已经收到。"); controller.abort(); yield idle;
    } }, { abortSession: async () => {
      receivedAtAbort = output.filter((item) => item.event === "token").map((item) => item.data.content).join("");
    } });
    expect(await run(runtime, output, { signal: controller.signal })).toMatchObject({ aborted: true });
    expect(receivedAtAbort).toBe("已经收到。");
    expect(output.some((item) => item.event === "final_answer" || item.event === "done")).toBe(false);
  });

  it("interrupts a paced tool transition before publishing the tool event", async () => {
    const output: Output[] = [];
    const controller = new AbortController();
    const text = "管道压力核查。".repeat(15);
    const runtime = createRuntime({ async *[Symbol.asyncIterator]() {
      yield textStart; yield textDelta(text);
      setTimeout(() => controller.abort(), 1);
      yield tool("query", "mcp_fixture_query", "running");
      throw new Error("must stop reading native events after cancellation");
    } });
    expect(await run(runtime, output, { signal: controller.signal })).toMatchObject({ aborted: true });
    expect(output.filter((item) => item.event === "token").map((item) => item.data.content).join("")).toBe(text);
    expect(output.some((item) => item.event === "progress" && item.data.id === "query")).toBe(false);
  });
});