"""High-level evaluation orchestration primitives."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

from .evaluator import evaluate_available_splits
from .exceptions import InferenceRunnerError, ModelLoadError
from .schemas import MetricPayload, validate_payload
from .status_handler import EvaluationStatusReporter


if TYPE_CHECKING:
    from collections.abc import Callable, Mapping, Sequence

    from .contracts import SplitEvaluatorProtocol
    from .runtime_registry import RuntimeTaskRegistry


@dataclass
class EvalContext:
    """Runtime context consumed by orchestrators."""

    task_type: str
    runtime_framework: str
    requested_splits: Sequence[str]
    split_to_loader: Mapping[str, object]
    index_to_labels: Mapping[str, str] = field(default_factory=dict)
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class EvaluationOrchestrator:
    """Execute split evaluation and payload normalization."""

    action_tracker: Any
    context: EvalContext
    evaluate_split: SplitEvaluatorProtocol
    validate_output_payload: bool = True

    def run(self) -> MetricPayload:
        """Execute evaluation over available splits and normalize payload.

        Args:
            None.

        Returns:
            Normalized metric payload aggregated across requested splits.
        """
        reporter = EvaluationStatusReporter(self.action_tracker)
        reporter.acknowledged()
        reporter.dataset_loaded()
        reporter.started()

        payload = evaluate_available_splits(
            requested_splits=self.context.requested_splits,
            split_to_loader=self.context.split_to_loader,
            evaluate_split=self.evaluate_split,
        )

        if self.validate_output_payload:
            payload = validate_payload(payload)
        return payload

    def persist(self, payload: MetricPayload) -> None:
        """Persist successful payload and emit completion status.

        Args:
            payload: Normalized metric payload to persist through tracker.

        Returns:
            None.
        """
        if hasattr(self.action_tracker, "save_evaluation_results"):
            self.action_tracker.save_evaluation_results(payload)
        reporter = EvaluationStatusReporter(self.action_tracker)
        reporter.completed()

    def fail(self, error: Exception) -> None:
        """Report evaluation failure details to tracking facilities.

        Args:
            error: Exception raised during evaluation flow.

        Returns:
            None.
        """
        reporter = EvaluationStatusReporter(self.action_tracker)
        reporter.failed(f"Error in evaluation: {error}")
        if hasattr(self.action_tracker, "log_error"):
            self.action_tracker.log_error(__file__, "eval/orchestrator", str(error))


def build_runtime_split_evaluator(
    registry: RuntimeTaskRegistry,
    *,
    action_tracker: Any,
    context: EvalContext,
    runner_kwargs: Mapping[str, Any] | None = None,
    split_formatter: Callable[[str, Any], MetricPayload] | None = None,
) -> Callable[[str, object], MetricPayload]:
    """Create split evaluator using runtime registry and optional formatter.

    Args:
        registry: Task/runtime registry containing loaders and runners.
        action_tracker: Tracker object passed through for API compatibility.
        context: Evaluation context containing task, runtime, and split metadata.
        runner_kwargs: Extra keyword arguments for runner invocation.
        split_formatter: Optional formatter converting runner output to payload.

    Returns:
        Callable with signature ``fn(split, loader) -> MetricPayload``.
    """
    del action_tracker
    kwargs = dict(runner_kwargs or {})
    runner = registry.resolve_runner(context.task_type, context.runtime_framework)

    def _evaluate_split(split: str, loader: object) -> MetricPayload:
        """Evaluate a single split via runtime runner and normalize output.

        Args:
            split: Split name passed by the orchestrator loop.
            loader: Loader instance for the requested split.

        Returns:
            Normalized metric payload records for the evaluated split.

        Raises:
            InferenceRunnerError: If runner execution fails or output is invalid.
        """
        try:
            raw_output = runner(loader, kwargs["model"], **kwargs.get("runner_extra", {}))
        except Exception as exc:
            raise InferenceRunnerError(
                f"Runner failed for split='{split}', runtime='{context.runtime_framework}'"
            ) from exc

        if split_formatter is None:
            if isinstance(raw_output, list):
                return validate_payload(raw_output)
            raise InferenceRunnerError(
                "Runner output is not payload list and no split_formatter was provided."
            )
        return split_formatter(split, raw_output)

    return _evaluate_split


def load_model_from_registry(
    registry: RuntimeTaskRegistry,
    *,
    action_tracker: Any,
    task_type: str,
    runtime_framework: str,
) -> Any:
    """Load model using task/runtime registry and raise typed errors.

    Args:
        registry: Task/runtime registry containing model loaders.
        action_tracker: Tracker-like object passed to loader implementation.
        task_type: Task identifier used to resolve loader namespace.
        runtime_framework: Runtime descriptor used to resolve loader variant.

    Returns:
        Loaded model artifact returned by resolved loader.
    """
    try:
        loader = registry.resolve_loader(task_type, runtime_framework)
        return loader(action_tracker)
    except Exception as exc:
        raise ModelLoadError(
            f"Model load failed for task='{task_type}', runtime='{runtime_framework}'"
        ) from exc

