import asyncio
import inspect
import shutil

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.orchestrators.base import BaseOrchestrator
from harbor.trial.trial import Trial
from harbor.utils.logger import logger


class QueueOrchestrator(BaseOrchestrator):
    """
    Queue-based producer-consumer orchestrator for submitting trials dynamically.

    This orchestrator allows trials to be submitted individually or in batches
    and returns futures for the results. It maintains a worker pool that processes
    trials with concurrency control and retry logic.
    """

    # Grace period between launching containers to avoid overwhelming the orchestration system
    CONTAINER_LAUNCH_GRACE_PERIOD_SEC = 2.0
    # Interval for logging queue status
    QUEUE_STATUS_LOG_INTERVAL_SEC = 120.0

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

    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,
        name: str | None = None,
    ):
        super().__init__(
            trial_configs=trial_configs,
            n_concurrent_trials=n_concurrent_trials,
            metrics=metrics,
            quiet=quiet,
            plain_output=plain_output,
            retry_config=retry_config,
        )
        self._logger = logger.getChild(__name__)
        self._queue: asyncio.Queue[tuple[TrialConfig, asyncio.Future[TrialResult]]] = (
            asyncio.Queue()
        )
        self._workers: list[asyncio.Task] = []
        self._is_running = False
        self._trial_results: list[TrialResult] = []
        self._queue_status_task: asyncio.Task | None = None
        self._container_launch_lock = asyncio.Lock()
        self._active_trials = 0

        # Auto-generate name from caller if not provided
        if name is None:
            frame = inspect.currentframe()
            if frame and frame.f_back:
                caller = frame.f_back
                module = caller.f_globals.get("__name__", "unknown")
                # Extract short module name (last component)
                short_module = module.rsplit(".", 1)[-1]
                name = f"{short_module}:{caller.f_lineno}"
        self._name = name or "queue"

    def _should_retry_exception(self, exception_type: str) -> bool:
        """Check if an exception should trigger a retry."""
        if (
            self._retry_config.exclude_exceptions
            and exception_type in self._retry_config.exclude_exceptions
        ):
            self._logger.debug(
                f"Exception {exception_type} is in exclude_exceptions, not retrying"
            )
            return False

        if (
            self._retry_config.include_exceptions
            and exception_type not in self._retry_config.include_exceptions
        ):
            self._logger.debug(
                f"Exception {exception_type} is not in include_exceptions, not retrying"
            )
            return False

        return True

    def _calculate_backoff_delay(self, attempt: int) -> float:
        """Calculate the backoff delay for a retry attempt."""
        delay = self._retry_config.min_wait_sec * (
            self._retry_config.wait_multiplier**attempt
        )
        return min(delay, self._retry_config.max_wait_sec)

    async def _log_queue_status(self):
        """Background task that logs queue status periodically."""
        while self._is_running:
            try:
                await asyncio.sleep(self.QUEUE_STATUS_LOG_INTERVAL_SEC)
                if self._is_running:
                    queue_size = self._queue.qsize()
                    completed = len(self._trial_results)
                    self._logger.info(
                        f"[{self._name}] Queue status: {queue_size} pending, {self._active_trials} active, "
                        f"{completed} completed, {self._n_concurrent_trials} workers"
                    )
            except asyncio.CancelledError:
                break
            except Exception as e:
                self._logger.warning(f"[{self._name}] Error logging queue status: {e}")

    async def _execute_trial_with_retries(
        self, trial_config: TrialConfig
    ) -> TrialResult:
        """Execute a trial with retry logic."""
        for attempt in range(self._retry_config.max_retries + 1):
            trial = Trial(trial_config)
            self._setup_orchestrator_hooks(trial)
            result = await trial.run()

            if result.exception_info is None:
                return result

            if not self._should_retry_exception(result.exception_info.exception_type):
                self._logger.debug(
                    "Not retrying trial because the exception is not in "
                    "include_exceptions or the maximum number of retries has been "
                    "reached"
                )
                return result
            if attempt == self._retry_config.max_retries:
                self._logger.debug(
                    "Not retrying trial because the maximum number of retries has been "
                    "reached"
                )
                return result

            shutil.rmtree(trial.trial_dir, ignore_errors=True)

            delay = self._calculate_backoff_delay(attempt)

            self._logger.debug(
                f"Trial {trial_config.trial_name} failed with exception "
                f"{result.exception_info.exception_type}. Retrying in "
                f"{delay:.2f} seconds..."
            )

            await asyncio.sleep(delay)

        raise RuntimeError(
            f"Trial {trial_config.trial_name} produced no result. This should never "
            "happen."
        )

    async def _worker(self):
        """Worker coroutine that processes trials from the queue."""
        while self._is_running or not self._queue.empty():
            try:
                # Use a timeout to periodically check if we should shutdown
                trial_config, future = await asyncio.wait_for(
                    self._queue.get(), timeout=1.0
                )
            except asyncio.TimeoutError:
                continue

            try:
                # Grace period between container launches to avoid overwhelming the system
                async with self._container_launch_lock:
                    await asyncio.sleep(self.CONTAINER_LAUNCH_GRACE_PERIOD_SEC)

                self._active_trials += 1
                result = await self._execute_trial_with_retries(trial_config)
                self._active_trials -= 1
                self._trial_results.append(result)

                # Set the result on the future
                # Note: Hooks are now invoked by the trial itself via _setup_orchestrator_hooks
                if not future.done():
                    future.set_result(result)

            except asyncio.CancelledError:
                self._active_trials = max(0, self._active_trials - 1)
                # Cancel the future when worker is cancelled mid-execution
                if not future.done():
                    future.cancel()
                raise
            except Exception as e:
                self._active_trials = max(0, self._active_trials - 1)
                # Set the exception on the future
                if not future.done():
                    future.set_exception(e)
                self._logger.error(
                    f"Worker encountered error processing trial {trial_config.trial_name}: {e}"
                )
            finally:
                self._queue.task_done()

    async def start(self):
        """Start the worker pool."""
        if self._is_running:
            raise RuntimeError("QueueOrchestrator is already running")

        self._is_running = True

        # Start worker tasks
        for _ in range(self._n_concurrent_trials):
            worker = asyncio.create_task(self._worker())
            self._workers.append(worker)

        # Start queue status logging task
        self._queue_status_task = asyncio.create_task(self._log_queue_status())

        self._logger.info(
            f"[{self._name}] Started {self._n_concurrent_trials} workers "
            f"(status every {self.QUEUE_STATUS_LOG_INTERVAL_SEC}s, "
            f"{self.CONTAINER_LAUNCH_GRACE_PERIOD_SEC}s launch grace period)"
        )

    async def submit(self, trial_config: TrialConfig) -> asyncio.Future[TrialResult]:
        """
        Submit a trial for execution.

        Args:
            trial_config: Configuration for the trial to execute

        Returns:
            A future that will be resolved with the TrialResult when the trial completes

        Raises:
            RuntimeError: If the orchestrator has not been started
        """
        if not self._is_running:
            raise RuntimeError(
                "QueueOrchestrator must be started before submitting trials"
            )

        future: asyncio.Future[TrialResult] = asyncio.Future()
        await self._queue.put((trial_config, future))
        return future

    async def submit_batch(
        self, configs: list[TrialConfig]
    ) -> list[asyncio.Future[TrialResult]]:
        """
        Submit multiple trials for execution.

        Args:
            configs: List of trial configurations to execute

        Returns:
            A list of futures that will be resolved with TrialResults when trials complete

        Raises:
            RuntimeError: If the orchestrator has not been started
        """
        if not self._is_running:
            raise RuntimeError(
                "QueueOrchestrator must be started before submitting trials"
            )

        futures: list[asyncio.Future[TrialResult]] = []
        for config in configs:
            future: asyncio.Future[TrialResult] = asyncio.Future()
            await self._queue.put((config, future))
            futures.append(future)
        return futures

    async def shutdown(self, wait: bool = True):
        """
        Gracefully shutdown the orchestrator.

        Args:
            wait: If True, wait for all pending trials to complete before shutting down.
                  If False, cancel running workers and pending trials immediately.
        """
        if not self._is_running:
            return

        self._is_running = False

        # Cancel queue status logging task
        if self._queue_status_task is not None:
            self._queue_status_task.cancel()
            try:
                await self._queue_status_task
            except asyncio.CancelledError:
                pass
            self._queue_status_task = None

        if not wait:
            # Cancel all running workers immediately
            for worker in self._workers:
                worker.cancel()

        if wait:
            # Wait for the queue to be empty
            await self._queue.join()

        # Wait for all worker tasks to complete (or be cancelled)
        # Use TaskGroup for proper cancellation propagation
        try:
            async with asyncio.TaskGroup() as tg:
                for worker in self._workers:
                    tg.create_task(self._await_worker(worker))
        except* asyncio.CancelledError:
            # Expected when workers are cancelled during shutdown
            pass

        # For non-waiting shutdown, drain queue and cancel pending futures
        if not wait:
            cancelled_count = 0
            while not self._queue.empty():
                try:
                    _config, future = self._queue.get_nowait()
                    if not future.done():
                        future.cancel()
                        cancelled_count += 1
                    self._queue.task_done()
                except asyncio.QueueEmpty:
                    break
            if cancelled_count > 0:
                self._logger.info(
                    f"[{self._name}] Cancelled {cancelled_count} pending trials"
                )

        self._workers.clear()

        self._logger.info(
            f"[{self._name}] Shutdown complete. Total completed: {len(self._trial_results)}"
        )

    async def _await_worker(self, worker: asyncio.Task) -> None:
        """Await a worker task, suppressing CancelledError."""
        try:
            await worker
        except asyncio.CancelledError:
            pass

    async def run(self) -> list[TrialResult]:
        """
        Run all trials that were provided during initialization.

        This method maintains compatibility with the BaseOrchestrator interface.
        It starts the orchestrator, submits all initial trial configs, waits for
        completion, and then shuts down.

        Returns:
            List of TrialResults for all trials
        """
        # Start the orchestrator
        await self.start()

        try:
            # Submit all initial trial configs
            futures = await self.submit_batch(self._trial_configs)

            # Wait for all futures to complete using TaskGroup for proper cancellation
            results: list[TrialResult] = []
            errors: list[tuple[int, Exception]] = []

            async with asyncio.TaskGroup() as tg:

                async def collect_result(idx: int, future: asyncio.Future[TrialResult]):
                    try:
                        result = await future
                        results.append((idx, result))  # type: ignore[arg-type]
                    except Exception as e:
                        errors.append((idx, e))

                for i, future in enumerate(futures):
                    tg.create_task(collect_result(i, future))

            # Check for errors
            if errors:
                # Report all errors, then raise the first one
                for idx, error in errors:
                    self._logger.error(
                        f"Trial {self._trial_configs[idx].trial_name} failed: {error}"
                    )
                raise errors[0][1]

            # Sort results by index to maintain order
            results.sort(key=lambda x: x[0])
            return [r[1] for r in results]  # type: ignore[misc]

        finally:
            # Always shutdown gracefully
            await self.shutdown(wait=True)
