"""DAPO integration built on top of DeepGym's reward engine.

DAPO is commonly run through ``verl`` recipes rather than a standalone trainer
API. DeepGym keeps the integration thin: expose reward callbacks and helper
config/module generators that fit naturally into an external DAPO training run.
"""

from __future__ import annotations

import asyncio
from collections.abc import Awaitable, Callable

from deepgym.core import DeepGym
from deepgym.models import Environment


def make_dapo_reward_fn(
    env: Environment,
    dg: DeepGym | None = None,
    max_parallel: int = 32,
) -> Callable[..., list[float]]:
    """Create a sync reward function for DAPO-style training loops.

    The callable accepts decoded completions plus arbitrary per-sample metadata
    and returns one scalar reward per completion.
    """
    _dg = dg or DeepGym(mode='auto')

    def reward_fn(completions: list[str], **kwargs: object) -> list[float]:
        if not completions:
            return []
        batch = _dg.run_batch(
            env,
            completions,
            max_parallel=min(len(completions), max_parallel),
            **kwargs,
        )
        return [result.score for result in batch.results]

    return reward_fn


def make_dapo_async_reward_fn(
    env: Environment,
    dg: DeepGym | None = None,
    max_parallel: int = 32,
) -> Callable[..., Awaitable[list[float]]]:
    """Create an async reward function for DAPO-style training loops."""
    _dg = dg or DeepGym(mode='auto')

    async def reward_fn(completions: list[str], **kwargs: object) -> list[float]:
        if not completions:
            return []
        loop = asyncio.get_running_loop()
        batch = await loop.run_in_executor(
            None,
            lambda: _dg.run_batch(
                env,
                completions,
                max_parallel=min(len(completions), max_parallel),
                **kwargs,
            ),
        )
        return [result.score for result in batch.results]

    return reward_fn


def make_dapo_shaped_reward_fn(
    env: Environment,
    dg: DeepGym | None = None,
    *,
    component: str | None = None,
    max_parallel: int = 32,
) -> Callable[..., list[float]]:
    """Return a specific shaped reward component for each completion.

    If ``component`` is ``None``, the aggregate DeepGym score is returned.
    """
    _dg = dg or DeepGym(mode='auto')

    def reward_fn(completions: list[str], **kwargs: object) -> list[float]:
        if not completions:
            return []
        batch = _dg.run_batch(
            env,
            completions,
            max_parallel=min(len(completions), max_parallel),
            **kwargs,
        )
        if component is None:
            return [result.score for result in batch.results]
        return [
            (result.reward_components or {}).get(component, result.score)
            for result in batch.results
        ]

    return reward_fn


_DAPO_VERL_CONFIG_TEMPLATE = """\
# verl DAPO config with DeepGym reward scoring
# Generated by deepgym.integrations.dapo
#
# Usage:
#   1. Write the generated reward module next to your training entrypoint
#   2. Point verl's custom reward hook at that module

algorithm:
  adv_estimator: dapo

data:
  train_files: {train_files}
  val_files: {val_files}
  prompt_key: {prompt_key}
  response_key: {response_key}

custom_reward_function:
  path: {reward_module_path}
  name: {reward_fn_name}

trainer:
  default_local_dir: {output_dir}
  total_epochs: {epochs}

rollout:
  n: {samples_per_prompt}

logging:
  project: {project_name}
"""


def generate_dapo_verl_config(
    *,
    train_files: str,
    reward_module_path: str,
    reward_fn_name: str = 'reward_fn',
    output_dir: str = 'outputs/dapo',
    val_files: str = '',
    prompt_key: str = 'prompt',
    response_key: str = 'response',
    epochs: int = 1,
    samples_per_prompt: int = 8,
    project_name: str = 'deepgym-dapo',
) -> str:
    """Generate a minimal verl-flavored DAPO config snippet.

    This is intentionally thin and focuses on the reward bridge fields DeepGym
    can own directly.
    """
    return _DAPO_VERL_CONFIG_TEMPLATE.format(
        train_files=train_files,
        val_files=val_files or train_files,
        prompt_key=prompt_key,
        response_key=response_key,
        reward_module_path=reward_module_path,
        reward_fn_name=reward_fn_name,
        output_dir=output_dir,
        epochs=epochs,
        samples_per_prompt=samples_per_prompt,
        project_name=project_name,
    )


_DAPO_REWARD_MODULE_TEMPLATE = """\
from deepgym import load_environment
from deepgym.integrations.dapo import make_dapo_reward_fn

env = load_environment({env_name!r})
reward_fn = make_dapo_reward_fn(env)
"""


def generate_dapo_reward_module(env_name: str) -> str:
    """Generate a tiny reward module for external DAPO/verl jobs."""
    return _DAPO_REWARD_MODULE_TEMPLATE.format(env_name=env_name)
