"""Universal reward function detection and wrapping.

Inspects any callable's signature and wraps it to the normalized
``(completion_text: str, task: dict) -> float`` interface. Handles all
common patterns from TRL, verifiers, verl, and raw Python functions.
"""

from __future__ import annotations

import asyncio
import inspect
from enum import Enum
from typing import Any, Callable, Optional, get_args, get_origin

from rewardprobe.types import NormalizedEnvironment, NormalizedRewardFn


class RewardFnPattern(str, Enum):
    """Detected signature pattern of a reward function."""

    SINGLE_FLEXIBLE = "single_flexible"  # fn(completion) -> float
    SINGLE_PAIR = "single_pair"  # fn(completion, answer) -> float
    BATCH_SIMPLE = "batch_simple"  # fn(completions, **kwargs) -> list[float]
    BATCH_WITH_PROMPTS = "batch_with_prompts"  # fn(prompts, completions, **kw)
    VERIFIERS_ASYNC = "verifiers_async"  # async fn(completion, answer, ...)


# Known parameter names that indicate batch (list) input
_KNOWN_BATCH_NAMES = {
    "completions",
    "responses",
    "outputs",
    "predictions",
    "texts",
    "prompts",
    "solutions",
    "answers",
    "labels",
    "samples",
    "generations",
    "results",
    "inputs",
}


def _has_list_annotation(param: inspect.Parameter) -> bool:
    """Check if a parameter has a list-like type annotation."""
    ann = param.annotation
    if ann is inspect.Parameter.empty:
        return False
    origin = get_origin(ann)
    if origin is list:
        return True
    # Handle str annotations like "list[str]"
    if isinstance(ann, str) and ann.startswith("list"):
        return True
    return False


def _is_batch_param(param: inspect.Parameter) -> bool:
    """Heuristic: does this param look like it receives a batch (list)?"""
    if _has_list_annotation(param):
        return True
    if param.name in _KNOWN_BATCH_NAMES:
        return True
    return False


def classify_reward_fn(fn: Callable) -> RewardFnPattern:
    """Classify a callable into one of the known reward function patterns."""

    # 1. Async → verifiers style
    if asyncio.iscoroutinefunction(fn):
        return RewardFnPattern.VERIFIERS_ASYNC

    sig = inspect.signature(fn)
    params = [
        p
        for p in sig.parameters.values()
        if p.kind
        in (
            inspect.Parameter.POSITIONAL_ONLY,
            inspect.Parameter.POSITIONAL_OR_KEYWORD,
        )
    ]
    # Filter out self/cls for bound methods
    params = [p for p in params if p.name not in ("self", "cls")]
    required = [p for p in params if p.default is inspect.Parameter.empty]
    required = [p for p in required if p.name not in ("self", "cls")]

    # 2. Batch patterns — first required param looks like a list
    if required and _is_batch_param(required[0]):
        # Prompts-first pattern: fn(prompts, completions, ...)
        if len(required) >= 2 and required[0].name == "prompts" and _is_batch_param(required[1]):
            return RewardFnPattern.BATCH_WITH_PROMPTS
        return RewardFnPattern.BATCH_SIMPLE

    # 3. Single-param patterns
    if len(required) == 1:
        return RewardFnPattern.SINGLE_FLEXIBLE

    # 4. Two+ required params without batch hints → single pair
    return RewardFnPattern.SINGLE_PAIR


def wrap_reward_fn(
    fn: Callable,
    tasks: list[dict[str, Any]],
    pattern: Optional[RewardFnPattern] = None,
) -> Callable[[str, dict], float]:
    """Wrap any reward function to ``(completion_text: str, task: dict) -> float``.

    Auto-detects the signature pattern if not provided.
    """
    if pattern is None:
        pattern = classify_reward_fn(fn)

    if pattern == RewardFnPattern.SINGLE_FLEXIBLE:
        return _wrap_single_flexible(fn)
    if pattern == RewardFnPattern.SINGLE_PAIR:
        return _wrap_single_pair(fn, tasks)
    if pattern == RewardFnPattern.BATCH_SIMPLE:
        return _wrap_batch_simple(fn, tasks)
    if pattern == RewardFnPattern.BATCH_WITH_PROMPTS:
        return _wrap_batch_with_prompts(fn, tasks)
    if pattern == RewardFnPattern.VERIFIERS_ASYNC:
        return _wrap_verifiers_async(fn)

    # Fallback
    return _wrap_single_pair(fn, tasks)


# ---------------------------------------------------------------------------
# Pattern wrappers
# ---------------------------------------------------------------------------


def _wrap_single_flexible(fn: Callable) -> Callable[[str, dict], float]:
    """fn(completion) -> float"""

    def wrapper(completion_text: str, task: dict) -> float:
        result = fn(completion_text)
        if result is None:
            return 0.0
        if isinstance(result, (list, tuple)):
            result = result[0]
        return float(result)

    return wrapper


def _wrap_single_pair(fn: Callable, tasks: list[dict[str, Any]]) -> Callable[[str, dict], float]:
    """fn(completion, answer, ...) -> float — all non-first params from task dict."""
    sig = inspect.signature(fn)
    params = [
        p
        for p in sig.parameters.values()
        if p.kind
        in (
            inspect.Parameter.POSITIONAL_ONLY,
            inspect.Parameter.POSITIONAL_OR_KEYWORD,
        )
    ]

    # Figure out what keys the second+ params map to in the task dict.
    # Strategy: try the param name directly, then common aliases.
    _ANSWER_ALIASES = {
        "answer": ["answer", "solution", "ground_truth", "target", "label"],
        "ground_truth": ["answer", "ground_truth", "solution", "target"],
        "solution": ["answer", "solution", "ground_truth"],
        "target": ["answer", "target", "ground_truth"],
        "label": ["answer", "label", "ground_truth"],
        "expected": ["answer", "expected", "ground_truth"],
    }

    def wrapper(completion_text: str, task: dict) -> float:
        kwargs = {}
        for p in params[1:]:
            if p.name in task:
                kwargs[p.name] = task[p.name]
            elif p.default is inspect.Parameter.empty:
                # Required, not in task — try aliases
                aliases = _ANSWER_ALIASES.get(p.name, [p.name, "answer"])
                for alias in aliases:
                    if alias in task:
                        kwargs[p.name] = task[alias]
                        break
                else:
                    kwargs[p.name] = ""
            # If has default and not in task, don't pass — let fn use its default
        result = fn(completion_text, **kwargs)
        if result is None:
            return 0.0
        if isinstance(result, (list, tuple)):
            result = result[0]
        return float(result)

    return wrapper


def _wrap_batch_simple(fn: Callable, tasks: list[dict[str, Any]]) -> Callable[[str, dict], float]:
    """fn(completions, **kwargs) -> list[float] — TRL standard/conversational."""
    sig = inspect.signature(fn)
    params = list(sig.parameters.values())

    # Collect non-completion param names for kwargs forwarding
    _SKIP_PARAMS = {"completions", "completion_ids", "prompts", "trainer_state"}
    kwarg_params = [
        p.name
        for p in params
        if p.name not in _SKIP_PARAMS
        and p.kind
        in (
            inspect.Parameter.POSITIONAL_OR_KEYWORD,
            inspect.Parameter.KEYWORD_ONLY,
        )
        and p.name != params[0].name
    ]

    # Detect if this fn expects conversational format (list[list[dict]])
    # Strategy: check type annotation first, then try calling with both formats
    first_param = params[0] if params else None
    is_conversational = False
    if first_param:
        ann = first_param.annotation
        if ann is not inspect.Parameter.empty:
            origin = get_origin(ann)
            args = get_args(ann)
            if origin is list and args:
                inner = args[0]
                inner_origin = get_origin(inner)
                if inner_origin is list or inner_origin is dict:
                    is_conversational = True

    # Heuristic: if no annotation, probe the fn with both formats
    if not is_conversational and tasks:
        try:
            test_kwargs: dict[str, Any] = {}
            for pname in kwarg_params:
                if pname in tasks[0]:
                    test_kwargs[pname] = [tasks[0][pname]]
            has_var_kw = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params)
            if has_var_kw:
                for key, val in tasks[0].items():
                    if key not in test_kwargs and key not in _SKIP_PARAMS:
                        test_kwargs[key] = [val]
            # Try string format first
            try:
                fn(["test"], **test_kwargs)
            except (TypeError, IndexError, AttributeError, KeyError):
                # String format failed — try conversational
                try:
                    fn(
                        [[{"role": "assistant", "content": "test"}]],
                        **test_kwargs,
                    )
                    is_conversational = True
                except Exception:
                    pass
        except Exception:
            pass

    # Collect non-completion param names that might be kwargs from task
    _SKIP_PARAMS = {"completions", "completion_ids", "prompts", "trainer_state"}
    kwarg_params = [
        p.name
        for p in params
        if p.name not in _SKIP_PARAMS
        and p.kind
        in (
            inspect.Parameter.POSITIONAL_OR_KEYWORD,
            inspect.Parameter.KEYWORD_ONLY,
        )
        and p.name != params[0].name
    ]

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

        # Auto-forward task keys that match param names
        kwargs: dict[str, Any] = {}
        for pname in kwarg_params:
            if pname in task:
                kwargs[pname] = [task[pname]]

        # Also forward any task keys matching **kwargs params
        has_var_kw = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params)
        if has_var_kw:
            for key, val in task.items():
                if key not in kwargs and key not in _SKIP_PARAMS:
                    kwargs[key] = [val]

        results = fn(completions, **kwargs)
        result = results[0]
        return float(result) if result is not None else 0.0

    return wrapper


def _wrap_batch_with_prompts(
    fn: Callable, tasks: list[dict[str, Any]]
) -> Callable[[str, dict], float]:
    """fn(prompts, completions, **kwargs) -> list[float] — prompts-first TRL."""
    sig = inspect.signature(fn)
    params = list(sig.parameters.values())

    # Detect conversational from second param (completions)
    is_conversational = False
    if len(params) >= 2:
        ann = params[1].annotation
        if ann is not inspect.Parameter.empty:
            origin = get_origin(ann)
            args = get_args(ann)
            if origin is list and args:
                inner_origin = get_origin(args[0])
                if inner_origin is list or inner_origin is dict:
                    is_conversational = True

    # Probe if no annotation
    if not is_conversational and tasks:
        try:
            test_kwargs: dict[str, Any] = {}
            for p in params[2:]:
                if p.kind == inspect.Parameter.VAR_KEYWORD:
                    continue
                if p.name in tasks[0]:
                    test_kwargs[p.name] = [tasks[0][p.name]]
            # Try string format
            try:
                fn(["test_prompt"], ["test_completion"], **test_kwargs)
            except (TypeError, IndexError, AttributeError, KeyError):
                try:
                    fn(
                        [[{"role": "user", "content": "test"}]],
                        [[{"role": "assistant", "content": "test"}]],
                        **test_kwargs,
                    )
                    is_conversational = True
                except Exception:
                    pass
        except Exception:
            pass

    # Collect forwarded kwargs (skip prompts + completions)
    _SKIP = {params[0].name, params[1].name if len(params) > 1 else ""}
    _SKIP.update({"completion_ids", "trainer_state"})

    def wrapper(completion_text: str, task: dict) -> float:
        prompt_val = task.get("prompt", task.get("question", ""))

        if is_conversational:
            if isinstance(prompt_val, list):
                prompts: Any = [prompt_val]
            else:
                prompts = [[{"role": "user", "content": str(prompt_val)}]]
            completions: Any = [[{"role": "assistant", "content": completion_text}]]
        else:
            prompts = [str(prompt_val)]
            completions = [completion_text]

        kwargs: dict[str, Any] = {}
        for p in params[2:]:
            if p.kind == inspect.Parameter.VAR_KEYWORD:
                continue
            if p.name in task:
                kwargs[p.name] = [task[p.name]]

        has_var_kw = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params)
        if has_var_kw:
            for key, val in task.items():
                if key not in kwargs and key not in _SKIP:
                    kwargs[key] = [val]

        results = fn(prompts, completions, **kwargs)
        result = results[0]
        return float(result) if result is not None else 0.0

    return wrapper


def _wrap_verifiers_async(fn: Callable) -> Callable[[str, dict], float]:
    """async fn(completion, answer, parser, ...) -> float — verifiers style."""
    sig = inspect.signature(fn)
    param_names = set(sig.parameters.keys())

    def wrapper(completion_text: str, task: dict) -> float:
        kwargs: dict[str, Any] = {}

        if "completion" in param_names:
            kwargs["completion"] = [{"role": "assistant", "content": completion_text}]
        if "answer" in param_names:
            kwargs["answer"] = task.get("answer", "")
        if "prompt" in param_names:
            raw_prompt = task.get("prompt", task.get("question", ""))
            if isinstance(raw_prompt, list):
                kwargs["prompt"] = raw_prompt
            else:
                kwargs["prompt"] = [{"role": "user", "content": str(raw_prompt)}]
        if "info" in param_names:
            kwargs["info"] = task.get("info", task)
        if "state" in param_names:
            kwargs["state"] = {}

        result = fn(**kwargs)

        if inspect.isawaitable(result):
            try:
                loop = asyncio.get_running_loop()
            except RuntimeError:
                loop = None
            if loop and loop.is_running():
                import concurrent.futures

                with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
                    result = pool.submit(asyncio.run, result).result()
            else:
                result = asyncio.run(result)

        if result is None:
            return 0.0
        if isinstance(result, (list, tuple)):
            result = result[0]
        return float(result)

    return wrapper


# ---------------------------------------------------------------------------
# High-level: adapt any callable to NormalizedEnvironment
# ---------------------------------------------------------------------------


def is_answer_aware(
    fn: NormalizedRewardFn,
    tasks: list[dict[str, Any]],
    n_probe: int = 3,
) -> bool:
    """Check if a reward fn's output depends on the task's answer.

    If swapping the answer in the task dict doesn't change the score,
    the function is format-only (doesn't check correctness).
    """
    from rewardprobe.families.base import _safe_call

    for task in tasks[:n_probe]:
        if "answer" not in task:
            continue
        original_answer = task["answer"]

        # Score with real answer
        score_real, err1 = _safe_call(fn, f"test {original_answer}", task)
        if err1 is not None:
            continue

        # Score with fake answer
        fake_task = {**task, "answer": "DEFINITELY_WRONG_XYZ_99999"}
        score_fake, err2 = _safe_call(fn, f"test {original_answer}", fake_task)
        if err2 is not None:
            continue

        if score_real != score_fake:
            return True

    return False


def auto_adapt(
    fn: Callable | list[Callable],
    tasks: list[dict[str, Any]],
    name: Optional[str] = None,
    weights: list[float] | None = None,
) -> NormalizedEnvironment:
    """Adapt any callable(s) + tasks into a NormalizedEnvironment.

    Auto-detects the signature pattern and wraps accordingly.
    No framework-specific flags needed.
    """
    fns = fn if isinstance(fn, list) else [fn]

    reward_fns: list[NormalizedRewardFn] = []
    for i, f in enumerate(fns):
        pattern = classify_reward_fn(f)
        wrapped = wrap_reward_fn(f, tasks, pattern)
        fn_name = getattr(f, "__name__", "reward_fn")
        weight = weights[i] if weights and i < len(weights) else 1.0
        reward_fns.append(NormalizedRewardFn(name=fn_name, fn=wrapped, weight=weight))

    env_name = name or reward_fns[0].name
    return NormalizedEnvironment(
        name=env_name,
        reward_fns=reward_fns,
        tasks=tasks,
        metadata={"adapter": "auto"},
    )
