# -*- coding: utf-8 -*-
"""Test for the workflow module"""
import asyncio
import os
import shutil
import threading
import time
import unittest
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Dict, Optional
from unittest import mock
from unittest.mock import AsyncMock, MagicMock

import openai
import ray
from parameterized import parameterized, parameterized_class
from torch import Tensor

from tests.tools import (
    CHAT_TEMPLATE,
    RayUnittestBaseAsync,
    get_model_path,
    get_template_config,
    get_unittest_dataset_config,
)
from trinity.common.config import InferenceModelConfig
from trinity.common.constants import LOG_DIR_ENV_VAR, LOG_LEVEL_ENV_VAR
from trinity.common.experience import EID, Experience
from trinity.common.models.allocator import Allocator
from trinity.common.models.model import ModelWrapper
from trinity.common.workflows import WORKFLOWS, Workflow
from trinity.common.workflows.customized_math_workflows import MathBoxedWorkflow
from trinity.common.workflows.eval_workflow import MathEvalWorkflow
from trinity.common.workflows.workflow import MathWorkflow, MultiTurnWorkflow, Task
from trinity.explorer.workflow_runner import WorkflowRunner


def deserialize_experiences(exp_payload: bytes) -> list[Experience]:
    if not exp_payload:
        return []
    return Experience.deserialize_many(exp_payload)


def patch_runner_models(*wrappers):
    return mock.patch(
        "trinity.explorer.workflow_runner.Allocator.get_model",
        side_effect=list(wrappers),
    )


@dataclass
class MockResponse:
    response_text: str
    reward: float = 0.0
    metrics: Optional[Dict[str, float]] = None
    info: Optional[Dict] = None
    unique_id: Optional[str] = "0"
    tokens: Optional[Tensor] = Tensor([0, 0])
    prompt_length: int = 1
    eid: EID = field(default_factory=EID)


class DummyWorkflow(Workflow):
    can_reset: bool = True
    can_repeat: bool = True

    def __init__(self, model, task: Task, auxiliary_models=None):
        super().__init__(task=task, model=model, auxiliary_models=auxiliary_models)
        self.obj = task.raw_task
        self.output_format = task.workflow_args["output_format"]
        self.repeat_times = task.rollout_args.n
        # Check self.auxiliary_models (OpenAI clients derived from ModelWrapper)
        if self.auxiliary_models is not None:
            for m in self.auxiliary_models:
                assert isinstance(m, openai.OpenAI)

    def reset(self, task: Task):
        self.obj = task.raw_task
        self.output_format = task.workflow_args["output_format"]

    def set_repeat_times(self, repeat_times, run_id_base):
        self.repeat_times = repeat_times
        self.run_id_base = run_id_base

    def run(self):
        exps = []
        if self.output_format == "json":
            import json

            for i in range(self.repeat_times):
                exp = Experience(tokens=Tensor([0, 1, 2, 3]), prompt_length=1)
                exp.response_text = json.dumps(self.obj)
                exps.append(exp)
            return exps
        elif self.output_format == "yaml":
            import yaml

            for i in range(self.repeat_times):
                exp = Experience(tokens=Tensor([0, 1, 2, 3]), prompt_length=1)
                exp.response_text = yaml.safe_dump(self.obj)
                exps.append(exp)
            return exps
        else:
            raise ValueError("Invalid output format")


class DummyAsyncWorkflow(Workflow):
    can_reset: bool = True
    can_repeat: bool = True
    is_async: bool = True

    def __init__(self, model, task: Task, auxiliary_models=None):
        super().__init__(task=task, model=model, auxiliary_models=auxiliary_models)
        self.obj = task.raw_task
        self.output_format = task.workflow_args["output_format"]
        self.repeat_times = task.rollout_args.n
        # Check self.auxiliary_models (AsyncOpenAI clients derived from ModelWrapper)
        if self.auxiliary_models is not None:
            for m in self.auxiliary_models:
                assert isinstance(m, openai.AsyncOpenAI)

    def reset(self, task: Task):
        self.obj = task.raw_task
        self.output_format = task.workflow_args["output_format"]

    def set_repeat_times(self, repeat_times, run_id_base):
        self.repeat_times = repeat_times
        self.run_id_base = run_id_base

    async def run_async(self):
        await asyncio.sleep(0.1)
        exps = []
        if self.output_format == "json":
            import json

            for i in range(self.repeat_times):
                exp = Experience(tokens=Tensor([0, 1, 2, 3]), prompt_length=1)
                exp.response_text = json.dumps(self.obj)
                exps.append(exp)
            return exps
        elif self.output_format == "yaml":
            import yaml

            for i in range(self.repeat_times):
                exp = Experience(tokens=Tensor([0, 1, 2, 3]), prompt_length=1)
                exp.response_text = yaml.safe_dump(self.obj)
                exps.append(exp)
            return exps
        else:
            raise ValueError("Invalid output format")


class DummyMultiTurnWorkflow(MultiTurnWorkflow):
    can_repeat: bool = True

    def __init__(self, model, task: Task, auxiliary_models=None):
        super().__init__(task=task, model=model, auxiliary_models=auxiliary_models)
        self.contents = task.raw_task["contents"]  # type: ignore

    def run(self):
        memory = [{"role": "system", "content": "You are a helpful assistant."}]
        experience_list = []
        for content in self.contents:
            memory.append({"role": "user", "content": content})
            memory.append({"role": "assistant", "content": content.upper()})
            experience = self.process_messages_to_experience(memory, 0, {})
            experience_list.append(experience)
        return experience_list


class DummyAsyncMultiTurnWorkflow(MultiTurnWorkflow):
    is_async: bool = True
    can_repeat: bool = True

    def __init__(self, model, task: Task, auxiliary_models=None):
        super().__init__(task=task, model=model, auxiliary_models=auxiliary_models)
        self.contents = task.raw_task["contents"]  # type: ignore

    async def run_async(self):
        memory = [{"role": "system", "content": "You are a helpful assistant."}]
        experience_list = []
        for content in self.contents:
            await asyncio.sleep(0.1)
            memory.append({"role": "user", "content": content})
            memory.append({"role": "assistant", "content": content.upper()})
            experience = await self.process_messages_to_experience_async(memory, 0, {})
            experience_list.append(experience)
        return experience_list


class WorkflowTest(unittest.TestCase):
    def test_math_workflow(self) -> None:
        model = MagicMock()
        model.chat.return_value = [
            MockResponse(r"\boxed{2}"),
            MockResponse(r"\boxted{3}"),
            MockResponse(r"2"),
            MockResponse("<think>\nThinking\n</think>\n<answer>\n3\n</answer>"),
            MockResponse("<think>\nThinking\n</think>\n<answer>\n\\boxed{2}\n</answer>"),
            MockResponse("<think>Missing closing</think><answer>\\boxed{2}"),
            MockResponse("<answer>\nOnly answer\n</answer>"),
            MockResponse("<think>\nOnly thinking\n</think>"),
            MockResponse("<think>Thinking</think><answer>Answer is not end</answer><answer>1"),
        ]
        taskset_config = get_unittest_dataset_config("countdown")
        task = Task(
            workflow=MathWorkflow,
            repeat_times=taskset_config.repeat_times,
            format_args=taskset_config.format,
            rollout_args=taskset_config.rollout_args,
            is_eval=False,
            raw_task={
                taskset_config.format.prompt_key: "1+1=",
                taskset_config.format.response_key: "2",
            },
        )
        workflow = task.to_workflow(model=model)
        experiences = workflow.run()
        self.assertEqual(len(experiences), 9)
        self.assertEqual(experiences[0].reward, 0.9)
        self.assertEqual(experiences[1].reward, -0.1)
        self.assertEqual(experiences[2].reward, 0.9)
        self.assertEqual(experiences[3].reward, 0.1)
        self.assertEqual(experiences[4].reward, 1.1)
        self.assertEqual(experiences[5].reward, 0.9)
        self.assertEqual(experiences[6].reward, -0.1)
        self.assertEqual(experiences[7].reward, -0.1)
        self.assertEqual(experiences[8].reward, -0.1)

    def test_math_fraction_workflow(self) -> None:
        model = MagicMock()
        model.chat.return_value = [
            MockResponse(r"\boxed{\frac{40}{400}}"),
            MockResponse(r"\boxed{\frac{1}{10}}"),
            MockResponse(r"\boxed{0.1}"),
            MockResponse(r"\boxed{0.1000}"),
            MockResponse(r"\boxed{\frac{1} {10}}"),
            MockResponse(r"The answer is \boxed{\frac{40}{400}}"),
        ]
        taskset_config = get_unittest_dataset_config("countdown")
        task = Task(
            workflow=MathWorkflow,
            repeat_times=taskset_config.repeat_times,
            format_args=taskset_config.format,
            rollout_args=taskset_config.rollout_args,
            is_eval=False,
            raw_task={
                taskset_config.format.prompt_key: r"\frac{40}{400}",
                taskset_config.format.response_key: r"\frac{40}{400}",
            },
        )
        workflow = task.to_workflow(model=model)
        experiences = workflow.run()
        self.assertEqual(len(experiences), 6)
        self.assertEqual(experiences[0].reward, 0.9)
        self.assertEqual(experiences[1].reward, 0.9)
        self.assertEqual(experiences[2].reward, 0.9)
        self.assertEqual(experiences[3].reward, 0.9)
        self.assertEqual(experiences[4].reward, 0.9)
        self.assertEqual(experiences[5].reward, 0.9)

    def test_math_complex_workflow(self) -> None:
        model = MagicMock()
        model.chat.return_value = [
            MockResponse(
                r"$\boxed{\dfrac{108 + 31\sqrt{5}}{216}} \quad \text{and} \quad \boxed{\dfrac{108 - 31\sqrt{5}}{216}}$"
            ),
        ]
        taskset_config = get_unittest_dataset_config("countdown")
        task = Task(
            workflow=MathWorkflow,
            repeat_times=taskset_config.repeat_times,
            format_args=taskset_config.format,
            rollout_args=taskset_config.rollout_args,
            is_eval=False,
            raw_task={
                taskset_config.format.prompt_key: "",
                taskset_config.format.response_key: r"$x_{1}=\frac{1}{2}+\frac{31\sqrt{5}}{216},\quadx_{2}=\frac{1}{2}-\frac{31\sqrt{5}}{216}$",
            },
        )
        workflow = task.to_workflow(model=model)
        experiences = workflow.run()
        self.assertEqual(len(experiences), 1)
        self.assertEqual(experiences[0].reward, 0.9)

    def test_math_boxed_workflow(self) -> None:
        model = MagicMock()
        model.chat.return_value = [
            MockResponse("<think> balabalabala 99 </think>\n \\boxed{36}"),
            MockResponse("answer is \\boxed{36 }"),
            MockResponse("Kim's total points are 6 + 30 =\\boxed{36}"),
            MockResponse("<think> balalaba </think> \\boxed{35.00}"),
        ]
        taskset_config = get_unittest_dataset_config("countdown")
        task = Task(
            workflow=MathBoxedWorkflow,
            repeat_times=taskset_config.repeat_times,
            format_args=taskset_config.format,
            rollout_args=taskset_config.rollout_args,
            workflow_args={
                "with_think": False,
                "format_score_coef": 0.2,
            },
            is_eval=False,
            raw_task={
                taskset_config.format.prompt_key: "",
                taskset_config.format.response_key: r"36",
            },
        )
        workflow = task.to_workflow(model=model)
        experiences = workflow.run()
        self.assertEqual(experiences[0].reward, 1.0)
        self.assertEqual(experiences[1].reward, 1.0)
        self.assertEqual(experiences[2].reward, 1.0)
        self.assertEqual(experiences[3].reward, 0.0)
        task_new = Task(
            workflow=MathBoxedWorkflow,
            repeat_times=taskset_config.repeat_times,
            format_args=taskset_config.format,
            rollout_args=taskset_config.rollout_args,
            workflow_args={
                "with_think": True,
                "format_score_coef": 0.2,
            },
            is_eval=False,
            raw_task={
                taskset_config.format.prompt_key: "",
                taskset_config.format.response_key: r"36",
            },
        )
        workflow.reset(task_new)
        workflow_new = task_new.to_workflow(model=model)
        experiences = workflow_new.run()
        self.assertEqual(experiences[0].reward, 1.0)
        self.assertEqual(experiences[1].reward, 0.8)
        self.assertEqual(experiences[2].reward, 0.8)
        self.assertEqual(experiences[3].reward, 0.0)

    def test_gsm8k_workflow(self) -> None:
        model = MagicMock()
        model.chat.return_value = [
            MockResponse("<think> balabalabala 99 </think>\n<answer> 36 </answer>"),
            MockResponse("<answer> 36.0 </answer>"),
            MockResponse("<answer>Kim's total points are 6 + 30 = 36 </answer>"),
            MockResponse("<think> balalaba </think><answer> 35.00 </answer>"),
        ]
        taskset_config = get_unittest_dataset_config("countdown")
        task = Task(
            workflow=MathWorkflow,
            repeat_times=taskset_config.repeat_times,
            format_args=taskset_config.format,
            rollout_args=taskset_config.rollout_args,
            is_eval=False,
            raw_task={
                taskset_config.format.prompt_key: "",
                taskset_config.format.response_key: r"36",
            },
        )
        workflow = task.to_workflow(model=model)
        experiences = workflow.run()
        self.assertEqual(experiences[0].reward, 1.1)
        self.assertEqual(experiences[1].reward, 0.9)
        self.assertEqual(experiences[2].reward, 0.9)
        self.assertEqual(experiences[3].reward, 0.1)
        task_new = Task(
            workflow=MathWorkflow,
            repeat_times=taskset_config.repeat_times,
            format_args=taskset_config.format,
            rollout_args=taskset_config.rollout_args,
            is_eval=False,
            raw_task={
                taskset_config.format.prompt_key: "",
                taskset_config.format.response_key: r"35",
            },
        )
        workflow.reset(task_new)
        workflow_new = task_new.to_workflow(model=model)
        experiences = workflow_new.run()
        self.assertEqual(experiences[0].reward, 0.1)
        self.assertEqual(experiences[1].reward, -0.1)
        self.assertEqual(experiences[2].reward, -0.1)
        self.assertEqual(experiences[3].reward, 1.1)

    def test_math_eval_workflow(self) -> None:
        model = MagicMock()
        model.chat.return_value = [
            MockResponse("My step-by-step reasoning leads to the answer \boxed{36}"),
            MockResponse("Here is the answer of \boxed{36.0}"),
            MockResponse("I made a mistake, the answer is \boxed{42}"),
            MockResponse("The answer is 36, but I forgot the box."),
        ]

        taskset_config = get_unittest_dataset_config("countdown")
        task = Task(
            workflow=MathEvalWorkflow,
            repeat_times=taskset_config.repeat_times,
            is_eval=True,
            format_args=taskset_config.format,
            raw_task={
                taskset_config.format.prompt_key: "",
                taskset_config.format.response_key: "36",
            },
        )

        workflow = task.to_workflow(model=model)
        experiences = workflow.run()
        self.assertEqual(len(experiences), 4)
        expected_accuracies = [1.0, 1.0, 0.0, 0.0]
        for i, (exp, expected_acc) in enumerate(zip(experiences, expected_accuracies)):
            with self.subTest(f"Response {i}"):
                self.assertEqual(exp.reward, 0.0)
                assert exp.metrics is not None, f"Metrics for response {i} should not be None"
                self.assertEqual(exp.metrics["accuracy"], expected_acc)

    @parameterized.expand([(DummyWorkflow,), (DummyAsyncWorkflow,)])
    def test_workflow_resettable(self, workflow_cls) -> None:
        model = MagicMock()
        json_task = Task(
            workflow=workflow_cls,
            repeat_times=1,
            raw_task={"a": 1},
            workflow_args={"output_format": "json"},
        )
        yaml_task = Task(
            workflow=workflow_cls,
            repeat_times=1,
            raw_task={"a": 1},
            workflow_args={"output_format": "yaml"},
        )
        workflow = json_task.to_workflow(model)
        if workflow.asynchronous:
            answer = asyncio.run(workflow.run_async())
        else:
            answer = workflow.run()
        self.assertEqual(answer[0].response_text, '{"a": 1}')
        workflow.reset(yaml_task)
        if workflow.asynchronous:
            answer = asyncio.run(workflow.run_async())
        else:
            answer = workflow.run()
        self.assertEqual(answer[0].response_text, "a: 1\n")

    @parameterized.expand([(DummyWorkflow,), (DummyAsyncWorkflow,)])
    def test_workflow_repeatable(self, workflow_cls) -> None:
        model = MagicMock()
        task = Task(
            workflow=workflow_cls,
            repeat_times=3,
            raw_task={"a": 1},
            workflow_args={"output_format": "json"},
        )
        workflow = task.to_workflow(model)
        workflow.set_repeat_times(2, run_id_base=0)
        self.assertEqual(workflow.repeat_times, 2)
        if workflow.asynchronous:
            answer = asyncio.run(workflow.run_async())
        else:
            answer = workflow.run()
        self.assertEqual(len(answer), 2)


@parameterized_class(
    ("workflow_cls",),
    [
        (DummyMultiTurnWorkflow,),
        (DummyAsyncMultiTurnWorkflow,),
    ],
)
class MultiTurnWorkflowTest(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        # configure the model
        self.config = get_template_config()
        self.config.mode = "explore"
        self.config.model.model_path = get_model_path()
        self.config.explorer.rollout_model.engine_num = 1  # self.engine_num
        self.config.explorer.rollout_model.tensor_parallel_size = 1  # self.tensor_parallel_size
        self.config.explorer.rollout_model.chat_template = CHAT_TEMPLATE
        self.config.algorithm.repeat_times = 2  # self.repeat_times
        self.config.explorer.rollout_model.enable_history = True  # self.enable_history
        self.config.check_and_update()
        allocator = Allocator(self.config.explorer)
        rollout_model, _ = await allocator.create_all_models()
        self.model_wrapper = rollout_model[0]
        await self.model_wrapper.prepare()

    async def test_multi_turn_workflow(self):
        task = Task(
            workflow=self.workflow_cls,
            repeat_times=3,
            raw_task={"contents": ["hello world!", "how are you?"]},
            workflow_args={"output_format": "json"},
        )
        workflow = task.to_workflow(self.model_wrapper)
        workflow.set_repeat_times(2, run_id_base=0)
        if workflow.asynchronous:
            answer = await workflow.run_async()
        else:
            answer = workflow.run()
        self.assertEqual(len(answer), 2)

    def tearDown(self):
        ray.shutdown(_exiting_interpreter=True)


class StateRecordingWorkflow(Workflow):
    is_async: bool = True

    def __init__(self, *, task, model: ModelWrapper, auxiliary_models):
        super().__init__(task=task, model=model, auxiliary_models=auxiliary_models)
        self.wait_time = task.workflow_args.get("wait_time", 1)

    async def run_async(self):
        for i in range(self.wait_time):
            await self.model.set_workflow_state({"step": i})
            await asyncio.sleep(1)
        return [Experience(tokens=Tensor([0, 1, 2]), prompt_length=1, reward=1.0)]


class TestWorkflowStateRecording(unittest.IsolatedAsyncioTestCase):
    async def test_workflow_state_recording(self):
        model = MagicMock()
        model_wrapper = ModelWrapper(model, config=InferenceModelConfig(model_path="dummy_model"))

        task = Task(
            workflow=StateRecordingWorkflow,
            repeat_times=3,
            raw_task={},
            workflow_args={"wait_time": 3},
        )
        workflow = task.to_workflow(model_wrapper)

        async def monitor_routine():
            old_state = {}
            count = 0
            for i in range(20):
                await asyncio.sleep(0.2)
                new_state = await model_wrapper.get_workflow_state()
                if new_state.get("step") != old_state.get("step"):
                    old_state = new_state
                    count += 1
            self.assertEqual(count, 3)
            return count

        await asyncio.gather(*[monitor_routine(), workflow.run_async()])


class TestAgentScopeWorkflowAdapter(unittest.IsolatedAsyncioTestCase):
    async def test_adapter_v1(self):
        try:
            from agentscope.model import ChatModelBase
            from agentscope.tuner import JudgeOutput, WorkflowOutput
        except ImportError:
            self.skipTest("agentscope >= 1.0.12 is not installed")

        async def as_workflow_func(task, model) -> WorkflowOutput:
            self.assertIsInstance(task, dict)
            self.assertIsInstance(model, ChatModelBase)
            return WorkflowOutput(
                reward=task["reward"],
                response=task["reward"],
                metrics={"workflow_metric_1": 0.0},
            )

        async def as_judge_func(task, response) -> JudgeOutput:
            self.assertIsInstance(task, dict)
            self.assertIsInstance(response, float)
            return JudgeOutput(
                reward=response,
                metrics={"judge_metric_1": 1.0},
            )

        model = MagicMock()
        openai_client = MagicMock()
        openai_client.model_path = "Qwen/Qwen3-8B"
        model.get_openai_async_client.return_value = openai_client
        model.extract_experience_from_history.return_value = [
            Experience(tokens=Tensor([0, 1, 2]), prompt_length=1, logprobs=Tensor([0.1, 0.2])),
        ]

        as_adapter_cls = WORKFLOWS.get("agentscope_workflow_adapter_v1")
        as_adapter = as_adapter_cls(
            task=Task(
                raw_task={"reward": 0.2},
                workflow_args={
                    "workflow_func": as_workflow_func,
                    "judge_func": as_judge_func,
                },
            ),
            model=model,
        )
        result = await as_adapter.run_async()
        self.assertEqual(len(result), 1)
        self.assertEqual(result[0].reward, 0.2)
        self.assertEqual(result[0].prompt_length, 1)
        metrics = result[-1].metrics
        self.assertEqual(len(metrics), 2)
        self.assertEqual(metrics["workflow_metric_1"], 0.0)
        self.assertEqual(metrics["judge_metric_1"], 1.0)


class DummyModelWrapper:
    def __init__(self, model, **kwargs):
        pass

    async def prepare(self):
        return

    def get_openai_client(self):
        return openai.OpenAI(api_key="EMPTY")

    def get_openai_async_client(self):
        return openai.AsyncOpenAI(api_key="EMPTY")

    async def clean_workflow_state(self):
        return

    @property
    async def model_version_async(self):
        return 0


class APIWorkflow(Workflow):
    is_async: bool = True

    def __init__(self, model: ModelWrapper, task: Task, auxiliary_models=None):
        super().__init__(task=task, model=model, auxiliary_models=auxiliary_models)
        self.client = model.get_openai_async_client()
        self.raise_except = task.raw_task.get("raise_except", False)

    async def run_async(self):
        _ = await self.client.chat.completions.create(
            model=self.client.model_path,
            messages=[{"role": "user", "content": "Hello!"}],
        )
        if self.raise_except:
            raise RuntimeError("Intentional Exception for testing.")
        exps = self.model.extract_experience_from_history()
        exps[0].reward = 0.5
        return exps


class PartialFailureWorkflow(Workflow):
    can_reset: bool = True

    _call_lock = threading.Lock()
    _call_count = 0

    def __init__(self, model: ModelWrapper, task: Task, auxiliary_models=None):
        super().__init__(task=task, model=model, auxiliary_models=auxiliary_models)
        self.fail_call_ids = set(task.raw_task.get("fail_call_ids", []))

    def reset(self, task: Task):
        self.fail_call_ids = set(task.raw_task.get("fail_call_ids", []))

    @classmethod
    def reset_call_count(cls):
        with cls._call_lock:
            cls._call_count = 0

    @classmethod
    def next_call_id(cls) -> int:
        with cls._call_lock:
            call_id = cls._call_count
            cls._call_count += 1
            return call_id

    def run(self):
        call_id = self.next_call_id()
        if call_id in self.fail_call_ids:
            raise RuntimeError(f"Intentional failure for run call {call_id}")

        exp = Experience(
            tokens=Tensor([0, 1, 2]),
            prompt_length=1,
            metrics={"run_metrics": float(call_id)},
        )
        exp.response_text = str(call_id)
        return [exp]


class TestWorkflowRunner(unittest.IsolatedAsyncioTestCase):
    async def test_workflow_runner(self):
        config = get_template_config()
        config.explorer.auxiliary_models = [
            config.explorer.rollout_model,
            config.explorer.rollout_model,
        ]

        with patch_runner_models(
            DummyModelWrapper(MagicMock()),
            DummyModelWrapper(MagicMock()),
            DummyModelWrapper(MagicMock()),
        ):
            runner = WorkflowRunner(
                config,
                rollout_model_id=0,
                auxiliary_model_ids=[0, 1],
                runner_id=0,
            )
            await runner.prepare()
            task = Task(
                workflow=DummyWorkflow,
                repeat_times=3,
                raw_task={"a": 1},
                workflow_args={"output_format": "json"},
            )

            status, exps = await runner.run_task(
                task, batch_id="test", repeat_times=3, run_id_base=0
            )
            exps = deserialize_experiences(exps)

            self.assertTrue(status.ok)
            self.assertIsInstance(exps, list)
            self.assertEqual(len(exps), 3)

            task = Task(
                workflow=DummyAsyncWorkflow,
                repeat_times=2,
                raw_task={"a": 1},
                workflow_args={"output_format": "yaml"},
            )

            status, exps = await runner.run_task(
                task, batch_id="test", repeat_times=2, run_id_base=0
            )
            exps = deserialize_experiences(exps)
            self.assertTrue(status.ok)
            self.assertIsInstance(exps, list)
            self.assertEqual(len(exps), 2)

    @parameterized.expand(
        [
            ("sequential", 2),
            ("asynchronous", 2),
            ("multi-threading", 2),
        ]
    )
    async def test_workflow_runner_partial_success_non_repeatable(
        self, concurrent_mode: str, expected_success_runs: int
    ):
        config = get_template_config()
        config.explorer.concurrent_mode = concurrent_mode

        with patch_runner_models(DummyModelWrapper(MagicMock())):
            runner = WorkflowRunner(
                config,
                rollout_model_id=0,
                runner_id=0,
            )
            await runner.prepare()
            PartialFailureWorkflow.reset_call_count()
            task = Task(
                workflow=PartialFailureWorkflow,
                repeat_times=3,
                raw_task={"fail_call_ids": [1]},
            )

            status, exps = await runner.run_task(
                task, batch_id="test", repeat_times=3, run_id_base=0
            )
            exps = deserialize_experiences(exps)

            self.assertFalse(status.ok)
            self.assertEqual(status.completed_runs, expected_success_runs)
            self.assertEqual(status.total_runs, 3)
            self.assertEqual(len(exps), expected_success_runs)

            # One internal run fails with call_id=1, so runner-level metrics should
            # retain only the successful runs from this single subtask: call_id=0 and 2.
            self.assertEqual(len(status.metrics), expected_success_runs)
            self.assertEqual(
                sorted(metric["run_metrics"] for metric in status.metrics),
                [0.0, 2.0],
            )

            # Experiences returned from the runner should match the same successful
            # run set, proving failed runs do not leak into partial-return outputs.
            self.assertEqual(
                sorted(exp.metrics["run_metrics"] for exp in exps if exp.metrics),
                [0.0, 2.0],
            )
            self.assertIn(f"{expected_success_runs}/3 runs completed successfully", status.message)  # type: ignore[arg-type]

    @parameterized.expand(
        [
            ("sequential",),
            ("asynchronous",),
            ("multi-threading",),
        ]
    )
    async def test_workflow_runner_fail_fast_without_partial_collection(self, concurrent_mode: str):
        config = get_template_config()
        config.explorer.concurrent_mode = concurrent_mode

        with patch_runner_models(DummyModelWrapper(MagicMock())):
            runner = WorkflowRunner(
                config,
                rollout_model_id=0,
                runner_id=0,
            )
            task = Task(
                workflow=PartialFailureWorkflow,
                repeat_times=3,
                raw_task={"fail_call_ids": []},
            )
            await runner.prepare()

            async def mock_execute_single_run(
                workflow: Workflow,
                task: Task,
                run_index: int,
                run_id_base: int,
            ):
                if run_index == 0:
                    await asyncio.sleep(0.01)
                    exp = Experience(
                        tokens=Tensor([0, 1, 2]),
                        prompt_length=1,
                        metrics={"run_metrics": 0.0},
                    )
                    return True, [exp], {"run_metrics": 0.0}, None
                if run_index == 1:
                    await asyncio.sleep(0.02)
                    return False, [], None, "planned failure"
                await asyncio.sleep(0.5)
                exp = Experience(
                    tokens=Tensor([0, 1, 2]),
                    prompt_length=1,
                    metrics={"run_metrics": 2.0},
                )
                return True, [exp], {"run_metrics": 2.0}, None

            runner._execute_single_run = AsyncMock(side_effect=mock_execute_single_run)

            status, exps = await runner.run_task(
                task,
                batch_id="test",
                repeat_times=3,
                run_id_base=0,
                collect_partial_runs=False,
            )
            exps = deserialize_experiences(exps)

            self.assertFalse(status.ok)
            self.assertEqual(status.completed_runs, 1)
            self.assertEqual(status.total_runs, 3)
            self.assertEqual(len(exps), 1)
            self.assertIn("1/3 runs completed successfully", status.message)  # type: ignore[arg-type]

    async def test_workflow_runner_get_state(self):
        config = get_template_config()

        async def mock_get_api_server_url_remote():
            return None

        async def mock_get_model_version_remote():
            return 1

        async def mock_get_api_key_remote():
            return "dummy_api_key"

        async def mock_get_model_config_remote():
            return InferenceModelConfig(model_path="dummy_model")

        model = MagicMock()
        model.get_api_server_url.remote = MagicMock(side_effect=mock_get_api_server_url_remote)
        model.get_model_version.remote = MagicMock(side_effect=mock_get_model_version_remote)
        model.get_api_key.remote = MagicMock(side_effect=mock_get_api_key_remote)
        model.get_model_config.remote = MagicMock(side_effect=mock_get_model_config_remote)

        with patch_runner_models(
            ModelWrapper(model, config=InferenceModelConfig(model_path="dummy_model"))
        ):
            runner = WorkflowRunner(
                config,
                rollout_model_id=0,
                runner_id=1,
            )
            await runner.prepare()
        task = Task(
            workflow=StateRecordingWorkflow,
            raw_task={},
            workflow_args={"wait_time": 2},
            batch_id=1,
            task_id=2,
        )

        async def monitor_routine():
            state_history = defaultdict(set)
            count = 0
            for i in range(20):
                await asyncio.sleep(0.4)
                new_state = await runner.get_runner_state()
                for k, v in new_state.items():
                    state_history[k].add(v)
            self.assertEqual(len(state_history["model_version"]), 1)
            self.assertEqual(len(state_history["workflow_id"]), 3)
            self.assertEqual(len(state_history["begin_time"]), 3)
            self.assertEqual(len(state_history["step"]), 2)
            return count

        await asyncio.gather(
            *[
                monitor_routine(),
                runner.run_task(task, batch_id="test", repeat_times=3, run_id_base=0),
            ]
        )

    async def test_workflow_with_openai(self):
        config = get_template_config()
        config.mode = "explore"
        config.model.model_path = get_model_path()
        config.explorer.rollout_model.engine_num = 1
        config.explorer.rollout_model.enable_openai_api = True
        config.explorer.rollout_model.enable_history = True
        config.check_and_update()
        allocator = Allocator(config.explorer)
        rollout_model, auxiliary_models = await allocator.create_all_models()
        runner = WorkflowRunner(
            config,
            rollout_model_id=0,
            runner_id=0,
        )
        await runner.prepare()
        tasks = [
            Task(
                workflow=APIWorkflow,
                raw_task={"raise_except": True},
                repeat_times=2,
            ),
            Task(
                workflow=APIWorkflow,
                raw_task={},
                repeat_times=2,
            ),
        ]

        status, exps = await runner.run_task(
            tasks[0], batch_id="test", repeat_times=2, run_id_base=0
        )  # test exception handling
        exps = deserialize_experiences(exps)
        self.assertEqual(status.ok, False)
        self.assertEqual(len(exps), 0)
        exps = runner.model_wrapper.extract_experience_from_history(clear_history=False)
        self.assertEqual(len(exps), 1)
        status, exps = await runner.run_task(
            tasks[1], batch_id="test", repeat_times=2, run_id_base=0
        )  # normal run
        exps = deserialize_experiences(exps)
        self.assertEqual(status.ok, True)
        self.assertEqual(len(exps), 2)
        exps = runner.model_wrapper.extract_experience_from_history(clear_history=False)
        self.assertEqual(len(exps), 0)
        self.assertEqual(len(rollout_model), 1)
        await rollout_model[0].shutdown()

    def tearDown(self):
        ray.shutdown(_exiting_interpreter=True)


class ConcurrentTestWorkflow(Workflow):
    is_async: bool = True

    def __init__(self, model: ModelWrapper, task: Task, auxiliary_models=None):
        super().__init__(task=task, model=model, auxiliary_models=auxiliary_models)
        self.client = self.model.get_openai_async_client()

    async def run_async(self):
        assert self.task.raw_task is not None
        _ = await self.model.chat_async([{"role": "user", "content": self.task.raw_task["text"]}])
        await asyncio.sleep(1.0)
        _ = await self.client.chat.completions.create(
            model=self.client.model_path,
            messages=[{"role": "user", "content": self.task.raw_task["text"]}],
        )
        history_exps = self.model.extract_experience_from_history()
        assert len(history_exps) == 2
        assert history_exps[0].prompt_length == history_exps[1].prompt_length
        prompt_length = history_exps[0].prompt_length
        assert (
            history_exps[0].tokens[:prompt_length].shape
            == history_exps[1].tokens[:prompt_length].shape
        )
        self.logger.debug("[DEBUG MESSAGE]")
        self.logger.info("[INFO MESSAGE]")
        self.logger.warning("[WARNING MESSAGE]")
        return history_exps


class TestConcurrentWorkflowRunner(RayUnittestBaseAsync):
    def setUp(self) -> None:
        config = get_template_config()
        config.mode = "explore"
        config.model.model_path = get_model_path()
        config.explorer.rollout_model.engine_num = 1
        config.explorer.rollout_model.enable_history = True
        config.explorer.rollout_model.enable_openai_api = True
        config.check_and_update()
        self.config = config
        os.makedirs(self.config.log.save_dir, exist_ok=True)

    async def test_concurrent_workflow_runner(self):
        allocator = Allocator(self.config.explorer)
        rollout_model, _ = await allocator.create_all_models()

        self.config.explorer.concurrent_mode = "sequential"
        sequential_runner = (
            ray.remote(WorkflowRunner)
            .options(
                runtime_env={
                    "env_vars": {
                        LOG_DIR_ENV_VAR: os.path.join(self.config.log.save_dir),
                        LOG_LEVEL_ENV_VAR: "DEBUG",
                    },
                }
            )
            .remote(
                self.config,
                rollout_model_id=0,
                runner_id=0,
            )
        )
        self.config.explorer.concurrent_mode = "asynchronous"
        async_runner = (
            ray.remote(WorkflowRunner)
            .options(
                runtime_env={
                    "env_vars": {
                        LOG_DIR_ENV_VAR: os.path.join(self.config.log.save_dir),
                        LOG_LEVEL_ENV_VAR: "INFO",
                    },
                }
            )
            .remote(
                self.config,
                rollout_model_id=0,
                runner_id=1,
            )
        )
        thread_runner = (
            ray.remote(WorkflowRunner)
            .options(
                runtime_env={
                    "env_vars": {
                        LOG_DIR_ENV_VAR: os.path.join(self.config.log.save_dir),
                        LOG_LEVEL_ENV_VAR: "WARNING",
                    },
                }
            )
            .remote(
                self.config,
                rollout_model_id=0,
                runner_id=2,
            )
        )
        await asyncio.gather(
            sequential_runner.prepare.remote(),
            async_runner.prepare.remote(),
            thread_runner.prepare.remote(),
        )
        task = Task(
            workflow=ConcurrentTestWorkflow,
            repeat_times=4,
            raw_task={"text": "Hello, world!"},
        )

        # warmup
        async_status, async_exps = await async_runner.run_task.remote(
            task, batch_id="test", repeat_times=2, run_id_base=0
        )

        st = time.time()
        async_status, async_exps = await async_runner.run_task.remote(
            task, batch_id="test", repeat_times=4, run_id_base=0
        )
        async_runtime = time.time() - st

        # warmup
        thread_status, thread_exps = await thread_runner.run_task.remote(
            task, batch_id="test", repeat_times=1, run_id_base=0
        )

        st = time.time()
        thread_status, thread_exps = await thread_runner.run_task.remote(
            task, batch_id="test", repeat_times=4, run_id_base=0
        )
        thread_runtime = time.time() - st
        st = time.time()
        sequential_status, sequential_exps = await sequential_runner.run_task.remote(
            task, batch_id="test", repeat_times=4, run_id_base=0
        )
        sequential_runtime = time.time() - st

        self.assertTrue(async_status.ok)
        self.assertTrue(thread_status.ok)
        self.assertTrue(sequential_status.ok)

        async_exps = deserialize_experiences(async_exps)
        thread_exps = deserialize_experiences(thread_exps)
        sequential_exps = deserialize_experiences(sequential_exps)

        self.assertEqual(len(async_exps), 8)
        self.assertEqual(len(thread_exps), 8)
        self.assertEqual(len(sequential_exps), 8)

        self.assertLessEqual(async_runtime * 2, sequential_runtime)
        self.assertLessEqual(thread_runtime * 2, sequential_runtime)

        # check log files
        sequential_log_path = os.path.join(self.config.log.save_dir, "explorer_runner_0.log")
        async_log_path = os.path.join(self.config.log.save_dir, "explorer_runner_1.log")
        thread_log_path = os.path.join(self.config.log.save_dir, "explorer_runner_2.log")
        with open(sequential_log_path, "r") as f:
            sequential_logs = f.read()
            assert "[DEBUG MESSAGE]" in sequential_logs
            assert "[INFO MESSAGE]" in sequential_logs
            assert "[WARNING MESSAGE]" in sequential_logs
            # count the occurrences of each log level
            debug_count = sequential_logs.count("[DEBUG MESSAGE]")
            info_count = sequential_logs.count("[INFO MESSAGE]")
            warning_count = sequential_logs.count("[WARNING MESSAGE]")
            assert debug_count == 4
            assert info_count == 4
            assert warning_count == 4
        with open(async_log_path, "r") as f:
            async_logs = f.read()
            assert "[DEBUG MESSAGE]" not in async_logs
            assert "[INFO MESSAGE]" in async_logs
            assert "[WARNING MESSAGE]" in async_logs
            info_count = async_logs.count("[INFO MESSAGE]")
            warning_count = async_logs.count("[WARNING MESSAGE]")
            assert info_count == 6
            assert warning_count == 6
        with open(thread_log_path, "r") as f:
            thread_logs = f.read()
            assert "[DEBUG MESSAGE]" not in thread_logs
            assert "[INFO MESSAGE]" not in thread_logs
            assert "[WARNING MESSAGE]" in thread_logs
            warning_count = thread_logs.count("[WARNING MESSAGE]")
            assert warning_count == 5
        await rollout_model[0].shutdown()

    def tearDown(self):
        shutil.rmtree(self.config.log.save_dir, ignore_errors=True)
