# SPDX-License-Identifier: Apache-2.0

import asyncio
import atexit
import os
import threading
from collections.abc import Callable
from concurrent.futures import ProcessPoolExecutor
from concurrent.futures.process import BrokenProcessPool
from functools import partial

from areal.utils import logging

logger = logging.getLogger("RewardAPI")


def _get_device_count_safely() -> int:
    """
    Safely get device count without initializing CUDA context.
    """
    gpu_types = ["nvidia", "davinci"]
    try:
        if os.path.exists("/dev"):
            for gpu_type in gpu_types:
                devices = [
                    f
                    for f in os.listdir("/dev")
                    if f.startswith(gpu_type) and f[len(gpu_type) :].isdigit()
                ]
                if devices:
                    return len(devices)
    except (OSError, ValueError) as e:
        # /dev doesn't exist or can't read (e.g., Windows, macOS)
        logger.debug(f"Could not read device list from /dev, using fallback: {e}")

    # Fallback: assume 8 devices for cautious max_workers calculation
    return 8


def reward_fn(
    prompt: str,
    completions: str,
    prompt_ids: list[int],
    completion_ids: list[int],
    **kwargs,
):
    """This function is a placeholder for the reward function that will be used in the RLVR pipeline.

    In general, there's no restriction on the signature and implementation of this function in customized rollout workflows.
    It would be convinent to follow this signature and directly use it in our predefined rollout workflows.

    :param prompt: The string representing the task to be completed.
    :param completions: The string representing the trajectory generated by the model.
    :param prompt_ids: The token IDs of the prompt.
    :param completion_ids: The token IDs of the trajectory generated by the model.
    :param kwargs: Other attributes of the data in the dataset, such as solutions, input_outputs, etc.
        Any other attributes in the dataset will be passed as keyword arguments to this function.
    :rtype: float
    """


class AsyncRewardWrapper:
    """Wraps a synchronous reward function for async execution with timeout and retries.

    Executors are shared by ``max_workers`` key and cleaned up via ``atexit``.
    Includes automatic recovery from broken process pools.

    The reward function and its arguments must be picklable since they
    are dispatched to worker processes via ``ProcessPoolExecutor``.
    """

    _executors: dict[int, ProcessPoolExecutor] = {}
    _lock = threading.Lock()
    _atexit_registered = False

    def __init__(
        self,
        reward_fn: Callable,
        timeout_seconds: float = 15,
        max_workers: int | None = None,
        max_retries: int = 3,
    ):
        self.reward_fn = reward_fn
        self.timeout_seconds = timeout_seconds
        if max_workers is None:
            cpu_count = os.cpu_count() or 1
            device_count = _get_device_count_safely()
            max_workers = max((cpu_count // device_count) // 2, 1)
        self.max_workers = max_workers
        self.max_retries = max_retries
        self._executor_key = max_workers

        with self._lock:
            if self._executor_key not in self._executors:
                self._executors[self._executor_key] = ProcessPoolExecutor(
                    max_workers=max_workers
                )
            if not AsyncRewardWrapper._atexit_registered:
                atexit.register(AsyncRewardWrapper._atexit_shutdown_all)
                AsyncRewardWrapper._atexit_registered = True

    @classmethod
    def _atexit_shutdown_all(cls):
        """Shut down all executors before ``_python_exit`` to prevent
        worker processes deadlocking in ``_finalize_join``.
        Must use ``wait=True`` so the result queue is fully drained.
        """
        with cls._lock:
            for executor in cls._executors.values():
                try:
                    executor.shutdown(wait=True, cancel_futures=True)
                except Exception as e:
                    logger.warning(f"Error shutting down executor at exit: {e}")
            cls._executors.clear()

    @classmethod
    def _recreate_executor(
        cls,
        executor_key: int,
        max_workers: int,
        broken: ProcessPoolExecutor,
    ) -> ProcessPoolExecutor | None:
        with cls._lock:
            current = cls._executors.get(executor_key)
            if current is not broken:
                return current
            try:
                broken.shutdown(wait=False)
            except Exception as e:
                logger.warning(f"Error shutting down broken executor: {e}")
            try:
                new_executor = ProcessPoolExecutor(max_workers=max_workers)
            except Exception:
                logger.exception("Failed to create replacement ProcessPoolExecutor")
                cls._executors.pop(executor_key, None)
                return None
            cls._executors[executor_key] = new_executor
            logger.info(f"Recreated ProcessPoolExecutor with {max_workers} workers")
            return new_executor

    async def __call__(self, *args, **kwargs) -> float:
        for attempt in range(self.max_retries + 1):
            executor = self._executors.get(self._executor_key)
            if executor is None:
                raise RuntimeError("ProcessPoolExecutor has been shut down")

            is_last = attempt == self.max_retries
            try:
                future = asyncio.get_running_loop().run_in_executor(
                    executor,
                    partial(self.reward_fn, *args, **kwargs),
                )
                return await asyncio.wait_for(future, timeout=self.timeout_seconds)
            except TimeoutError:
                logger.warning(
                    f"Computing reward timeout after {self.timeout_seconds}s "
                    f"(attempt {attempt + 1}/{self.max_retries + 1}). "
                    f"{'Returning 0.' if is_last else 'Retrying...'}"
                )
                if is_last:
                    return 0
            except BrokenProcessPool:
                logger.warning(
                    f"ProcessPoolExecutor broken (attempt {attempt + 1}/{self.max_retries + 1}). "
                    "Attempting to recreate..."
                )
                if is_last:
                    raise
                if (
                    self._recreate_executor(
                        self._executor_key, self.max_workers, executor
                    )
                    is None
                ):
                    raise
            except Exception:
                logger.exception(
                    f"Reward computation error (attempt {attempt + 1}/{self.max_retries + 1})"
                )
                if is_last:
                    raise

        return 0
