"""Base class for evaluation rules."""

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Optional

from agenteval.models import Case, Turn

if TYPE_CHECKING:
    from agenteval.model_gateway import ModelGateway
    from agenteval.services.model_configs import ModelRuntimeConfig


@dataclass
class RuleResult:
    """Result of applying an evaluation rule."""

    passed: bool
    score: Optional[float] = None
    reason: str = ""
    details: Optional[dict[str, Any]] = None


class EvalRule(ABC):
    """Abstract evaluation rule. All implementations must be async."""

    name: str = ""

    def __init__(
        self,
        params: dict[str, Any],
        model_config: "ModelRuntimeConfig | None" = None,
        gateway: "ModelGateway | None" = None,
    ):
        self.params = params
        self.model_config = model_config
        self.gateway = gateway
        # 本规则执行期间消耗的评测侧 token 用量（引擎按评测岗位归集）
        self.llm_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}

    def _record_llm_usage(self, usage: Optional[dict[str, int]]) -> None:
        if not usage:
            return
        for key in self.llm_usage:
            self.llm_usage[key] += int(usage.get(key) or 0)

    @abstractmethod
    async def evaluate(self, case: Case, dialog: list[Turn]) -> RuleResult:
        """Evaluate the dialog against the case expectation."""
        ...


_RULE_REGISTRY: dict[str, type[EvalRule]] = {}


def register_rule(rule_cls: type[EvalRule]) -> type[EvalRule]:
    _RULE_REGISTRY[rule_cls.name] = rule_cls
    return rule_cls


def get_rule(
    rule_type: str,
    params: dict[str, Any],
    model_config: "ModelRuntimeConfig | None" = None,
    gateway: "ModelGateway | None" = None,
) -> EvalRule:
    """Instantiate a rule by type name."""
    if rule_type not in _RULE_REGISTRY:
        raise ValueError(f"unknown rule type: {rule_type}. Available: {list(_RULE_REGISTRY.keys())}")
    return _RULE_REGISTRY[rule_type](params, model_config=model_config, gateway=gateway)


def list_rule_types() -> list[str]:
    return list(_RULE_REGISTRY.keys())