import type { ChatTimelineMessageItem } from '@easyflow/common-ui';

import { describe, expect, it } from 'vitest';

import {
  appendWorkflowStreamDelta,
  appendWorkflowThinkingDelta,
  createWorkflowStreamMessage,
  updateWorkflowStreamStatus,
} from './workflowChatStreamMessage';

function createMessage(): ChatTimelineMessageItem {
  return appendWorkflowStreamDelta(
    createWorkflowStreamMessage('llm-stream-1', '大模型'),
    '首包',
  );
}

describe('workflowChatStreamMessage', () => {
  it('appends later chunks with a new message and text-part reference', () => {
    const message = createMessage();
    const nextMessage = appendWorkflowStreamDelta(message, '后续内容');
    const answerPart = nextMessage.parts.find(
      (part) => part.id === 'llm-stream-1-answer',
    );

    expect(nextMessage).not.toBe(message);
    expect(answerPart).not.toBe(message.parts[0]);
    expect(answerPart?.content).toBe('**大模型**\n\n首包后续内容');
    expect(message.parts[0]?.content).toBe('**大模型**\n\n首包');
  });

  it('streams thinking separately and ends it when answer starts', () => {
    const message = createWorkflowStreamMessage('llm-stream-1', '大模型');
    const thinkingMessage = appendWorkflowThinkingDelta(message, '先分析');
    const nextThinkingMessage = appendWorkflowThinkingDelta(
      thinkingMessage,
      '再作答',
    );
    const answerMessage = appendWorkflowStreamDelta(
      nextThinkingMessage,
      '结论',
    );

    expect(nextThinkingMessage.parts).toEqual([
      {
        content: '先分析再作答',
        id: 'llm-stream-1-thinking',
        status: 'thinking',
        type: 'thinking',
      },
      {
        content: '**大模型**',
        id: 'llm-stream-1-answer',
        type: 'text',
      },
    ]);
    expect(answerMessage.parts).toEqual([
      {
        content: '先分析再作答',
        id: 'llm-stream-1-thinking',
        status: 'end',
        type: 'thinking',
      },
      {
        content: '**大模型**\n\n结论',
        id: 'llm-stream-1-answer',
        type: 'text',
      },
    ]);
  });

  it('ignores late thinking after answer output has started', () => {
    const message = createMessage();

    expect(appendWorkflowThinkingDelta(message, '迟到内容')).toBe(message);
  });

  it('completes thinking with a new message reference', () => {
    const message = appendWorkflowThinkingDelta(
      createWorkflowStreamMessage('llm-stream-1', '大模型'),
      '思考内容',
    );
    const nextMessage = updateWorkflowStreamStatus(message, 'done');

    expect(nextMessage).not.toBe(message);
    expect(nextMessage.status).toBe('done');
    expect(nextMessage.parts[0]).toMatchObject({
      status: 'end',
      type: 'thinking',
    });
    expect(message.status).toBe('streaming');
  });

  it('marks thinking as error when the stream fails', () => {
    const message = appendWorkflowThinkingDelta(
      createWorkflowStreamMessage('llm-stream-1', '大模型'),
      '思考内容',
    );

    const nextMessage = updateWorkflowStreamStatus(message, 'error');

    expect(nextMessage.parts[0]).toMatchObject({
      status: 'error',
      type: 'thinking',
    });
  });
});
