"""Function decorator for turning plain functions into eval checks."""

from __future__ import annotations

import inspect
from collections.abc import Callable
from typing import Any, overload

from agentbyte.eval.base import BaseEvalJudge
from agentbyte.eval.types import AnswerStrategy, EvalTrajectory

from .types import CheckResult, EvalCheck

_KNOWN_PARAMS = frozenset(
    {"task", "messages", "response", "expected_output", "query", "usage", "success"}
)


class _AnswerExtractorJudge(BaseEvalJudge):
    """Minimal judge wrapper used only for answer extraction."""

    async def score(
        self,
        trajectory: EvalTrajectory,
        criteria: list[str] | None = None,
        cancellation_token: Any = None,
    ) -> Any:
        raise NotImplementedError


def _coerce_result(value: Any, check_name: str) -> CheckResult:
    if isinstance(value, CheckResult):
        if value.check_name == check_name:
            return value
        return value.model_copy(update={"check_name": check_name})
    if isinstance(value, bool):
        return CheckResult(
            passed=value,
            reason="passed" if value else "failed",
            check_name=check_name,
        )
    if isinstance(value, (int, float)):
        score = float(value)
        return CheckResult(
            passed=score >= 0.5,
            reason=f"score={score:.3f}",
            check_name=check_name,
            score=score,
        )
    if isinstance(value, dict):
        if "score" in value:
            score = float(value["score"])
            return CheckResult(
                passed=bool(value.get("passed", score >= 0.5)),
                reason=str(value.get("reason", f"score={score:.3f}")),
                check_name=check_name,
                score=score,
            )
        if "passed" in value:
            passed = bool(value["passed"])
            return CheckResult(
                passed=passed,
                reason=str(value.get("reason", "passed" if passed else "failed")),
                check_name=check_name,
            )
    raise TypeError(
        f"Function evaluator '{check_name}' returned unsupported type "
        f"{type(value).__name__}. Expected bool, float, dict, or CheckResult."
    )


def _resolve_args(
    fn: Callable[..., Any],
    trajectory: EvalTrajectory,
    *,
    param_names: set[str] | frozenset[str],
    extractor: _AnswerExtractorJudge,
) -> dict[str, Any]:
    field_map: dict[str, Any] = {
        "task": trajectory.task,
        "messages": trajectory.messages,
        "response": extractor.extract_answer(trajectory),
        "expected_output": trajectory.task.expected_output,
        "query": trajectory.task.input,
        "usage": trajectory.usage,
        "success": trajectory.success,
    }
    return {name: field_map[name] for name in param_names if name in field_map}


@overload
def evaluator(fn: Callable[..., Any], /) -> EvalCheck: ...


@overload
def evaluator(
    *,
    answer_strategy: AnswerStrategy = "last_non_empty",
    name: str | None = None,
    source_filter: str | None = None,
) -> Callable[[Callable[..., Any]], EvalCheck]: ...


def evaluator(
    fn: Callable[..., Any] | None = None,
    *,
    answer_strategy: AnswerStrategy = "last_non_empty",
    name: str | None = None,
    source_filter: str | None = None,
) -> EvalCheck | Callable[[Callable[..., Any]], EvalCheck]:
    """Wrap a plain function as an EvalCheck with parameter injection.

    Pass ``source_filter`` to restrict the injected ``response`` value to messages
    from a specific agent (e.g. ``source_filter="writer"`` in a multi-agent pipeline).
    """

    def _wrap(func: Callable[..., Any]) -> EvalCheck:
        check_name = name or getattr(func, "__name__", None) or "evaluator"
        signature = inspect.signature(func)
        required_unknown = {
            param_name
            for param_name, param in signature.parameters.items()
            if param_name not in _KNOWN_PARAMS
            and param.default is inspect.Parameter.empty
        }
        if required_unknown:
            raise TypeError(
                f"Function evaluator '{func.__name__}' has unknown required "
                f"parameter(s) {sorted(required_unknown)}. Supported: {sorted(_KNOWN_PARAMS)}"
            )
        param_names = {
            param_name
            for param_name in signature.parameters
            if param_name in _KNOWN_PARAMS
        }

        async def _check(trajectory: EvalTrajectory) -> CheckResult:
            _extractor = _AnswerExtractorJudge(
                name="extractor",
                answer_strategy=answer_strategy,
                source_filter=source_filter,
            )
            kwargs = _resolve_args(
                func,
                trajectory,
                param_names=param_names,
                extractor=_extractor,
            )
            result = func(**kwargs)
            if inspect.isawaitable(result):
                result = await result
            return _coerce_result(result, check_name)

        _check.__name__ = check_name  # type: ignore[attr-defined,assignment]
        _check.__doc__ = func.__doc__
        return _check

    if fn is not None:
        return _wrap(fn)
    return _wrap


__all__ = ["evaluator"]
