"""RawAdapter -- wrap any Python callable as a reward function."""

from __future__ import annotations

import inspect
from typing import Any, Callable, Optional

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


def _wrap_callable(
    fn: Callable,
) -> Callable[[str, dict], float]:
    """Inspect *fn*'s signature and return a ``(completion_text, task) -> float`` wrapper.

    Supported raw signatures
    ~~~~~~~~~~~~~~~~~~~~~~~~
    * ``fn(completion: str, answer: str) -> float``  -- most common
    * ``fn(completion: str) -> float``                -- no ground truth
    * ``fn(completion: str, **kwargs) -> float``      -- flexible
    * ``fn(completion, answer, question, ...) -> float`` -- named params pulled
      from the task dict
    """
    sig = inspect.signature(fn)
    params = list(sig.parameters.values())

    has_var_keyword = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params)

    # Single positional param -- completion only
    if len(params) == 1 and not has_var_keyword:

        def _wrap_single(completion_text: str, task: dict) -> float:
            return float(fn(completion_text))

        return _wrap_single

    # Known completion-like first param names
    _COMPLETION_NAMES = {
        "completion",
        "completion_text",
        "response",
        "text",
        "output",
        "solution_str",
        "solution",
        "pred",
        "prediction",
        "generated",
        "model_output",
    }
    # Known answer-like second param names
    _ANSWER_NAMES = {
        "answer",
        "expected",
        "ground_truth",
        "target",
        "label",
        "solution",
        "gold",
        "reference",
        "correct",
    }

    # Two positional params (completion, answer) -- most common raw form.
    # Heuristic: if exactly 2 required params with no **kwargs, assume
    # first = completion, second = answer — even if names don't match.
    non_default_params = [
        p
        for p in params
        if p.default is inspect.Parameter.empty
        and p.kind not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD)
    ]

    if len(non_default_params) == 2 and not has_var_keyword:
        p1_name = non_default_params[1].name

        def _wrap_pair(completion_text: str, task: dict) -> float:
            # Try the exact param name, then fall back to "answer"
            answer_val = task.get(p1_name, task.get("answer", ""))
            return float(fn(completion_text, answer_val))

        return _wrap_pair

    if (
        len(non_default_params) == 2
        and has_var_keyword
        and non_default_params[0].name in _COMPLETION_NAMES
        and non_default_params[1].name in _ANSWER_NAMES
    ):

        def _wrap_pair_kwargs(completion_text: str, task: dict) -> float:
            p1_name = non_default_params[1].name
            answer_val = task.get(p1_name, task.get("answer", ""))
            extra = {k: v for k, v in task.items() if k != p1_name}
            return float(fn(completion_text, answer_val, **extra))

        return _wrap_pair_kwargs

    # Has **kwargs -- pass completion as first arg, rest from task
    if has_var_keyword:
        positional_names = [
            p.name
            for p in params
            if p.kind
            in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD)
        ]

        def _wrap_kwargs(completion_text: str, task: dict) -> float:
            kwargs: dict[str, Any] = {}
            args: list[Any] = []
            for i, name in enumerate(positional_names):
                if i == 0:
                    args.append(completion_text)
                else:
                    args.append(task.get(name, ""))
            # Forward remaining task keys as kwargs
            forwarded = set(positional_names)
            for k, v in task.items():
                if k not in forwarded:
                    kwargs[k] = v
            return float(fn(*args, **kwargs))

        return _wrap_kwargs

    # General case: first param gets completion_text, remaining named params
    # are looked up from the task dict.
    def _wrap_general(completion_text: str, task: dict) -> float:
        bound: dict[str, Any] = {}
        for i, p in enumerate(params):
            if p.kind == inspect.Parameter.VAR_KEYWORD:
                continue
            if i == 0:
                bound[p.name] = completion_text
            else:
                default = p.default if p.default is not inspect.Parameter.empty else ""
                bound[p.name] = task.get(p.name, default)
        return float(fn(**bound))

    return _wrap_general


class RawAdapter(BaseAdapter):
    """Wrap one or more plain Python callables as reward functions.

    Parameters
    ----------
    fn:
        A callable (or list of callables) to use as the reward function(s).
    tasks:
        A list of task dicts.  Each dict should at minimum contain an
        ``"answer"`` key (ground truth).
    name:
        Human-readable name for this environment.  Falls back to *fn*'s
        ``__name__``.
    weights:
        Optional per-function weights (defaults to 1.0 each).
    """

    def __init__(
        self,
        fn: Callable | list[Callable],
        tasks: list[dict[str, Any]],
        name: Optional[str] = None,
        weights: Optional[list[float]] = None,
    ) -> None:
        fns = fn if isinstance(fn, list) else [fn]
        self._fns = fns
        self._tasks = tasks
        self._name = name or fns[0].__name__
        self._weights = weights or [1.0] * len(fns)

    def adapt(self) -> NormalizedEnvironment:
        reward_fns: list[NormalizedRewardFn] = []
        for func, weight in zip(self._fns, self._weights):
            wrapped = _wrap_callable(func)
            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,
        )
