from typing import Any, Callable, Hashable, Iterator

from pydantic import BaseModel


class ComputedValue(BaseModel):
    """Return type for metric functions that need to attach debug info."""

    model_config = {"frozen": True}

    value: float
    debug: dict[str, Any] = {}


class MetricDefinition(BaseModel):
    """Definition of a metric before computation."""

    key: str
    name: str = ""
    fn: Callable[..., float | ComputedValue]
    unit: str = ""
    description: str = ""
    higher_is_better: bool = True
    normalize: Callable[[float], float] | None = None
    metadata: Any = {}
    help: dict[str, str] = {}

    model_config = {"arbitrary_types_allowed": True}


class ConstraintDefinition(BaseModel):
    """Definition of a constraint before computation."""

    key: str
    name: str = ""
    fn: Callable[..., float | ComputedValue]
    threshold: float
    unit: str = ""
    description: str = ""
    higher_is_better: bool = False
    metadata: Any = {}
    help: dict[str, str] = {}

    model_config = {"arbitrary_types_allowed": True}


class Metric(BaseModel):
    """A computed measurement. Created by MetricSet.compute(), not by user code."""

    key: str
    name: str = ""
    value: float
    unit: str = ""
    description: str = ""
    higher_is_better: bool = True
    normalize: Callable[[float], float] | None = None
    debug: dict[str, Any] = {}
    metadata: Any = {}
    help: dict[str, str] = {}

    model_config = {"frozen": True, "arbitrary_types_allowed": True}


class Constraint(BaseModel):
    """A computed check with explicit pass/fail. Created by MetricSet.compute(), not by user code."""

    model_config = {"frozen": True}

    key: str
    name: str = ""
    value: float
    threshold: float
    passed: bool
    unit: str = ""
    description: str = ""
    higher_is_better: bool = False
    debug: dict[str, Any] = {}
    metadata: Any = {}
    help: dict[str, str] = {}


class EvalResult(BaseModel):
    """Output of MetricSet.compute(). Holds all metrics and constraints."""

    model_config = {"frozen": True}

    metrics: dict[str, Metric]
    constraints: dict[str, Constraint]

    @property
    def violations(self) -> list[Constraint]:
        """Constraints that did not pass."""
        return [c for c in self.constraints.values() if not c.passed]

    @property
    def constraints_ok(self) -> bool:
        """True if all constraints passed."""
        return len(self.violations) == 0

    def assert_constraints(self) -> None:
        """Raise ValueError if any constraint failed."""
        if not self.constraints_ok:
            messages = [
                f"{c.key}: {c.value} > {c.threshold}" for c in self.violations
            ]
            raise ValueError("Constraint violations: " + ", ".join(messages))

    def __len__(self) -> int:
        return len(self.metrics) + len(self.constraints)

    def __iter__(self) -> Iterator[str]:
        yield from self.metrics
        yield from self.constraints

    def __getitem__(self, key: str) -> Metric | Constraint:
        """Get a Metric or Constraint by key. Raises KeyError if not found."""
        if key in self.metrics:
            return self.metrics[key]
        if key in self.constraints:
            return self.constraints[key]
        raise KeyError(key)

    def score(self, key: str) -> float:
        """Get value of a metric or constraint by key. Raises KeyError if not found."""
        if key in self.metrics:
            return self.metrics[key].value
        if key in self.constraints:
            return self.constraints[key].value
        raise KeyError(key)


class BatchResult:
    """Output of MetricSet.compute_batch(). Maps keys to EvalResults."""

    def __init__(self, results: dict[Hashable, EvalResult], key_names: list[str] | None = None) -> None:
        self._results = results
        self.key_names = key_names

    def __getitem__(self, key: Hashable) -> EvalResult:
        return self._results[key]

    def __iter__(self) -> Iterator[Hashable]:
        return iter(self._results)

    def __len__(self) -> int:
        return len(self._results)

    def items(self) -> Iterator[tuple[Hashable, EvalResult]]:
        """Iterate over (key, EvalResult) pairs."""
        return self._results.items()

    def scores(self, metric_key: str) -> dict[Hashable, float]:
        """Return {key: float} for a given metric across all batch entries."""
        return {key: result.score(metric_key) for key, result in self._results.items()}
