/**
 * Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
 * SPDX-License-Identifier: MIT
 */

import { describe, expect, it } from 'vitest';
import { IContainer, IEngine, WorkflowStatus } from '@flowgram.ai/runtime-interface';

import { WorkflowRuntimeContainer } from '../../container';
import { TestSchemas } from '.';

const container: IContainer = WorkflowRuntimeContainer.instance;

describe('WorkflowRuntime branch schema', () => {
  it('should execute a workflow with branch 1', async () => {
    const engine = container.get<IEngine>(IEngine);
    const { context, processing } = engine.invoke({
      schema: TestSchemas.branchTwoLayersSchema,
      inputs: {
        model_id: 1,
        prompt: 'Tell me a joke',
      },
    });
    expect(context.statusCenter.workflow.status).toBe(WorkflowStatus.Processing);
    const result = await processing;
    expect(context.statusCenter.workflow.status).toBe(WorkflowStatus.Succeeded);
    expect(result).toStrictEqual({
      m3_res:
        'Hi, I am an AI model, my name is AI_MODEL_3, temperature is 0.5, system prompt is "I\'m Model 3", prompt is "Hi, I am an AI model, my name is AI_MODEL_1, temperature is 0.5, system prompt is "I\'m Model 1", prompt is "Tell me a joke""',
    });
    const report = context.reporter.export();
    expect(report.workflowStatus.status).toBe(WorkflowStatus.Succeeded);
    expect(report.reports.start_0.status).toBe(WorkflowStatus.Succeeded);
    expect(report.reports.condition_0.status).toBe(WorkflowStatus.Succeeded);
    expect(report.reports.llm_1.status).toBe(WorkflowStatus.Succeeded);
    expect(report.reports.llm_3.status).toBe(WorkflowStatus.Succeeded);
    expect(report.reports.end_0.status).toBe(WorkflowStatus.Succeeded);
    expect(report.reports.llm_2).toBeUndefined();
    expect(report.reports.llm_4).toBeUndefined();
  });

  it('should execute a workflow with branch 2', async () => {
    const engine = container.get<IEngine>(IEngine);
    const { context, processing } = engine.invoke({
      schema: TestSchemas.branchTwoLayersSchema,
      inputs: {
        model_id: 2,
        prompt: 'Tell me a story',
      },
    });
    expect(context.statusCenter.workflow.status).toBe(WorkflowStatus.Processing);
    const result = await processing;
    expect(context.statusCenter.workflow.status).toBe(WorkflowStatus.Succeeded);
    expect(result).toStrictEqual({
      m4_res:
        'Hi, I am an AI model, my name is AI_MODEL_4, temperature is 0.5, system prompt is "I\'m Model 4", prompt is "Hi, I am an AI model, my name is AI_MODEL_2, temperature is 0.6, system prompt is "I\'m Model 2", prompt is "Tell me a story""',
    });
    const report = context.reporter.export();
    expect(report.workflowStatus.status).toBe(WorkflowStatus.Succeeded);
    expect(report.reports.start_0.status).toBe(WorkflowStatus.Succeeded);
    expect(report.reports.condition_0.status).toBe(WorkflowStatus.Succeeded);
    expect(report.reports.llm_2.status).toBe(WorkflowStatus.Succeeded);
    expect(report.reports.llm_4.status).toBe(WorkflowStatus.Succeeded);
    expect(report.reports.end_0.status).toBe(WorkflowStatus.Succeeded);
    expect(report.reports.llm_1).toBeUndefined();
    expect(report.reports.llm_3).toBeUndefined();
  });
});
