import assert from "node:assert/strict";
import test from "node:test";
import {
  appendStreamingDelta,
  asStreamingAssistantMetadata,
  clearStreamingAssistantTrace,
  createStreamingAssistantMetadata,
  createStreamingAttemptState,
  finalizeStreamingAssistant,
  promoteStreamingAssistantTrace,
  setStreamingAssistantSegment,
  transitionStreamingToolCall,
} from "../src/lib/chat-stream-presentation.ts";

function assistantMessage(content = "") {
  return {
    id: "temp-assistant-test",
    createdAt: "2026-08-30T00:00:00.000Z",
    role: "assistant",
    content,
    name: null,
    metadata: createStreamingAssistantMetadata(),
  };
}

function metadata(messages) {
  return asStreamingAssistantMetadata(messages.at(-1)?.metadata);
}

test("a tool boundary promotes the current segment and the next delta starts the answer in its place", () => {
  let messages = [assistantMessage()];
  let attempt = appendStreamingDelta(createStreamingAttemptState(), "I'll search for that.");
  messages = setStreamingAssistantSegment(messages, "temp-assistant-", attempt.segmentText);

  const toolTransition = transitionStreamingToolCall(attempt, {
    toolCallId: "search-1",
    status: "initiated",
  });
  attempt = toolTransition.attempt;
  messages = promoteStreamingAssistantTrace(messages, "temp-assistant-", toolTransition.traceText ?? "");

  assert.equal(messages.at(-1).content, "");
  assert.deepEqual(metadata(messages), {
    kind: "streaming_assistant",
    ephemeralTrace: "I'll search for that.",
    traceRevision: 1,
  });

  attempt = appendStreamingDelta(attempt, "Here is the answer");
  messages = setStreamingAssistantSegment(messages, "temp-assistant-", attempt.segmentText);

  assert.equal(messages.at(-1).content, "Here is the answer");
  assert.equal(metadata(messages)?.ephemeralTrace, null);

  messages = finalizeStreamingAssistant(messages, "temp-assistant-", "Here is the answer.");
  assert.equal(messages.at(-1).content, "Here is the answer.");
  assert.equal(metadata(messages)?.ephemeralTrace, null);
});

test("a later tool round replaces the previous trace and parallel initiated calls do not erase it", () => {
  let messages = [assistantMessage("First trace")];
  let attempt = appendStreamingDelta(createStreamingAttemptState(), "First trace");
  let transition = transitionStreamingToolCall(attempt, { toolCallId: "tool-1", status: "initiated" });
  attempt = transition.attempt;
  messages = promoteStreamingAssistantTrace(messages, "temp-assistant-", transition.traceText ?? "");

  attempt = appendStreamingDelta(attempt, "I need one more lookup.");
  messages = setStreamingAssistantSegment(messages, "temp-assistant-", attempt.segmentText);
  transition = transitionStreamingToolCall(attempt, { toolCallId: "tool-2", status: "initiated" });
  attempt = transition.attempt;
  messages = promoteStreamingAssistantTrace(messages, "temp-assistant-", transition.traceText ?? "");

  assert.equal(metadata(messages)?.ephemeralTrace, "I need one more lookup.");
  assert.equal(metadata(messages)?.traceRevision, 2);

  transition = transitionStreamingToolCall(attempt, { toolCallId: "tool-3", status: "initiated" });
  assert.equal(transition.startedNewToolCall, true);
  assert.equal(transition.traceText, null);
  messages = promoteStreamingAssistantTrace(messages, "temp-assistant-", transition.traceText ?? "");

  assert.equal(metadata(messages)?.ephemeralTrace, "I need one more lookup.");
  assert.equal(metadata(messages)?.traceRevision, 2);
});

test("terminal and duplicate tool events never promote or reset the in-progress segment", () => {
  const initial = appendStreamingDelta(createStreamingAttemptState(), "candidate");
  const completed = transitionStreamingToolCall(initial, { toolCallId: "tool-1", status: "completed" });

  assert.equal(completed.startedNewToolCall, false);
  assert.strictEqual(completed.attempt, initial);

  const initiated = transitionStreamingToolCall(initial, { toolCallId: "tool-1", status: "initiated" });
  const withNextSegment = appendStreamingDelta(initiated.attempt, "next segment");
  const duplicate = transitionStreamingToolCall(withNextSegment, { toolCallId: "tool-1", status: "initiated" });

  assert.equal(duplicate.startedNewToolCall, false);
  assert.strictEqual(duplicate.attempt, withNextSegment);
  assert.equal(duplicate.attempt.segmentText, "next segment");
});

test("a fresh per-attempt accumulator can replay and reclassify an already-seen tool round", () => {
  let messages = [assistantMessage("stale partial answer")];
  let replayAttempt = appendStreamingDelta(createStreamingAttemptState(), "Replayed trace");
  messages = setStreamingAssistantSegment(messages, "temp-assistant-", replayAttempt.segmentText);

  const replayedTool = transitionStreamingToolCall(replayAttempt, { toolCallId: "same-tool-id", status: "initiated" });
  messages = promoteStreamingAssistantTrace(messages, "temp-assistant-", replayedTool.traceText ?? "");

  assert.equal(replayedTool.startedNewToolCall, true);
  assert.equal(messages.at(-1).content, "");
  assert.equal(metadata(messages)?.ephemeralTrace, "Replayed trace");
});

test("message helpers preserve identity when no matching temporary assistant exists", () => {
  const messages = [
    {
      id: "persisted-assistant",
      createdAt: "2026-08-30T00:00:00.000Z",
      role: "assistant",
      content: "Persisted",
      name: null,
      metadata: null,
    },
  ];

  assert.strictEqual(setStreamingAssistantSegment(messages, "temp-assistant-", "new"), messages);
  assert.strictEqual(promoteStreamingAssistantTrace(messages, "temp-assistant-", "trace"), messages);
  assert.strictEqual(finalizeStreamingAssistant(messages, "temp-assistant-", "done"), messages);
  assert.strictEqual(clearStreamingAssistantTrace(messages, "temp-assistant-"), messages);
});

test("a terminal error clears the ephemeral trace without discarding a partial answer", () => {
  let messages = [
    {
      ...assistantMessage("Partial answer"),
      metadata: {
        ...createStreamingAssistantMetadata(),
        ephemeralTrace: "Still checking one source.",
        traceRevision: 1,
      },
    },
  ];

  messages = clearStreamingAssistantTrace(messages, "temp-assistant-");

  assert.equal(messages.at(-1).content, "Partial answer");
  assert.equal(metadata(messages)?.ephemeralTrace, null);
});