"""
Task-Agnostic Strategy Protocols.

Defines the structural contracts that task-specific components must satisfy.
Each protocol represents one dimension of task-specific behavior:

- **Criterion**: Loss function construction
- **Validation metrics**: Per-epoch metric computation
- **Eval pipeline**: Full evaluation orchestration (inference, visualization,
  reporting) as a single cohesive unit

Concrete implementations live under ``orchard.tasks.<task_type>/``
and are registered in the :mod:`orchard.core.task_registry`.
"""

from __future__ import annotations

from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Callable, Protocol, runtime_checkable

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

if TYPE_CHECKING:  # pragma: no cover
    from ..tracking import TrackerProtocol
    from .config import (
        AugmentationConfig,
        DatasetConfig,
        EvaluationConfig,
        TrainingConfig,
    )
    from .paths import RunPaths


@runtime_checkable
class TaskTrainingStep(Protocol):
    """Protocol for task-specific training forward pass and loss computation."""

    def compute_training_loss(
        self,
        model: nn.Module,
        inputs: torch.Tensor,
        targets: torch.Tensor,
        criterion: nn.Module,
        mixup_fn: Callable[..., Any] | None = None,
    ) -> torch.Tensor:
        """
        Execute the forward pass and compute the training loss.

        Encapsulates task-specific differences in how models are called
        and how losses are computed. For classification, this means
        ``model(inputs)`` returning logits with optional MixUp blending.
        For detection, the model returns a loss dict that is summed.

        Args:
            model: Neural network in training mode.
            inputs: Batch of input tensors (already on device).
            targets: Batch of target tensors (already on device).
            criterion: Loss function module.
            mixup_fn: Optional MixUp augmentation callable.

        Returns:
            Scalar loss tensor for backward pass.
        """
        ...  # pragma: no cover


@runtime_checkable
class TaskCriterionFactory(Protocol):
    """Protocol for task-specific loss function construction."""

    def get_criterion(
        self,
        training: TrainingConfig,
        class_weights: torch.Tensor | None = None,
    ) -> nn.Module:
        """
        Build a loss function from training config.

        Args:
            training: Training sub-config with criterion parameters.
            class_weights: Optional per-class weights for imbalanced datasets.

        Returns:
            Configured loss module.
        """
        ...  # pragma: no cover


@runtime_checkable
class TaskValidationMetrics(Protocol):
    """Protocol for per-epoch validation metric computation."""

    def compute_validation_metrics(
        self,
        model: nn.Module,
        val_loader: DataLoader[Any],
        criterion: nn.Module,
        device: torch.device,
    ) -> Mapping[str, float]:
        """
        Compute task-specific validation metrics.

        Must return an immutable mapping with at least a ``"loss"`` key.

        Args:
            model: Neural network model to evaluate.
            val_loader: Validation data provider.
            criterion: Loss function.
            device: Hardware target (CUDA/MPS/CPU).

        Returns:
            Immutable mapping of metric name to value.
        """
        ...  # pragma: no cover


@runtime_checkable
class TaskEvalPipeline(Protocol):
    """Protocol for end-to-end evaluation (inference + visualization + reporting)."""

    def run_evaluation(
        self,
        model: nn.Module,
        test_loader: DataLoader[Any],
        train_losses: list[float],
        val_metrics_history: list[Mapping[str, float]],
        class_names: list[str],
        paths: RunPaths,
        training: TrainingConfig,
        dataset: DatasetConfig,
        augmentation: AugmentationConfig,
        evaluation: EvaluationConfig,
        arch_name: str,
        aug_info: str = "N/A",
        tracker: TrackerProtocol | None = None,
    ) -> Mapping[str, float]:
        """
        Execute the complete evaluation pipeline.

        Coordinates inference, visualization, and structured reporting
        as a single task-specific unit.

        Args:
            model: Trained model (already on target device).
            test_loader: DataLoader for test set.
            train_losses: Training loss history per epoch.
            val_metrics_history: Validation metrics history per epoch.
            class_names: List of class label strings.
            paths: RunPaths for artifact output.
            training: Training sub-config.
            dataset: Dataset sub-config.
            augmentation: Augmentation sub-config.
            evaluation: Evaluation sub-config.
            arch_name: Architecture identifier.
            aug_info: Augmentation description string.
            tracker: Optional experiment tracker for final metrics.

        Returns:
            Mapping of task-specific metric names to float values.
        """
        ...  # pragma: no cover
