"""Wrapper to integrate our toolset evaluators with pydantic-evals Dataset API."""

from __future__ import annotations

import asyncio
from typing import Any, Callable

from pydantic_evals import Case, Dataset
from pydantic_evals.evaluators import Evaluator

from pydantic_ai_toolsets.evals.base import AgentRunner, StorageInspector
from pydantic_ai_toolsets.evals.config import EvaluationConfig
from pydantic_ai_toolsets.evals.evaluators import (
    EfficiencyEvaluator,
    OutputQualityEvaluator,
    StorageStateEvaluator,
    ToolUsageEvaluator,
)


class ToolsetTaskWrapper:
    """Wrapper to create a task function for pydantic-evals from a toolset evaluator."""

    def __init__(
        self,
        toolset_name: str,
        toolset_factory: Callable[[], Any],
        system_prompt: str,
        runner: AgentRunner,
    ):
        """Initialize the wrapper.

        Args:
            toolset_name: Name of the toolset.
            toolset_factory: Function that creates the toolset (closure with storage).
            system_prompt: System prompt for the agent.
            runner: AgentRunner instance.
        """
        self.toolset_name = toolset_name
        self.toolset_factory = toolset_factory
        self.system_prompt = system_prompt
        self.runner = runner

    async def __call__(self, inputs: dict[str, Any]) -> str:
        """Run the agent task.

        Args:
            inputs: Dictionary with 'prompt' key containing the user prompt.

        Returns:
            Agent output as string.
        """
        prompt = inputs.get("prompt", "")
        if not prompt:
            return ""

        # Create storage (the factory already captures it, but we need it for inspection)
        # Actually, the factory is a closure, so storage is already captured
        # We just need to run the agent
        run_result = await self.runner.run_agent(
            toolset_factory=self.toolset_factory,
            prompt=prompt,
            storage=None,  # Storage is captured in the factory closure
            system_prompt=self.system_prompt,
        )

        if run_result.error:
            raise RuntimeError(f"Agent execution failed: {run_result.error}")

        return run_result.output


def create_custom_evaluators(
    case: Any,
    output_evaluator: OutputQualityEvaluator,
    tool_evaluator: ToolUsageEvaluator,
    storage_evaluator: StorageStateEvaluator,
    efficiency_evaluator: EfficiencyEvaluator,
) -> list[Evaluator]:
    """Create custom evaluators for a case.

    Args:
        case: Test case with metadata.
        output_evaluator: Output quality evaluator.
        tool_evaluator: Tool usage evaluator.
        storage_evaluator: Storage state evaluator.
        efficiency_evaluator: Efficiency evaluator.

    Returns:
        List of evaluator instances.
    """
    evaluators = []

    # Create a custom evaluator that uses our existing evaluators
    class CombinedEvaluator(Evaluator):
        def evaluate(self, ctx) -> float:
            # This will be called by pydantic-evals
            # We need to extract the result from the context
            # For now, return a placeholder - we'll implement this properly
            return 1.0

    # For now, return empty list - we'll implement custom evaluators properly
    # The main goal is to get experiments showing up in Logfire first
    return evaluators


def create_dataset_from_cases(
    cases: list[Any],
    toolset_name: str,
    toolset_factory: Callable[[], Any],
    system_prompt: str,
    runner: AgentRunner,
    config: EvaluationConfig,
) -> Dataset:
    """Create a pydantic-evals Dataset from test cases.

    Args:
        cases: List of test case objects.
        toolset_name: Name of the toolset.
        toolset_factory: Function that creates the toolset.
        system_prompt: System prompt for the agent.
        runner: AgentRunner instance.
        config: Evaluation configuration.

    Returns:
        Dataset instance ready to evaluate.
    """
    # Create task wrapper
    task_wrapper = ToolsetTaskWrapper(
        toolset_name=toolset_name,
        toolset_factory=toolset_factory,
        system_prompt=system_prompt,
        runner=runner,
    )

    # Convert test cases to pydantic-evals Cases
    pydantic_cases = []
    for case in cases:
        pydantic_cases.append(
            Case(
                name=case.name,
                inputs={"prompt": case.prompt},
                metadata={
                    "difficulty": case.difficulty,
                    "expected_tools": case.expected_tools or [],
                    "expected_storage_keys": case.expected_storage_keys or [],
                    "min_storage_items": case.min_storage_items,
                },
            )
        )

    # Create dataset
    dataset = Dataset(
        name=f"{toolset_name}_evaluation",
        cases=pydantic_cases,
        # We can add evaluators here later
        evaluators=[],
    )

    return dataset, task_wrapper

