import inspect
import re
import textwrap
from typing import Any, Callable, Generic, TypeVar

from dot_metrics.models import (
    BatchResult,
    ComputedValue,
    Constraint,
    ConstraintDefinition,
    EvalResult,
    Metric,
    MetricDefinition,
)
from pydantic import BaseModel

T = TypeVar("T")
M = TypeVar("M")


def _parse_docstring(fn) -> dict[str, str]:
    """Parse a docstring into a structured help dict.

    Format: First paragraph is summary. Named sections (Range:, Interpretation:,
    Notes:) are captured as free text.
    """
    doc = inspect.getdoc(fn)
    if not doc:
        return {}

    known_sections = {"range", "interpretation", "notes"}
    result: dict[str, str] = {}
    current_section: str | None = None
    summary_lines: list[str] = []
    section_lines: list[str] = []
    past_summary = False

    for line in doc.splitlines():
        m = re.match(r"^(\w+):\s*(.*)", line)
        if m and m.group(1).lower() in known_sections:
            if current_section is not None:
                result[current_section] = textwrap.dedent("\n".join(section_lines)).strip()
            elif summary_lines:
                result["summary"] = " ".join(summary_lines).strip()
                past_summary = True
            current_section = m.group(1).lower()
            section_lines = [m.group(2)] if m.group(2).strip() else []
        elif current_section is not None:
            section_lines.append(line)
        elif not past_summary:
            if line.strip():
                summary_lines.append(line.strip())
            elif summary_lines:
                past_summary = True

    if current_section is not None:
        result[current_section] = textwrap.dedent("\n".join(section_lines)).strip()
    elif summary_lines:
        result["summary"] = " ".join(summary_lines).strip()

    return result


class MetricSet(Generic[T, M]):
    """A set of metrics and constraints evaluated against a single input type.

    ``T`` is the type of the data object passed to :meth:`compute` and
    :meth:`compute_batch`.  ``M`` is the type of the metadata attached to
    each metric and constraint definition.

    Annotate your set with the concrete type so static type checkers can
    verify that all registered functions are compatible::

        @dataclass
        class Prediction:
            expected: float
            actual: float

        metric_set: MetricSet[Prediction] = MetricSet()
        metric_set.add("error", lambda p: abs(p.expected - p.actual))

        result = metric_set.compute(Prediction(expected=1.0, actual=0.9))

    For typed metadata, supply a second type parameter::

        class MyMeta(TypedDict):
            category: str
            priority: int

        ms: MetricSet[dict, MyMeta] = MetricSet()
        ms.add("score", lambda d: 1.0, metadata=MyMeta(category="perf", priority=1))

    Any type works for ``T`` — a dataclass, a plain dict, a string, or a
    primitive.  If your metrics need several inputs, bundle them into a single
    object (dataclass, TypedDict, namedtuple, …) rather than passing multiple
    arguments.
    """

    def __init__(self) -> None:
        self.metrics: dict[str, MetricDefinition] = {}
        self.constraints: dict[str, ConstraintDefinition] = {}

    def add(
        self,
        key: str,
        fn: Callable[[T], float | ComputedValue],
        *,
        name: str = "",
        unit: str = "",
        description: str = "",
        higher_is_better: bool = True,
        normalize: Callable[[float], float] | None = None,
        metadata: M | None = None,
        help: dict[str, str] | None = None,
    ) -> None:
        """Register a metric function. Raises ValueError if key already exists."""
        if key in self.metrics:
            raise ValueError(f"Metric '{key}' is already registered")
        if metadata is None:
            metadata = {}
        if help is None:
            help = {}
        self.metrics[key] = MetricDefinition(
            key=key,
            name=name,
            fn=fn,
            unit=unit,
            description=description,
            higher_is_better=higher_is_better,
            normalize=normalize,
            metadata=metadata,
            help=help,
        )

    def add_constraint(
        self,
        key: str,
        fn: Callable[[T], float | ComputedValue],
        *,
        name: str = "",
        threshold: float,
        unit: str = "",
        description: str = "",
        higher_is_better: bool = False,
        metadata: M | None = None,
        help: dict[str, str] | None = None,
    ) -> None:
        """Register a constraint function. Raises ValueError if key already exists."""
        if key in self.constraints:
            raise ValueError(f"Constraint '{key}' is already registered")
        if metadata is None:
            metadata = {}
        if help is None:
            help = {}
        self.constraints[key] = ConstraintDefinition(
            key=key,
            name=name,
            fn=fn,
            threshold=threshold,
            unit=unit,
            description=description,
            higher_is_better=higher_is_better,
            metadata=metadata,
            help=help,
        )

    def metric(
        self,
        key: str,
        *,
        name: str = "",
        unit: str = "",
        description: str = "",
        higher_is_better: bool = True,
        normalize: Callable[[float], float] | None = None,
        metadata: M | None = None,
    ) -> Callable[[Callable[[T], float | ComputedValue]], Callable[[T], float | ComputedValue]]:
        """Decorator to register a metric function.

        Usage::

            metric_set = MetricSet()

            @metric_set.metric("coverage")
            def coverage(data):
                return len(data["covered"]) / len(data["total"])

        The decorated function is returned unchanged and can still be called directly.
        """

        def decorator(
            fn: Callable[[T], float | ComputedValue],
        ) -> Callable[[T], float | ComputedValue]:
            self.add(
                key,
                fn,
                name=name,
                unit=unit,
                description=description,
                higher_is_better=higher_is_better,
                normalize=normalize,
                metadata=metadata,
                help=_parse_docstring(fn),
            )
            return fn

        return decorator

    def constraint(
        self,
        key: str,
        *,
        name: str = "",
        threshold: float,
        unit: str = "",
        description: str = "",
        higher_is_better: bool = False,
        metadata: M | None = None,
    ) -> Callable[[Callable[[T], float | ComputedValue]], Callable[[T], float | ComputedValue]]:
        """Decorator to register a constraint function.

        Usage::

            metric_set = MetricSet()

            @metric_set.constraint("violations", threshold=0)
            def violations(data):
                return count_issues(data)

        The decorated function is returned unchanged and can still be called directly.
        """

        def decorator(
            fn: Callable[[T], float | ComputedValue],
        ) -> Callable[[T], float | ComputedValue]:
            self.add_constraint(
                key,
                fn,
                name=name,
                threshold=threshold,
                unit=unit,
                description=description,
                higher_is_better=higher_is_better,
                metadata=metadata,
                help=_parse_docstring(fn),
            )
            return fn

        return decorator

    def compute(self, data: T) -> EvalResult:
        """Evaluate all metrics and constraints against data. Returns an EvalResult."""
        metrics: dict[str, Metric] = {}
        for defn in self.metrics.values():
            raw = defn.fn(data)
            value, debug = _unpack(raw)
            metrics[defn.key] = Metric(
                key=defn.key,
                name=defn.name,
                value=value,
                unit=defn.unit,
                description=defn.description,
                higher_is_better=defn.higher_is_better,
                normalize=defn.normalize,
                debug=debug,
                metadata=defn.metadata,
                help=defn.help,
            )

        constraints: dict[str, Constraint] = {}
        for defn in self.constraints.values():
            raw = defn.fn(data)
            value, debug = _unpack(raw)
            if defn.higher_is_better:
                passed = value >= defn.threshold
            else:
                passed = value <= defn.threshold
            constraints[defn.key] = Constraint(
                key=defn.key,
                name=defn.name,
                value=value,
                threshold=defn.threshold,
                passed=passed,
                unit=defn.unit,
                description=defn.description,
                higher_is_better=defn.higher_is_better,
                debug=debug,
                metadata=defn.metadata,
                help=defn.help,
            )

        return EvalResult(metrics=metrics, constraints=constraints)

    def compute_batch(self, items: list[T] | dict[Any, T], *, key_names: list[str] | None = None) -> BatchResult:
        """Evaluate all metrics/constraints for each item. Returns a BatchResult."""
        if isinstance(items, dict):
            pairs = items.items()
        elif isinstance(items, list):
            pairs = enumerate(items)
        else:
            raise TypeError(f"compute_batch expects a list or dict, got {type(items).__name__}")
        return BatchResult({key: self.compute(data) for key, data in pairs}, key_names=key_names)



def _unpack(raw: float | ComputedValue) -> tuple[float, dict[str, Any]]:
    if isinstance(raw, ComputedValue):
        return raw.value, raw.debug
    return raw, {}
