"""Parallel subplan execution scheduler with max_parallel concurrency control.

This module provides a dedicated scheduler for managing parallel execution of
subplans with configurable concurrency limits. It wraps the SubplanExecutionService
and provides a higher-level interface for orchestrating subplan execution.

The scheduler supports:
- Configurable max_parallel concurrency limit (1-50)
- Sequential, parallel, and dependency-ordered execution modes
- Automatic queuing of subplans when max_parallel limit is reached
- Parent plan blocking until all subplans complete
- Comprehensive failure handling and retry logic
- Merge strategy selection for combining subplan outputs

Design:
    The scheduler delegates actual execution to SubplanExecutionService while
    providing queue management and concurrency control at a higher level.
    All state is immutable and carried through the execution result objects.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import TYPE_CHECKING

from cleveragents.application.services.subplan_execution_service import (
    SubplanExecutionResult,
    SubplanExecutionService,
    SubplanExecutorFn,
)
from cleveragents.domain.models.core.plan import (
    ExecutionMode,
    SubplanConfig,
    SubplanStatus,
)

if TYPE_CHECKING:
    from cleveragents.application.services.subplan_merge_service import (
        SubplanMergeService,
    )

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class SubplanQueue:
    """Queue of subplans waiting to execute.

    Attributes:
        pending: Subplans waiting to start execution.
        active: Subplans currently executing.
        completed: Subplans that have finished (success or failure).
    """

    pending: list[SubplanStatus] = field(default_factory=list)
    active: list[SubplanStatus] = field(default_factory=list)
    completed: list[SubplanStatus] = field(default_factory=list)

    @property
    def total_count(self) -> int:
        """Total number of subplans (pending + active + completed)."""
        return len(self.pending) + len(self.active) + len(self.completed)

    @property
    def is_empty(self) -> bool:
        """Check if all queues are empty."""
        return len(self.pending) == 0 and len(self.active) == 0

    @property
    def all_done(self) -> bool:
        """Check if all subplans have completed."""
        return len(self.pending) == 0 and len(self.active) == 0

    def move_to_active(self, count: int) -> SubplanQueue:
        """Move up to *count* pending subplans to active.

        Args:
            count: Maximum number of subplans to move.

        Returns:
            New queue with updated pending and active lists.
        """
        to_move = self.pending[:count]
        remaining_pending = self.pending[count:]
        new_active = [*self.active, *to_move]
        return SubplanQueue(
            pending=remaining_pending,
            active=new_active,
            completed=self.completed,
        )

    def move_to_completed(self, subplan_id: str, status: SubplanStatus) -> SubplanQueue:
        """Move a subplan from active to completed.

        Args:
            subplan_id: The ID of the subplan to move.
            status: The updated status of the subplan.

        Returns:
            New queue with updated active and completed lists.
        """
        new_active = [s for s in self.active if s.subplan_id != subplan_id]
        new_completed = [*self.completed, status]
        return SubplanQueue(
            pending=self.pending,
            active=new_active,
            completed=new_completed,
        )


@dataclass(frozen=True)
class SchedulerState:
    """Immutable state of the parallel subplan scheduler.

    Attributes:
        queue: Current queue state (pending, active, completed).
        max_parallel: Maximum concurrent subplans allowed.
        execution_mode: How subplans should be executed.
        started_at: When scheduling started.
        completed_at: When all subplans completed (None if still running).
    """

    queue: SubplanQueue = field(default_factory=SubplanQueue)
    max_parallel: int = 5
    execution_mode: ExecutionMode = ExecutionMode.PARALLEL
    started_at: datetime | None = None
    completed_at: datetime | None = None

    @property
    def is_running(self) -> bool:
        """Check if scheduler is currently running."""
        return self.started_at is not None and self.completed_at is None

    @property
    def available_slots(self) -> int:
        """Number of available execution slots."""
        return max(0, self.max_parallel - len(self.queue.active))

    @property
    def can_start_more(self) -> bool:
        """Check if more subplans can be started."""
        return len(self.queue.pending) > 0 and self.available_slots > 0


class ParallelSubplanScheduler:
    """Scheduler for parallel subplan execution with max_parallel concurrency control.

    This scheduler manages the execution of multiple subplans with a configurable
    concurrency limit. It ensures that no more than max_parallel subplans execute
    simultaneously, queuing additional subplans until execution slots become available.

    The scheduler supports three execution modes:
    - SEQUENTIAL: Execute subplans one at a time
    - PARALLEL: Execute up to max_parallel subplans concurrently
    - DEPENDENCY_ORDERED: Execute respecting DAG dependencies with concurrent waves

    Args:
        config: Subplan execution configuration including max_parallel limit.
        executor_fn: Callable that executes a single subplan.
        merge_service: Optional service for merging subplan outputs.
        parent_plan_id: Optional parent plan identifier for logging/checkpoints.

    Raises:
        ValueError: If config or executor_fn is None.
    """

    def __init__(
        self,
        config: SubplanConfig,
        executor_fn: SubplanExecutorFn,
        merge_service: SubplanMergeService | None = None,
        parent_plan_id: str = "",
    ) -> None:
        if config is None:
            raise ValueError("config must not be None")
        if executor_fn is None:
            raise ValueError("executor_fn must not be None")

        self._config = config
        self._executor_fn = executor_fn
        self._merge_service = merge_service
        self._parent_plan_id = parent_plan_id
        self._state = SchedulerState(
            max_parallel=config.max_parallel,
            execution_mode=config.execution_mode,
        )

    @property
    def config(self) -> SubplanConfig:
        """The subplan execution configuration."""
        return self._config

    @property
    def state(self) -> SchedulerState:
        """Current scheduler state."""
        return self._state

    @property
    def max_parallel(self) -> int:
        """Maximum concurrent subplans allowed."""
        return self._config.max_parallel

    @property
    def execution_mode(self) -> ExecutionMode:
        """Current execution mode."""
        return self._config.execution_mode

    def schedule(
        self,
        subplan_statuses: list[SubplanStatus],
        base_files: dict[str, str],
        dependency_graph: dict[str, list[str]] | None = None,
    ) -> SubplanExecutionResult:
        """Schedule and execute all subplans with max_parallel concurrency control.

        This method orchestrates the execution of subplans, ensuring that no more
        than max_parallel subplans execute simultaneously. Additional subplans are
        queued and started as execution slots become available.

        The parent plan blocks until all subplans complete, regardless of execution
        mode or concurrency limit.

        Args:
            subplan_statuses: Status objects for each subplan to execute.
            base_files: File contents before subplans (for merge base).
            dependency_graph: For DEPENDENCY_ORDERED mode, maps each subplan_id
                to the list of subplan_ids it depends on.

        Returns:
            A SubplanExecutionResult with updated statuses and merge outcome.

        Raises:
            ValueError: If subplan_statuses is empty.
            ValueError: If DEPENDENCY_ORDERED mode but no dependency_graph.
        """
        if not subplan_statuses:
            raise ValueError("subplan_statuses must not be empty")

        if (
            self._config.execution_mode == ExecutionMode.DEPENDENCY_ORDERED
            and dependency_graph is None
        ):
            raise ValueError("dependency_graph is required for DEPENDENCY_ORDERED mode")

        # Initialize scheduler state
        self._state = SchedulerState(
            queue=SubplanQueue(pending=subplan_statuses),
            max_parallel=self._config.max_parallel,
            execution_mode=self._config.execution_mode,
            started_at=datetime.now(tz=UTC),
        )

        logger.info(
            "scheduler_started max_parallel=%d mode=%s subplan_count=%d",
            self._config.max_parallel,
            self._config.execution_mode.value,
            len(subplan_statuses),
        )

        # Delegate to SubplanExecutionService for actual execution
        service = SubplanExecutionService(
            config=self._config,
            executor_fn=self._executor_fn,
            merge_service=self._merge_service,
            parent_plan_id=self._parent_plan_id,
        )

        result = service.execute_all(
            subplan_statuses=subplan_statuses,
            base_files=base_files,
            dependency_graph=dependency_graph,
        )

        # Update final state
        self._state = SchedulerState(
            queue=SubplanQueue(
                pending=[],
                active=[],
                completed=result.statuses,
            ),
            max_parallel=self._config.max_parallel,
            execution_mode=self._config.execution_mode,
            started_at=self._state.started_at,
            completed_at=datetime.now(tz=UTC),
        )

        logger.info(
            "scheduler_completed total_duration_ms=%d all_succeeded=%s failed_count=%d",
            result.total_duration_ms,
            result.all_succeeded,
            len(result.failed_subplan_ids),
        )

        return result

    def get_queue_status(self) -> dict[str, int]:
        """Get current queue status.

        Returns:
            Dictionary with pending, active, and completed counts.
        """
        return {
            "pending": len(self._state.queue.pending),
            "active": len(self._state.queue.active),
            "completed": len(self._state.queue.completed),
            "total": self._state.queue.total_count,
        }

    def get_available_slots(self) -> int:
        """Get number of available execution slots.

        Returns:
            Number of subplans that can start immediately.
        """
        return self._state.available_slots

    def can_accept_more(self) -> bool:
        """Check if more subplans can be queued.

        Returns:
            True if there are pending subplans waiting to execute.
        """
        return len(self._state.queue.pending) > 0
