from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, Awaitable, Callable

from harbor.metrics.base import BaseMetric
from harbor.models.job.config import RetryConfig
from harbor.models.orchestrator_type import OrchestratorType
from harbor.models.trial.config import TrialConfig
from harbor.models.trial.result import TrialResult
from harbor.trial.hooks import TrialEvent, TrialHookEvent

if TYPE_CHECKING:
    from harbor.trial.trial import Trial

# Backward compatibility alias: OrchestratorEvent was merged into TrialEvent
OrchestratorEvent = TrialEvent

# Type alias for hook callbacks
HookCallback = Callable[[TrialHookEvent], Awaitable[Any]]


class BaseOrchestrator(ABC):
    """
    Handles deployment of multiple trials based on trial configs.

    Needs to handle resuming, canceling, (pausing?), and finalizing.

    TODO: Should this start an active process on the computer that runs until
    completion? For now, yes. Otherwise how could it "finalize" the run?
    """

    def __init__(
        self,
        trial_configs: list[TrialConfig],
        n_concurrent_trials: int,
        metrics: dict[str, list[BaseMetric]],
        quiet: bool = False,
        plain_output: bool = False,
        retry_config: RetryConfig | None = None,
    ):
        self._trial_configs = trial_configs
        self._n_concurrent_trials = n_concurrent_trials
        self._metrics = metrics
        self._quiet = quiet
        self._plain_output = plain_output
        self._retry_config = retry_config or RetryConfig()
        self._hooks: dict[TrialEvent, list[HookCallback]] = {
            event: [] for event in TrialEvent
        }

    def add_hook(self, event: TrialEvent, hook: HookCallback) -> None:
        """Add a hook to be called when the specified event occurs."""
        self._hooks[event].append(hook)

    def _setup_orchestrator_hooks(self, trial: "Trial") -> None:
        """Wire orchestrator-level hooks to the trial.

        This ensures that hooks registered with the orchestrator via add_hook()
        are triggered when the trial fires the corresponding events.
        """

        for event, hooks in self._hooks.items():
            for hook in hooks:
                trial.add_hook(event, hook)

    @staticmethod
    @abstractmethod
    def type() -> OrchestratorType:
        """The type of orchestrator."""

    @abstractmethod
    async def run(self) -> list[TrialResult]:
        pass
