"""Efficiency evaluator with toolset metrics integration."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from pydantic_ai_toolsets.evals.base import EvaluationResult

if TYPE_CHECKING:
    try:
        from _shared.metrics import UsageMetrics
    except ImportError:
        UsageMetrics = Any


class EfficiencyEvaluator:
    """Evaluates efficiency: tokens, time, cost.

    Can integrate with the new UsageMetrics from toolset storage for
    more accurate toolset-specific token tracking.
    """

    def evaluate_token_usage(
        self, result: EvaluationResult, max_tokens: int | None = None
    ) -> dict[str, Any]:
        """Evaluate token usage.

        Args:
            result: Evaluation result.
            max_tokens: Maximum acceptable tokens.

        Returns:
            Dictionary with token metrics.
        """
        token_usage = result.token_usage or {}
        input_tokens = token_usage.get("input_tokens", 0)
        output_tokens = token_usage.get("output_tokens", 0)
        total_tokens = token_usage.get("total_tokens", input_tokens + output_tokens)

        efficiency_score = 1.0
        if max_tokens and total_tokens > max_tokens:
            efficiency_score = max_tokens / total_tokens

        return {
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "total_tokens": total_tokens,
            "efficiency_score": efficiency_score,
        }

    def evaluate_toolset_metrics(
        self, storage: Any, max_toolset_tokens: int | None = None
    ) -> dict[str, Any]:
        """Evaluate toolset-specific metrics from storage.

        Uses the new UsageMetrics tracking if available.

        Args:
            storage: Storage object from toolset (may have metrics property).
            max_toolset_tokens: Maximum acceptable toolset tokens.

        Returns:
            Dictionary with toolset metrics, or empty dict if not available.
        """
        if storage is None:
            return {}

        # Check if storage has metrics
        metrics = getattr(storage, "metrics", None)
        if metrics is None:
            return {}

        # Extract metrics data
        result: dict[str, Any] = {
            "toolset_input_tokens": metrics.total_input_tokens(),
            "toolset_output_tokens": metrics.total_output_tokens(),
            "toolset_total_tokens": metrics.total_tokens(),
            "toolset_invocation_count": len(metrics.invocations),
            "toolset_invocations_by_tool": metrics.invocation_count(),
            "toolset_total_duration_ms": metrics.total_duration_ms(),
        }

        # Calculate efficiency if max specified
        if max_toolset_tokens and result["toolset_total_tokens"] > 0:
            result["toolset_efficiency_score"] = min(
                1.0, max_toolset_tokens / result["toolset_total_tokens"]
            )
        else:
            result["toolset_efficiency_score"] = 1.0

        # Average tokens per invocation
        if result["toolset_invocation_count"] > 0:
            result["toolset_avg_tokens_per_call"] = (
                result["toolset_total_tokens"] / result["toolset_invocation_count"]
            )
        else:
            result["toolset_avg_tokens_per_call"] = 0.0

        return result

    def evaluate_execution_time(
        self, result: EvaluationResult, max_time: float | None = None
    ) -> dict[str, Any]:
        """Evaluate execution time.

        Args:
            result: Evaluation result.
            max_time: Maximum acceptable time in seconds.

        Returns:
            Dictionary with time metrics.
        """
        execution_time = result.execution_time

        efficiency_score = 1.0
        if max_time and execution_time > max_time:
            efficiency_score = max_time / execution_time

        return {
            "execution_time": execution_time,
            "efficiency_score": efficiency_score,
        }

    def estimate_cost(
        self,
        result: EvaluationResult,
        cost_per_1k_input_tokens: float = 0.0025,
        cost_per_1k_output_tokens: float = 0.01,
    ) -> dict[str, Any]:
        """Estimate cost based on token usage.

        Args:
            result: Evaluation result.
            cost_per_1k_input_tokens: Cost per 1000 input tokens.
            cost_per_1k_output_tokens: Cost per 1000 output tokens.

        Returns:
            Dictionary with cost metrics.
        """
        token_usage = result.token_usage or {}
        input_tokens = token_usage.get("input_tokens", 0)
        output_tokens = token_usage.get("output_tokens", 0)

        input_cost = (input_tokens / 1000) * cost_per_1k_input_tokens
        output_cost = (output_tokens / 1000) * cost_per_1k_output_tokens
        total_cost = input_cost + output_cost

        return {
            "estimated_input_cost": input_cost,
            "estimated_output_cost": output_cost,
            "estimated_total_cost": total_cost,
            "cost_per_1k_input_tokens": cost_per_1k_input_tokens,
            "cost_per_1k_output_tokens": cost_per_1k_output_tokens,
        }

    def evaluate_all(
        self,
        result: EvaluationResult,
        storage: Any | None = None,
        max_tokens: int | None = None,
        max_toolset_tokens: int | None = None,
        max_time: float | None = None,
        cost_per_1k_input_tokens: float = 0.0025,
        cost_per_1k_output_tokens: float = 0.01,
    ) -> dict[str, Any]:
        """Evaluate all efficiency metrics including toolset metrics.

        Args:
            result: Evaluation result.
            storage: Optional storage object for toolset metrics.
            max_tokens: Maximum acceptable LLM tokens.
            max_toolset_tokens: Maximum acceptable toolset tokens.
            max_time: Maximum acceptable time.
            cost_per_1k_input_tokens: Cost per 1000 input tokens.
            cost_per_1k_output_tokens: Cost per 1000 output tokens.

        Returns:
            Dictionary with all efficiency metrics.
        """
        token_metrics = self.evaluate_token_usage(result, max_tokens)
        time_metrics = self.evaluate_execution_time(result, max_time)
        cost_metrics = self.estimate_cost(
            result, cost_per_1k_input_tokens, cost_per_1k_output_tokens
        )

        # Add toolset metrics if storage is provided
        toolset_metrics = {}
        if storage is not None:
            toolset_metrics = self.evaluate_toolset_metrics(storage, max_toolset_tokens)

        return {
            **token_metrics,
            **time_metrics,
            **cost_metrics,
            **toolset_metrics,
        }


def format_metrics_report(metrics: dict[str, Any]) -> str:
    """Format metrics as a readable report.

    Args:
        metrics: Dictionary of metrics from evaluate_all.

    Returns:
        Formatted string report.
    """
    lines: list[str] = ["Efficiency Report", "=" * 40]

    # LLM tokens
    lines.append("")
    lines.append("LLM Token Usage:")
    lines.append(f"  Input:  {metrics.get('input_tokens', 0):,}")
    lines.append(f"  Output: {metrics.get('output_tokens', 0):,}")
    lines.append(f"  Total:  {metrics.get('total_tokens', 0):,}")

    # Toolset tokens (if available)
    if metrics.get("toolset_total_tokens"):
        lines.append("")
        lines.append("Toolset Token Usage:")
        lines.append(f"  Input:  {metrics.get('toolset_input_tokens', 0):,}")
        lines.append(f"  Output: {metrics.get('toolset_output_tokens', 0):,}")
        lines.append(f"  Total:  {metrics.get('toolset_total_tokens', 0):,}")
        lines.append(f"  Calls:  {metrics.get('toolset_invocation_count', 0)}")
        lines.append(f"  Avg/call: {metrics.get('toolset_avg_tokens_per_call', 0):.1f}")

        by_tool = metrics.get("toolset_invocations_by_tool", {})
        if by_tool:
            lines.append("  By tool:")
            for tool, count in sorted(by_tool.items()):
                lines.append(f"    {tool}: {count}")

    # Time
    lines.append("")
    lines.append(f"Execution Time: {metrics.get('execution_time', 0):.2f}s")
    if metrics.get("toolset_total_duration_ms"):
        lines.append(f"Toolset Duration: {metrics.get('toolset_total_duration_ms', 0):.1f}ms")

    # Cost
    lines.append("")
    lines.append("Estimated Cost:")
    lines.append(f"  Input:  ${metrics.get('estimated_input_cost', 0):.6f}")
    lines.append(f"  Output: ${metrics.get('estimated_output_cost', 0):.6f}")
    lines.append(f"  Total:  ${metrics.get('estimated_total_cost', 0):.6f}")

    # Efficiency scores
    lines.append("")
    lines.append("Efficiency Scores:")
    lines.append(f"  Token: {metrics.get('efficiency_score', 1.0):.1%}")
    if metrics.get("toolset_efficiency_score") is not None:
        lines.append(f"  Toolset: {metrics.get('toolset_efficiency_score', 1.0):.1%}")

    return "\n".join(lines)
