# -*- coding: utf-8 -*-
"""The Workflow Runner Module."""

import asyncio
import os
import time
import traceback
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

from trinity.buffer import get_buffer_reader, get_buffer_writer
from trinity.common.config import Config, StorageConfig
from trinity.common.constants import LOG_DIR_ENV_VAR, LOG_LEVEL_ENV_VAR
from trinity.common.experience import Experience
from trinity.common.models.allocator import Allocator
from trinity.common.models.model import ModelWrapper
from trinity.common.workflows import Task, Workflow
from trinity.utils.log import get_logger
from trinity.utils.metrics import aggregate_run_level_metrics


@dataclass(frozen=True)
class Status:
    """Status of the task running result."""

    completed_runs: int
    total_runs: int
    metrics: List[Dict[str, float]]
    # A list of metric dictionaries, where each dictionary is from a single run.
    message: Optional[str] = None

    @property
    def ok(self) -> bool:
        return self.completed_runs == self.total_runs


@dataclass(frozen=True)
class RunnerExecutionResult:
    """Execution result for one runner task."""

    status: Status
    experiences: List[Experience]


class WorkflowRunner:
    """A Ray remote actor to run the workflow and generate experiences."""

    def __init__(
        self,
        config: Config,
        rollout_model_id: int,
        auxiliary_model_ids: Optional[List[int]] = None,
        runner_id: Optional[int] = None,
    ) -> None:
        self.name = f"{config.explorer.name}_runner_{runner_id}"
        self.logger = get_logger(self.name, in_ray_actor=True)
        self.config = config
        allocator = Allocator(config.explorer)
        self.model_wrapper: ModelWrapper = allocator.get_model(
            config.explorer.rollout_model, "rollout", rollout_model_id
        )
        self.auxiliary_model_wrappers: List[ModelWrapper] = [
            allocator.get_model(
                config.explorer.auxiliary_models[index], f"auxiliary_{index}", auxiliary_model_id
            )
            for index, auxiliary_model_id in enumerate(auxiliary_model_ids or [])
        ]
        self.workflow_instance: Workflow = None
        self.runner_id = runner_id
        self.runner_state = {
            "workflow_id": None,
            "model_version": None,
            "begin_time": 0,
            "terminate_time": 0,
        }
        self.concurrent_mode = config.explorer.concurrent_mode
        if self.concurrent_mode == "sequential":
            self.concurrent_run_fn = self._sequential_run
        elif self.concurrent_mode == "asynchronous":
            self.concurrent_run_fn = self._asynchronous_run
        elif self.concurrent_mode == "multi-threading":
            self.concurrent_run_fn = self._multi_threading_run
        else:
            self.logger.warning(
                f"Unknown concurrent_mode {self.concurrent_mode}, defaulting to sequential."
            )
            self.concurrent_run_fn = self._sequential_run
        self.logger.info(
            f"WorkflowRunner [{self.name}]({self.concurrent_mode}) initialized:\n"
            f"  > rollout model: {self.config.explorer.rollout_model.model_path}\n"
            f"  > auxiliary models: {[aux_model_config.model_path for aux_model_config in self.config.explorer.auxiliary_models]}"
        )

    async def prepare(self) -> None:
        """Prepare the runner."""
        await asyncio.gather(
            self.model_wrapper.prepare(),
            *(aux_model.prepare() for aux_model in self.auxiliary_model_wrappers),
        )
        self.logger.info(f"WorkflowRunner [{self.name}] is prepared and ready to run tasks.")

    def is_alive(self):
        return True

    def _create_workflow_instance(self, task: Task) -> Workflow:
        if task.workflow is None:
            raise ValueError("Workflow is not set in the task.")
        if (
            self.workflow_instance is None
            or not self.workflow_instance.__class__ == task.workflow
            or not self.workflow_instance.resettable
        ):
            # Pass ModelWrapper directly; Workflow.__init__ will get OpenAI clients automatically
            self.workflow_instance = task.to_workflow(
                self.model_wrapper,
                self.auxiliary_model_wrappers,
            )
        else:
            self.workflow_instance.reset(task)
        return self.workflow_instance

    async def _run_workflow(self, workflow_instance: Workflow) -> List[Experience]:
        if workflow_instance.asynchronous:
            exps = await workflow_instance.run_async()
        else:
            exps = workflow_instance.run()
        return exps

    def _create_isolated_workflow_instance(self, task: Task) -> Workflow:
        return task.to_workflow(
            (
                self.model_wrapper.clone_with_isolated_history()
                if self.config.explorer.rollout_model.enable_history
                else self.model_wrapper
            ),
            self.auxiliary_model_wrappers,
        )

    def _build_execution_result(
        self,
        total_runs: int,
        completed_runs: int,
        metrics: List[Dict[str, float]],
        experiences: List[Experience],
        first_error: Optional[str] = None,
    ) -> RunnerExecutionResult:
        if first_error is None:
            message = None
        elif completed_runs > 0:
            message = (
                f"{completed_runs}/{total_runs} runs completed successfully. "
                f"First error: {first_error}"
            )
        else:
            message = first_error

        return RunnerExecutionResult(
            status=Status(
                completed_runs=completed_runs,
                total_runs=total_runs,
                metrics=list(metrics),
                message=message,
            ),
            experiences=experiences,
        )

    def _aggregate_run_results(
        self,
        total_runs: int,
        results: List[Tuple[bool, List[Experience], Optional[Dict[str, float]], Optional[str]]],
    ) -> RunnerExecutionResult:
        exps = []
        run_metrics = []
        first_error = None

        for ok, new_exps, run_metric, error in results:
            if ok:
                exps.extend(new_exps)
                if run_metric is not None:
                    run_metrics.append(run_metric)
                continue
            if first_error is None:
                first_error = error

        return self._build_execution_result(
            total_runs=total_runs,
            completed_runs=len(run_metrics),
            metrics=run_metrics,
            experiences=exps,
            first_error=first_error,
        )

    async def _run_parallel_runs(
        self,
        task: Task,
        repeat_times: int,
        run_id_base: int,
        collect_partial_runs: bool = True,
        use_threads: bool = False,
    ) -> RunnerExecutionResult:
        async def run_single(
            i: int,
        ) -> Tuple[bool, List[Experience], Optional[Dict[str, float]], Optional[str]]:
            workflow = self._create_isolated_workflow_instance(task)
            return await self._execute_single_run(workflow, task, i, run_id_base)

        if collect_partial_runs:
            if use_threads:
                results = await asyncio.gather(
                    *(
                        asyncio.to_thread(lambda idx=i: asyncio.run(run_single(idx)))  # type: ignore[misc]
                        for i in range(repeat_times)
                    )
                )
            else:
                results = await asyncio.gather(*(run_single(i) for i in range(repeat_times)))
            return self._aggregate_run_results(repeat_times, results)

        future_to_run_index = {}
        for i in range(repeat_times):
            if use_threads:
                future = asyncio.create_task(
                    asyncio.to_thread(lambda idx=i: asyncio.run(run_single(idx)))  # type: ignore[misc]
                )
            else:
                future = asyncio.create_task(run_single(i))
            future_to_run_index[future] = i

        results = []
        while future_to_run_index:
            done, pending = await asyncio.wait(
                future_to_run_index.keys(),
                return_when=asyncio.FIRST_COMPLETED,
            )
            should_stop = False
            for future in done:
                future_to_run_index.pop(future)
                result = future.result()
                results.append(result)
                ok, _, _, _ = result
                if not ok:
                    should_stop = True
            if should_stop:
                for future in pending:
                    future.cancel()
                if pending:
                    await asyncio.gather(*pending, return_exceptions=True)
                break

        return self._aggregate_run_results(repeat_times, results)

    async def _execute_single_run(
        self,
        workflow: Workflow,
        task: Task,
        run_index: int,
        run_id_base: int,
    ) -> Tuple[bool, List[Experience], Optional[Dict[str, float]], Optional[str]]:
        st = time.time()
        await self.model_wrapper.clean_workflow_state()
        self.runner_state["workflow_id"] = f"{task.batch_id}/{task.task_id}/{run_index}"
        self.runner_state["terminate_time"] = None
        self.runner_state["begin_time"] = st
        try:
            new_exps = await self._run_workflow(workflow)
            et = time.time()
            self.runner_state["terminate_time"] = et
            run_metric = aggregate_run_level_metrics(
                [exp.metrics for exp in new_exps if exp.metrics]
            )
            run_metric["time/run_execution"] = et - st
            for exp in new_exps:
                exp.eid.run = run_id_base + run_index
            return True, new_exps, run_metric, None
        except Exception as exc:
            self.runner_state["terminate_time"] = time.time()
            error_trace_back = traceback.format_exc()
            self.logger.error(
                "WorkflowRunner single run error: " f"{exc}\nTraceback:\n{error_trace_back}"
            )
            return False, [], None, error_trace_back.rstrip()

    async def _run_task(
        self,
        task: Task,
        repeat_times: int,
        run_id_base: int,
        collect_partial_runs: bool = True,
    ) -> RunnerExecutionResult:
        """Init workflow from the task and run it."""
        if task.workflow.can_repeat:
            workflow_instance = self._create_workflow_instance(task)
            workflow_instance.set_repeat_times(repeat_times, run_id_base)
            st = time.time()
            await self.model_wrapper.clean_workflow_state()
            self.runner_state["workflow_id"] = f"{task.batch_id}/{task.task_id}/{run_id_base}"
            self.runner_state["terminate_time"] = None
            self.runner_state["begin_time"] = st
            exps = await self._run_workflow(workflow_instance)
            et = time.time()
            self.runner_state["terminate_time"] = et
            # repeatable workflow cannot calculate run level metrics, we use experience level metrics directly
            run_metrics = [exp.metrics for exp in exps if exp.metrics]
            for metric in run_metrics:
                metric["time/run_execution"] = et - st
            return self._build_execution_result(
                total_runs=repeat_times,
                completed_runs=repeat_times,
                metrics=run_metrics,
                experiences=exps,
            )
        else:
            return await self.concurrent_run_fn(
                task,
                repeat_times,
                run_id_base,
                collect_partial_runs=collect_partial_runs,
            )

    async def _sequential_run(
        self,
        task: Task,
        repeat_times: int,
        run_id_base: int,
        collect_partial_runs: bool = True,
    ) -> RunnerExecutionResult:
        results = []
        for i in range(repeat_times):
            workflow = self._create_workflow_instance(task)
            result = await self._execute_single_run(workflow, task, i, run_id_base)
            results.append(result)
            if collect_partial_runs:
                continue
            ok, _, _, _ = result
            if ok:
                continue
            break
        return self._aggregate_run_results(repeat_times, results)

    async def _asynchronous_run(
        self,
        task: Task,
        repeat_times: int,
        run_id_base: int,
        collect_partial_runs: bool = True,
    ) -> RunnerExecutionResult:
        return await self._run_parallel_runs(
            task,
            repeat_times,
            run_id_base,
            collect_partial_runs=collect_partial_runs,
        )

    async def _multi_threading_run(
        self,
        task: Task,
        repeat_times: int,
        run_id_base: int,
        collect_partial_runs: bool = True,
    ) -> RunnerExecutionResult:
        return await self._run_parallel_runs(
            task,
            repeat_times,
            run_id_base,
            collect_partial_runs=collect_partial_runs,
            use_threads=True,
        )

    async def get_runner_state(self) -> Dict:
        """Get the runner state."""
        runner_state = self.runner_state.copy()
        runner_state.update(await self.model_wrapper.get_workflow_state())
        return runner_state

    async def run_task(
        self,
        task: Task,
        batch_id: str,
        repeat_times: int = 1,
        run_id_base: int = 0,
        collect_partial_runs: bool = True,
    ) -> Tuple[Status, bytes]:
        """Run the task and return the states."""
        # TODO: avoid sending the experiences back to the scheduler to reduce the communication overhead
        st = time.time()
        try:
            model_version = await self.model_wrapper.model_version_async
            self.runner_state["model_version"] = model_version
            self.logger.info(
                f"Starting task: step={batch_id}, model_version={model_version}, repeat_times={repeat_times}, run_id_base={run_id_base}"
            )
            execution_result = await self._run_task(
                task,
                repeat_times,
                run_id_base,
                collect_partial_runs=collect_partial_runs,
            )
            exps = execution_result.experiences
            if execution_result.status.completed_runs > 0:
                assert exps is not None and len(exps) > 0, "An empty experience is generated"
            # set eid for each experience
            for exp in exps:
                exp.eid.batch = task.batch_id
                # keep exp.eid.task if it has been set before (e.g., in workflow)
                if exp.eid.task == "":  # "" is the default value
                    exp.eid.task = task.task_id
                if not hasattr(exp, "info") or exp.info is None:
                    exp.info = {}
                exp.info["model_version"] = model_version
                exp.info["use_count"] = 0
                exp.info["task_index"] = task.index

                if not hasattr(exp, "metrics") or exp.metrics is None:
                    exp.metrics = {}

            status = execution_result.status

            if task.is_eval:
                # If the task is an evaluation task, we do not record the experiences to the buffer
                return status, b""
            else:
                exp_payload = Experience.serialize_many(exps)
                return status, exp_payload

        except Exception as e:
            error_trace_back = traceback.format_exc()
            self.logger.error(f"WorkflowRunner run task error: {e}\nTraceback:\n{error_trace_back}")
            return (
                Status(
                    completed_runs=0,
                    total_runs=repeat_times,
                    metrics=[{"time/run_execution": time.time() - st}],
                    message=error_trace_back.rstrip(),
                ),
                b"",
            )


class DebugWorkflowRunner(WorkflowRunner):
    """A WorkflowRunner for debugging."""

    def __init__(
        self,
        config: Config,
        output_dir: str = "debug_output",
        enable_profiling: bool = False,
        disable_overwrite: bool = False,
    ) -> None:
        if disable_overwrite:
            # if output dir is not empty, change to a new dir with datetime suffix
            if os.path.isdir(output_dir) and os.listdir(output_dir):
                suffix = time.strftime("%Y%m%d%H%M%S", time.localtime())
                output_dir = f"{output_dir}_{suffix}"
        os.environ[LOG_DIR_ENV_VAR] = os.path.join(output_dir, "log")
        os.environ[LOG_LEVEL_ENV_VAR] = "DEBUG"
        super().__init__(
            config=config,
            rollout_model_id=0,
            auxiliary_model_ids=[0] * len(config.explorer.auxiliary_models),
            runner_id=0,
        )
        self.taskset = get_buffer_reader(config.buffer.explorer_input.tasksets[0])
        self.output_dir = output_dir
        self.enable_profiling = enable_profiling
        self.logger.info(f"Debug output directory: {self.output_dir}")
        os.makedirs(self.output_dir, exist_ok=True)
        self.output_profiling_file = os.path.join(
            self.output_dir,
            "profiling.html",
        )
        self.output_sqlite_file = "sqlite:///" + os.path.join(
            self.output_dir,
            "experiences.db",
        )
        self.sqlite_writer = get_buffer_writer(
            StorageConfig(
                name="debug_buffer",
                schema_type="experience",
                path=self.output_sqlite_file,
                storage_type="sql",
                batch_size=1,
                wrap_in_ray=False,
            )
        )

    async def debug(self) -> None:
        """Run the debug workflow."""
        tasks = await self.taskset.read_async(batch_size=1)
        task = tasks[0]
        self.logger.info(f"Start debugging task:\n{task.raw_task}")
        if not self.enable_profiling:
            status, exp_payload = await self.run_task(
                task=task, batch_id="debug", repeat_times=1, run_id_base=0
            )
        else:
            from viztracer import VizTracer

            with VizTracer(output_file=self.output_profiling_file):
                status, exp_payload = await self.run_task(
                    task=task, batch_id="debug", repeat_times=1, run_id_base=0
                )
        experiences = Experience.deserialize_many(exp_payload) if exp_payload else []
        if not status.ok and not experiences:
            experiences = self.model_wrapper.extract_experience_from_history()
            self.logger.info(
                f"Debugging failed, extracting {len(experiences)} experiences from history."
            )
        await self.sqlite_writer.write_async(experiences)
        if status.ok:
            print(f"Task {task.task_id} completed successfully with metrics:\n{status.metrics}")
        else:
            self.logger.error(f"Task {task.task_id} failed with message: {status.message}")
        self.logger.info("Debugging completed.")
