"""VerifiersAdapter -- load a verifiers environment file and normalize it."""

from __future__ import annotations

import asyncio
import importlib.util
import inspect
import sys
from pathlib import Path
from typing import Any, Callable, Optional

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


def _get_or_create_event_loop() -> asyncio.AbstractEventLoop:
    """Return the running loop or create a new one."""
    try:
        loop = asyncio.get_running_loop()
    except RuntimeError:
        loop = None

    if loop is not None and loop.is_running():
        # We are inside an already-running loop (e.g. Jupyter).
        # Create a new loop in a thread instead -- handled by the caller.
        return loop

    try:
        return asyncio.get_event_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        return loop


def _run_sync(coro: Any) -> Any:
    """Run an async coroutine synchronously, even from inside a running loop."""
    try:
        loop = asyncio.get_running_loop()
    except RuntimeError:
        loop = None

    if loop is not None and loop.is_running():
        # Already inside an event loop (e.g. Jupyter / pytest-asyncio).
        import concurrent.futures

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

    return asyncio.run(coro)


def _wrap_verifiers_fn(
    fn: Callable,
    parser: Optional[Callable] = None,
    class_objects: Optional[dict[str, Any]] = None,
) -> Callable[[str, dict], float]:
    """Wrap a verifiers reward fn to sync ``(completion_text, task) -> float``.

    Verifiers reward functions use injectable params:
    ``completion`` (list[dict] messages), ``answer`` (str),
    ``info`` (dict), ``state`` (dict), ``prompt`` (list[dict]),
    ``parser`` (Parser object), plus any rubric class_objects.
    """
    sig = inspect.signature(fn)
    param_names = set(sig.parameters.keys())
    _class_objects = class_objects or {}

    def _call_sync(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"] = {}

        # Inject parser if the fn expects it
        if "parser" in param_names and parser is not None:
            kwargs["parser"] = parser

        # Inject any rubric class_objects (e.g. judge, tools, etc.)
        for name, obj in _class_objects.items():
            if name in param_names and name not in kwargs:
                kwargs[name] = obj

        result = fn(**kwargs)

        if inspect.isawaitable(result):
            result = _run_sync(result)

        return float(result)

    return _call_sync


def _load_module_from_path(file_path: str) -> Any:
    """Dynamically import a Python module from a file path."""
    path = Path(file_path).resolve()
    if not path.exists():
        raise FileNotFoundError(f"Environment file not found: {path}")

    module_name = f"_rewardprobe_env_{path.stem}"
    spec = importlib.util.spec_from_file_location(module_name, str(path))
    if spec is None or spec.loader is None:
        raise ImportError(f"Cannot create module spec from {path}")

    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module
    spec.loader.exec_module(module)
    return module


def _resolve_dataset(dataset: Any) -> list[dict[str, Any]]:
    """Convert a verifiers dataset (HF Dataset, callable, or list) to a list of dicts."""
    # If it's a callable (lazy loader), call it first.
    if callable(dataset) and not hasattr(dataset, "__getitem__"):
        dataset = dataset()

    # HuggingFace Dataset
    try:
        from datasets import Dataset

        if isinstance(dataset, Dataset):
            return [dict(row) for row in dataset]
    except ImportError:
        pass

    # Already a list of dicts
    if isinstance(dataset, list):
        return dataset

    # Iterable fallback
    try:
        return [dict(row) for row in dataset]
    except (TypeError, ValueError):
        raise TypeError(
            f"Cannot convert dataset of type {type(dataset).__name__} to list[dict]. "
            "Expected a HuggingFace Dataset, list of dicts, or callable returning one."
        )


class VerifiersAdapter(BaseAdapter):
    """Load a verifiers environment file and normalize it.

    Parameters
    ----------
    file_path:
        Path to a verifiers environment Python file (must expose a
        ``load_environment()`` function).
    name:
        Override the environment name.  Defaults to the file stem.
    max_tasks:
        Cap the number of tasks loaded from the dataset.  ``None`` means
        load all.
    """

    def __init__(
        self,
        file_path: str,
        name: Optional[str] = None,
        max_tasks: Optional[int] = None,
    ) -> None:
        try:
            import verifiers  # noqa: F401
        except ImportError:
            raise ImportError(
                "The 'verifiers' package is required for VerifiersAdapter. "
                "Install it with: pip install verifiers"
            )

        self._file_path = file_path
        self._name = name or Path(file_path).stem
        self._max_tasks = max_tasks

    def adapt(self) -> NormalizedEnvironment:
        module = _load_module_from_path(self._file_path)

        if not hasattr(module, "load_environment"):
            raise AttributeError(
                f"Environment file {self._file_path} does not expose a "
                "'load_environment()' function."
            )

        env = module.load_environment()

        # --- Reward functions from the rubric ---
        rubric = env.rubric

        # verifiers wraps user rubrics in a RubricGroup — recurse into
        # sub-rubrics to extract all reward functions.
        raw_fns: list = []
        weights: list[float] = []
        all_class_objects: dict = {}

        sub_rubrics = getattr(rubric, "rubrics", None)
        if sub_rubrics:
            for sub in sub_rubrics:
                sub_fns = list(getattr(sub, "funcs", []))
                sub_weights = list(getattr(sub, "weights", [1.0] * len(sub_fns)))
                raw_fns.extend(sub_fns)
                weights.extend(sub_weights)
                co = getattr(sub, "class_objects", {})
                if co:
                    all_class_objects.update(co)
        else:
            raw_fns = list(getattr(rubric, "funcs", []))
            weights = list(getattr(rubric, "weights", [1.0] * len(raw_fns)))

        parser = getattr(rubric, "parser", None)
        co_top = getattr(rubric, "class_objects", {}) or {}
        all_class_objects.update(co_top)

        reward_fns: list[NormalizedRewardFn] = []
        for i, (fn, weight) in enumerate(zip(raw_fns, weights)):
            fn_name = getattr(fn, "__name__", f"reward_{i}")

            # Skip zero-weight internal metrics (e.g. num_turns) —
            # these are monitoring metrics, not reward functions, and
            # probing them generates false positives.
            if weight == 0.0:
                continue

            wrapped = _wrap_verifiers_fn(fn, parser=parser, class_objects=all_class_objects)
            reward_fns.append(NormalizedRewardFn(name=fn_name, fn=wrapped, weight=weight))

        # --- Dataset / tasks ---
        # verifiers envs use .dataset or .eval_dataset
        raw_dataset = getattr(env, "dataset", None)
        if raw_dataset is None:
            raw_dataset = getattr(env, "eval_dataset", [])
        if raw_dataset is None:
            raw_dataset = []
        tasks = _resolve_dataset(raw_dataset)
        if self._max_tasks is not None:
            tasks = tasks[: self._max_tasks]

        # Detect environment type and tool capabilities for stateful attacks
        env_type = type(env).__name__
        has_tools = hasattr(env, "call_tool") or hasattr(env, "tool_defs")

        return NormalizedEnvironment(
            name=self._name,
            reward_fns=reward_fns,
            tasks=tasks,
            parser=parser,
            metadata={
                "source": "verifiers",
                "file": str(Path(self._file_path).resolve()),
                "env_type": env_type,
                "has_tools": has_tools,
            },
        )
