"""TRLAdapter -- wrap TRL-style reward functions."""

from __future__ import annotations

from typing import Any, Callable, Literal, Optional

from rewardprobe.adapters.base import BaseAdapter
from rewardprobe.types import NormalizedEnvironment, NormalizedRewardFn


def _wrap_trl_fn(
    fn: Callable,
    format: Literal["standard", "conversational"] = "standard",
    dataset_columns: Optional[list[str]] = None,
) -> Callable[[str, dict], float]:
    """Wrap a TRL batch reward function to the sync ``(completion_text, task) -> float`` interface.

    TRL reward functions have several common signatures:
    * ``fn(completions, **kwargs) -> list[float]``
    * ``fn(prompts, completions, answer, **kwargs) -> list[float]``
    * ``fn(completions, solution, **kwargs) -> list[float]``

    We call them with single-element lists and return the single result.
    """
    import inspect

    columns = dataset_columns or []
    sig = inspect.signature(fn)
    param_names = [
        p.name
        for p in sig.parameters.values()
        if p.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD)
    ]

    # Detect if fn takes 'prompts' as first arg (common TRL pattern)
    has_prompts_first = len(param_names) >= 2 and param_names[0] == "prompts"

    def _call_sync(completion_text: str, task: dict) -> float:
        if format == "conversational":
            completions: Any = [[{"role": "assistant", "content": completion_text}]]
        else:
            completions = [completion_text]

        # Build kwargs from dataset columns
        kwargs: dict[str, Any] = {}
        for col in columns:
            if col in task:
                kwargs[col] = [task[col]]

        # Also forward any param names that match task keys
        for pname in param_names:
            if pname in ("prompts", "completions", "completion_ids"):
                continue
            if pname in task and pname not in kwargs:
                kwargs[pname] = [task[pname]]

        if has_prompts_first:
            # fn(prompts, completions, **kwargs)
            prompt_val = task.get("prompt", task.get("question", ""))
            if format == "conversational":
                prompts: Any = [[{"role": "user", "content": str(prompt_val)}]]
            else:
                prompts = [str(prompt_val)]
            results = fn(prompts, completions, **kwargs)
        else:
            results = fn(completions, **kwargs)

        result = results[0]
        return float(result) if result is not None else 0.0

    return _call_sync


class TRLAdapter(BaseAdapter):
    """Wrap TRL-style reward functions into a :class:`NormalizedEnvironment`.

    Parameters
    ----------
    fn:
        A reward function (or list) with the TRL batch signature:
        ``fn(completions, **kwargs) -> list[float]``.
    tasks:
        A list of task dicts.
    name:
        Human-readable environment name.
    format:
        ``"standard"`` passes completions as ``list[str]``;
        ``"conversational"`` passes them as ``list[list[dict]]``.
    reward_weights:
        Optional per-function weights.
    dataset_columns:
        Column names from the task dict to forward as keyword arguments
        to the reward function (each value wrapped in a single-element list).
    """

    def __init__(
        self,
        fn: Callable | list[Callable],
        tasks: list[dict[str, Any]],
        name: Optional[str] = None,
        format: Literal["standard", "conversational"] = "standard",
        reward_weights: Optional[list[float]] = None,
        dataset_columns: Optional[list[str]] = None,
    ) -> None:
        fns = fn if isinstance(fn, list) else [fn]
        self._fns = fns
        self._tasks = tasks
        self._name = name or fns[0].__name__
        self._format = format
        self._weights = reward_weights or [1.0] * len(fns)
        self._dataset_columns = dataset_columns

    def adapt(self) -> NormalizedEnvironment:
        reward_fns: list[NormalizedRewardFn] = []
        for func, weight in zip(self._fns, self._weights):
            wrapped = _wrap_trl_fn(
                func,
                format=self._format,
                dataset_columns=self._dataset_columns,
            )
            reward_fns.append(
                NormalizedRewardFn(
                    name=getattr(func, "__name__", "reward"),
                    fn=wrapped,
                    weight=weight,
                )
            )

        return NormalizedEnvironment(
            name=self._name,
            reward_fns=reward_fns,
            tasks=self._tasks,
            metadata={"source": "trl", "format": self._format},
        )
