"""Base tracker protocol and framework-agnostic tool call tracker."""

from __future__ import annotations

import functools
import inspect
from abc import ABC, abstractmethod
from collections.abc import Callable
from typing import Any, TypeVar, overload

from evalkit.client import EvalKitClient, Evaluation
from evalkit.engine import EvaluationEngine
from evalkit.judges.base import BaseJudge
from evalkit.models import (
    EvaluationReport,
    Suite,
    TestCase,
    ToolCall,
)
from evalkit.suites.agent_tool_call import AgentToolCallSuite
from evalkit.suites.base import EvaluationSuite
from evalkit.suites.live_qa import LiveQASuite
from evalkit.suites.rag_qa import RAGQASuite

F = TypeVar("F", bound=Callable[..., Any])


class BaseTracker(ABC):
    """Protocol shared by all framework trackers.

    Provides raw data access (question, tool_calls, context, answer) and
    convenience methods for building test cases, running local evaluations,
    and uploading results to an evalkit server.
    """

    @property
    @abstractmethod
    def question(self) -> str:
        """The input question/prompt for this invocation."""
        ...

    @property
    @abstractmethod
    def tool_calls(self) -> list[ToolCall]:
        """Tool calls captured during this invocation."""
        ...

    @property
    @abstractmethod
    def context(self) -> list[str]:
        """RAG context chunks captured during this invocation."""
        ...

    @property
    @abstractmethod
    def answer(self) -> str:
        """Final answer produced by the agent/pipeline."""
        ...

    def to_test_case(
        self,
        expected_tool_calls: list[ToolCall] | None = None,
        expected_answer: str | None = None,
    ) -> TestCase:
        """Build a TestCase from captured data and optional expectations."""
        return TestCase(
            question=self.question,
            context=self.context,
            answer=self.answer,
            expected_answer=expected_answer,
            expected_tool_calls=expected_tool_calls,
            actual_tool_calls=self.tool_calls or None,
        )

    def evaluate(
        self,
        judge: BaseJudge,
        suite: EvaluationSuite | None = None,
        expected_tool_calls: list[ToolCall] | None = None,
        expected_answer: str | None = None,
    ) -> EvaluationReport:
        """Build TestCase, run evaluation locally, return report.

        If no suite is provided, auto-selects based on available data:
        - expected_tool_calls present -> AgentToolCallSuite
        - expected_answer present -> RAGQASuite
        - Neither -> LiveQASuite
        """
        tc = self.to_test_case(expected_tool_calls, expected_answer)
        if suite is None:
            suite = self._default_suite(tc)
        engine = EvaluationEngine(judge=judge)
        return engine.run(test_cases=[tc], suite=suite)

    def upload(
        self,
        project_name: str,
        server_url: str = "http://localhost:8000",
        suite: Suite = Suite.LIVE_QA,
        expected_tool_calls: list[ToolCall] | None = None,
        expected_answer: str | None = None,
        client: EvalKitClient | None = None,
    ) -> Evaluation:
        """Build TestCase, submit to evalkit server for evaluation.

        The server's configured judge handles evaluation. Pass a pre-configured
        ``EvalKitClient`` for batch workflows to avoid creating a new client on
        every call.
        """
        tc = self.to_test_case(expected_tool_calls, expected_answer)
        if client is None:
            client = EvalKitClient(server_url)
        project = client.create_project(project_name)
        return client.evaluate(project.id, [tc], suite=suite)

    @staticmethod
    def _default_suite(tc: TestCase) -> EvaluationSuite:
        """Pick suite based on what data is available."""
        if tc.expected_tool_calls:
            return AgentToolCallSuite()
        if tc.expected_answer:
            return RAGQASuite()
        return LiveQASuite()


class TrackerRun(BaseTracker):
    """Per-invocation capture returned by ``ToolCallTracker.capture()``.

    Implements ``BaseTracker`` so all convenience methods are available.
    """

    def __init__(self, question: str, calls: list[ToolCall]) -> None:
        self._question = question
        self._calls = calls
        self._context: list[str] = []
        self._answer: str = ""

    @property
    def question(self) -> str:
        return self._question

    @property
    def tool_calls(self) -> list[ToolCall]:
        return list(self._calls)

    @property
    def context(self) -> list[str]:
        return list(self._context)

    @context.setter
    def context(self, value: list[str]) -> None:
        self._context = value

    @property
    def answer(self) -> str:
        return self._answer

    @answer.setter
    def answer(self, value: str) -> None:
        self._answer = value


class ToolCallTracker:
    """Decorator-based tracker that captures tool calls made by any agent framework.

    Wrap your tool functions with :meth:`wrap` before registering them with your
    agent framework.  Every time the framework invokes a wrapped tool the call is
    recorded as an :class:`~evalkit.models.ToolCall`.

    For per-invocation scoping, use the :meth:`capture` context manager::

        tracker = ToolCallTracker()

        @tracker.wrap
        def get_weather(city: str) -> str:
            return f"Sunny in {city}"

        with tracker.capture("Weather in Paris?") as run:
            get_weather(city="Paris")
            run.answer = "Sunny in Paris"

        run.tool_calls  # [ToolCall(name='get_weather', parameters={'city': 'Paris'})]
        run.evaluate(judge=my_judge)
    """

    def __init__(self) -> None:
        self._calls: list[ToolCall] = []
        self._active_run: TrackerRun | None = None

    @property
    def tool_calls(self) -> list[ToolCall]:
        """All tool calls captured since the last :meth:`reset`."""
        return list(self._calls)

    def reset(self) -> None:
        """Clear all captured tool calls."""
        self._calls.clear()

    def capture(self, question: str) -> _CaptureContext:
        """Context manager that scopes tool call capture to a single invocation.

        Returns a :class:`TrackerRun` (via ``__enter__``) that holds the
        captured data and provides ``BaseTracker`` convenience methods.
        """
        if self._active_run is not None:
            raise RuntimeError(
                "ToolCallTracker.capture() does not support nesting. " "A capture() context is already active."
            )
        return _CaptureContext(self, question)

    @overload
    def wrap(self, fn: F) -> F: ...
    @overload
    def wrap(self, fn: None = None, *, name: str | None = None) -> Callable[[F], F]: ...

    def wrap(
        self,
        fn: F | None = None,
        *,
        name: str | None = None,
    ) -> F | Callable[[F], F]:
        """Decorate a tool function to record calls.

        Can be used directly (``@tracker.wrap``) or with options
        (``@tracker.wrap(name="custom_name")``).
        """
        if fn is None:
            return lambda f: self._apply(f, name)  # type: ignore[return-value]
        return self._apply(fn, name)  # type: ignore[return-value]

    def _apply(self, fn: F, override_name: str | None) -> F:
        tracker = self
        tool_name = override_name or fn.__name__
        sig = inspect.signature(fn)

        if inspect.iscoroutinefunction(fn):

            @functools.wraps(fn)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                params = _bind_params(sig, args, kwargs)
                tracker._calls.append(ToolCall(name=tool_name, parameters=params))
                return await fn(*args, **kwargs)

            return async_wrapper  # type: ignore[return-value]

        @functools.wraps(fn)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            params = _bind_params(sig, args, kwargs)
            tracker._calls.append(ToolCall(name=tool_name, parameters=params))
            return fn(*args, **kwargs)

        return wrapper  # type: ignore[return-value]


class _CaptureContext:
    """Context manager implementation for ``ToolCallTracker.capture()``."""

    def __init__(self, tracker: ToolCallTracker, question: str) -> None:
        self._tracker = tracker
        self._question = question
        self._saved_calls: list[ToolCall] | None = None
        self._run: TrackerRun | None = None

    def __enter__(self) -> TrackerRun:
        self._saved_calls = self._tracker._calls
        scoped_calls: list[ToolCall] = []
        self._tracker._calls = scoped_calls
        self._run = TrackerRun(self._question, scoped_calls)
        self._tracker._active_run = self._run
        return self._run

    def __exit__(self, *exc: Any) -> None:
        if self._saved_calls is not None:
            self._tracker._calls = self._saved_calls
        self._tracker._active_run = None


def _bind_params(
    sig: inspect.Signature,
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
) -> dict[str, Any]:
    """Bind call arguments to the function signature without applying defaults."""
    try:
        bound = sig.bind(*args, **kwargs)
        return dict(bound.arguments)
    except TypeError:
        return dict(kwargs) if kwargs else {}
