"""Subplan execution scheduler service.

Orchestrates the execution of child subplans within a parent plan, supporting
three scheduling modes (sequential, parallel, dependency-ordered) and four
merge strategies for combining subplan outputs.

The service delegates individual subplan execution to a caller-provided
executor callable and uses :class:`SubplanMergeService` for combining
sandbox outputs.  Failure handling follows :class:`SubplanFailureHandler`
semantics (retry, stop-others).

When a ``CheckpointService`` is injected and the ``on_subplan_spawn``
trigger is active, a checkpoint is created immediately after each child
plan is spawned (before its execution begins).

Design decisions:
    - Dependency injection: executor and merge service are injected.
    - Thread-based parallelism: uses ``concurrent.futures.ThreadPoolExecutor``
      for PARALLEL mode (consistent with the synchronous codebase).
    - Stateless: all state is carried in the ``SubplanConfig`` and
      ``SubplanStatus`` objects on the parent ``Plan``.
"""

from __future__ import annotations

import logging
import time
from collections import deque
from collections.abc import Callable
from concurrent.futures import CancelledError, Future, ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import TYPE_CHECKING

from cleveragents.application.services.subplan_merge_service import (
    MergeConflictError,
    SubplanMergeResult,
    SubplanMergeService,
)
from cleveragents.domain.models.core.plan import (
    ExecutionMode,
    ProcessingState,
    SubplanAttempt,
    SubplanConfig,
    SubplanFailureHandler,
    SubplanStatus,
)
from cleveragents.tool.runner import DEFAULT_AUTO_TRIGGERS

if TYPE_CHECKING:
    from cleveragents.application.services.checkpoint_service import CheckpointService

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Value objects
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class SubplanExecutionOutput:
    """Output produced by executing a single subplan.

    Attributes:
        subplan_id: The ULID of the executed subplan.
        success: Whether execution succeeded.
        files: Mapping of relative file paths to their content.
        error: Error message if execution failed.
        files_changed: Number of files modified.
        changeset_summary: Brief description of changes.
    """

    subplan_id: str
    success: bool
    files: dict[str, str] = field(default_factory=dict)
    error: str | None = None
    files_changed: int = 0
    changeset_summary: str | None = None


# ---------------------------------------------------------------------------
# Types
# ---------------------------------------------------------------------------

SubplanExecutorFn = Callable[[SubplanStatus], SubplanExecutionOutput]
"""Callable that executes a single subplan and returns its output."""


@dataclass(frozen=True)
class SubplanExecutionResult:
    """Aggregate result of executing all subplans for a parent plan.

    Attributes:
        all_succeeded: ``True`` if every subplan completed successfully.
        statuses: Updated status objects for each subplan.
        merge_result: Result of merging subplan outputs (None if merge
            was not attempted, e.g. all subplans failed).
        total_duration_ms: Wall-clock time for the entire execution in ms.
        failed_subplan_ids: IDs of subplans that failed.
    """

    all_succeeded: bool
    statuses: list[SubplanStatus]
    merge_result: SubplanMergeResult | None = None
    total_duration_ms: int = 0
    failed_subplan_ids: list[str] = field(default_factory=list)


# ---------------------------------------------------------------------------
# Service
# ---------------------------------------------------------------------------


class SubplanExecutionService:
    """Execute and merge subplans according to a parent plan's configuration.

    Args:
        config: The parent plan's subplan configuration.
        executor_fn: Callable that executes a single subplan.
        merge_service: Service for merging subplan outputs.  If ``None``,
            a default merge service is created from the config's strategy.
        failure_handler: Handler for failure decisions.  If ``None``,
            a default handler is created.
        checkpoint_service: Optional ``CheckpointService`` for automatic
            ``on_subplan_spawn`` checkpoint creation.  When provided and
            the trigger is active, a checkpoint is created immediately
            after each child plan is spawned.
        auto_checkpoint_triggers: Set of active trigger names.  Defaults
            to all four triggers when ``None``.
        parent_plan_id: Optional parent plan identifier used when creating
            ``on_subplan_spawn`` checkpoints.

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

    def __init__(
        self,
        config: SubplanConfig,
        executor_fn: SubplanExecutorFn,
        merge_service: SubplanMergeService | None = None,
        failure_handler: SubplanFailureHandler | None = None,
        checkpoint_service: CheckpointService | None = None,
        auto_checkpoint_triggers: frozenset[str] | 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: SubplanConfig = config
        self._executor_fn: SubplanExecutorFn = executor_fn
        self._merge_service: SubplanMergeService = merge_service or SubplanMergeService(
            config.merge_strategy
        )
        self._failure_handler: SubplanFailureHandler = (
            failure_handler or SubplanFailureHandler()
        )
        self._checkpoint_service = checkpoint_service
        self._auto_checkpoint_triggers: frozenset[str] = (
            auto_checkpoint_triggers
            if auto_checkpoint_triggers is not None
            else DEFAULT_AUTO_TRIGGERS
        )
        self._parent_plan_id = parent_plan_id

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

    @property
    def merge_service(self) -> SubplanMergeService:
        """The merge service used for combining outputs."""
        return self._merge_service

    def execute_all(
        self,
        subplan_statuses: list[SubplanStatus],
        base_files: dict[str, str],
        dependency_graph: dict[str, list[str]] | None = None,
    ) -> SubplanExecutionResult:
        """Execute all subplans and merge their outputs.

        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 :class:`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")

        start = time.monotonic()

        mode = self._config.execution_mode
        if mode == ExecutionMode.SEQUENTIAL:
            statuses, outputs = self._execute_sequential(subplan_statuses)
        elif mode == ExecutionMode.PARALLEL:
            statuses, outputs = self._execute_parallel(subplan_statuses)
        else:
            statuses, outputs = self._execute_dependency_ordered(
                subplan_statuses, dependency_graph or {}
            )

        duration_ms = int((time.monotonic() - start) * 1000)

        failed_ids = [
            s.subplan_id for s in statuses if s.status == ProcessingState.ERRORED
        ]
        all_succeeded = len(failed_ids) == 0

        # Merge outputs if any subplans produced results
        merge_result: SubplanMergeResult | None = None
        successful_outputs = [(sid, files) for sid, files in outputs if files]
        if successful_outputs:
            try:
                merge_result = self._merge_service.merge(base_files, successful_outputs)
            except MergeConflictError:
                logger.warning("Merge conflict after subplan execution")
                all_succeeded = False

        return SubplanExecutionResult(
            all_succeeded=all_succeeded,
            statuses=statuses,
            merge_result=merge_result,
            total_duration_ms=duration_ms,
            failed_subplan_ids=failed_ids,
        )

    # -- scheduling modes ----------------------------------------------------

    def _execute_sequential(
        self,
        statuses: list[SubplanStatus],
    ) -> tuple[list[SubplanStatus], list[tuple[str, dict[str, str]]]]:
        """Execute subplans one at a time in order.

        When ``timeout_per_subplan_seconds`` is configured, each subplan
        is executed with a per-subplan wall-clock timeout.
        """
        updated: list[SubplanStatus] = []
        outputs: list[tuple[str, dict[str, str]]] = []
        timeout = self._config.timeout_per_subplan_seconds

        for status in statuses:
            if timeout is not None:
                result_status, output = self._execute_with_timeout(status, timeout)
            else:
                result_status, output = self._execute_one_with_retry(status)
            updated.append(result_status)
            outputs.append((result_status.subplan_id, output))

            if (
                result_status.status == ProcessingState.ERRORED
                and self._failure_handler.should_stop_others(
                    self._config, result_status
                )
            ):
                # Cancel remaining subplans
                for remaining in statuses[len(updated) :]:
                    cancelled = self._cancel_status(remaining)
                    updated.append(cancelled)
                    outputs.append((cancelled.subplan_id, {}))
                break

        return updated, outputs

    def _execute_parallel(
        self,
        statuses: list[SubplanStatus],
    ) -> tuple[list[SubplanStatus], list[tuple[str, dict[str, str]]]]:
        """Execute subplans concurrently up to max_parallel.

        Outputs are returned in completion order so that ordering-sensitive
        merge strategies (e.g. ``SEQUENTIAL_APPLY``) apply changes in the
        order subplans actually finished.  Futures cancelled via fail-fast
        are reported with ``CANCELLED`` status, not ``ERRORED``.

        When ``timeout_per_subplan_seconds`` is configured, each subplan
        is executed with a per-subplan wall-clock timeout.
        """
        max_workers = min(self._config.max_parallel, len(statuses))
        results_map: dict[str, tuple[SubplanStatus, dict[str, str]]] = {}
        completion_order: list[str] = []
        stop_flag = False
        timeout = self._config.timeout_per_subplan_seconds

        with ThreadPoolExecutor(max_workers=max_workers) as pool:
            future_to_id: dict[Future[tuple[SubplanStatus, dict[str, str]]], str] = {}
            for status in statuses:
                if timeout is not None:
                    future: Future[tuple[SubplanStatus, dict[str, str]]] = pool.submit(
                        self._execute_with_timeout, status, timeout
                    )
                else:
                    future = pool.submit(self._execute_one_with_retry, status)
                future_to_id[future] = status.subplan_id

            for future in as_completed(future_to_id):
                subplan_id = future_to_id[future]
                try:
                    result_status, output = future.result()
                except CancelledError:
                    result_status = self._cancel_status(
                        next(s for s in statuses if s.subplan_id == subplan_id)
                    )
                    output = {}
                except Exception as exc:  # pragma: no cover - defensive
                    result_status = self._error_status(
                        next(s for s in statuses if s.subplan_id == subplan_id),
                        str(exc),
                    )
                    output = {}

                results_map[subplan_id] = (result_status, output)
                completion_order.append(subplan_id)

                if (
                    result_status.status == ProcessingState.ERRORED
                    and self._failure_handler.should_stop_others(
                        self._config, result_status
                    )
                ):
                    stop_flag = True
                    # Cancel remaining futures
                    for f in future_to_id:
                        if not f.done():
                            f.cancel()

        # Build results in completion order (not original input order)
        updated: list[SubplanStatus] = []
        outputs: list[tuple[str, dict[str, str]]] = []
        for subplan_id in completion_order:
            s, o = results_map[subplan_id]
            updated.append(s)
            outputs.append((s.subplan_id, o))

        # Safety net: append any subplans not captured by as_completed
        if stop_flag:  # pragma: no branch - as_completed always yields all futures
            for status in statuses:
                if status.subplan_id not in results_map:  # pragma: no cover - defensive
                    cancelled = self._cancel_status(status)
                    updated.append(cancelled)
                    outputs.append((cancelled.subplan_id, {}))

        return updated, outputs

    def _execute_dependency_ordered(
        self,
        statuses: list[SubplanStatus],
        dependency_graph: dict[str, list[str]],
    ) -> tuple[list[SubplanStatus], list[tuple[str, dict[str, str]]]]:
        """Execute subplans respecting dependency order with concurrent waves.

        Independent subplans (those whose dependencies are all satisfied)
        within the same topological wave are executed concurrently using
        :meth:`_execute_wave`.  Subsequent waves wait for the previous
        wave to finish.
        """
        node_ids = [s.subplan_id for s in statuses]
        # Validate no cycles (raises ValueError if cycle detected)
        self._topological_sort(node_ids, dependency_graph)

        status_map = {s.subplan_id: s for s in statuses}
        timeout = self._config.timeout_per_subplan_seconds

        # Build dependency tracking
        node_set = set(node_ids)
        remaining_deps: dict[str, set[str]] = {}
        dependents: dict[str, list[str]] = {nid: [] for nid in node_ids}

        for nid in node_ids:
            deps = dependency_graph.get(nid, [])
            remaining_deps[nid] = {d for d in deps if d in node_set}
            for dep in deps:
                if dep in node_set:
                    dependents[dep].append(nid)

        updated: list[SubplanStatus] = []
        outputs: list[tuple[str, dict[str, str]]] = []
        completed: set[str] = set()
        stop_flag = False

        while len(completed) < len(node_ids):
            ready = [
                nid
                for nid in node_ids
                if nid not in completed and len(remaining_deps[nid]) == 0
            ]
            if not ready:  # pragma: no cover - defensive
                break  # Shouldn't happen after _topological_sort validation

            if stop_flag:
                for nid in ready:
                    cancelled = self._cancel_status(status_map[nid])
                    updated.append(cancelled)
                    outputs.append((cancelled.subplan_id, {}))
                    completed.add(nid)
                    for dep_nid in dependents[nid]:
                        remaining_deps[dep_nid].discard(nid)
                continue

            wave_results = self._execute_wave(
                [status_map[nid] for nid in ready], timeout
            )
            for result_status, output in wave_results:
                nid = result_status.subplan_id
                updated.append(result_status)
                outputs.append((nid, output))
                completed.add(nid)
                for dep_nid in dependents[nid]:
                    remaining_deps[dep_nid].discard(nid)

                if (
                    result_status.status == ProcessingState.ERRORED
                    and self._failure_handler.should_stop_others(
                        self._config, result_status
                    )
                ):
                    stop_flag = True

        return updated, outputs

    # -- single subplan execution with retry ---------------------------------

    def _try_create_spawn_checkpoint(self, subplan_id: str) -> None:
        """Create an ``on_subplan_spawn`` checkpoint (non-fatal).

        Fires immediately after a child plan is spawned, before its
        execution begins.  Silently skips when no ``CheckpointService``
        is configured or when the trigger is not active.

        Args:
            subplan_id: The ULID of the spawned child plan.
        """
        if self._checkpoint_service is None:
            return
        if "on_subplan_spawn" not in self._auto_checkpoint_triggers:
            return
        plan_id = self._parent_plan_id or subplan_id
        try:
            self._checkpoint_service.create_checkpoint(
                plan_id=plan_id,
                sandbox_ref=f"subplan-spawn-{subplan_id}",
                reason="Automatic checkpoint: on_subplan_spawn",
                source_tool="subplan_execution_service",
                phase="execute",
                checkpoint_type="pre_write",
            )
        except Exception:
            logger.debug(
                "subplan_spawn_checkpoint_failed plan_id=%s subplan_id=%s",
                plan_id,
                subplan_id,
                exc_info=True,
            )

    def _execute_one_with_retry(
        self,
        status: SubplanStatus,
    ) -> tuple[SubplanStatus, dict[str, str]]:
        """Execute a single subplan with retry logic.

        Returns:
            Tuple of (updated_status, output_files).
        """
        current_status = status.model_copy(
            update={
                "status": ProcessingState.PROCESSING,
                "started_at": datetime.now(tz=UTC),
            }
        )

        # on_subplan_spawn: create checkpoint immediately after spawning
        # the child plan (before its first execution attempt).
        self._try_create_spawn_checkpoint(status.subplan_id)

        while True:
            attempt_start = datetime.now(tz=UTC)
            try:
                output = self._executor_fn(current_status)
                completed_status = current_status.model_copy(
                    update={
                        "status": (
                            ProcessingState.COMPLETE
                            if output.success
                            else ProcessingState.ERRORED
                        ),
                        "completed_at": datetime.now(tz=UTC),
                        "error": output.error,
                        "files_changed": output.files_changed,
                        "changeset_summary": output.changeset_summary,
                    }
                )

                if not output.success:
                    completed_status = self._record_attempt(
                        completed_status, attempt_start, output.error
                    )
                    if self._failure_handler.should_retry(
                        self._config, completed_status
                    ):
                        current_status = completed_status.model_copy(
                            update={
                                "status": ProcessingState.PROCESSING,
                                "attempt_number": (completed_status.attempt_number + 1),
                                "error": None,
                                "completed_at": None,
                            }
                        )
                        continue
                    return completed_status, output.files

                return completed_status, output.files

            except Exception as exc:
                error_msg = f"{type(exc).__name__}: {exc}"
                errored_status = current_status.model_copy(
                    update={
                        "status": ProcessingState.ERRORED,
                        "completed_at": datetime.now(tz=UTC),
                        "error": error_msg,
                    }
                )
                errored_status = self._record_attempt(
                    errored_status, attempt_start, error_msg
                )

                if self._failure_handler.should_retry(self._config, errored_status):
                    current_status = errored_status.model_copy(
                        update={
                            "status": ProcessingState.PROCESSING,
                            "attempt_number": errored_status.attempt_number + 1,
                            "error": None,
                            "completed_at": None,
                        }
                    )
                    continue
                return errored_status, {}

    # -- timeout & wave helpers -----------------------------------------------

    def _execute_with_timeout(
        self,
        status: SubplanStatus,
        timeout_seconds: int,
    ) -> tuple[SubplanStatus, dict[str, str]]:
        """Execute a single subplan with a per-subplan wall-clock timeout.

        Submits the retry loop to a dedicated thread and waits up to
        *timeout_seconds*.  If the deadline expires the subplan is marked
        ``ERRORED`` with a descriptive timeout message.

        Args:
            status: The subplan status to execute.
            timeout_seconds: Maximum wall-clock seconds for this subplan.

        Returns:
            Tuple of (updated_status, output_files).
        """
        pool = ThreadPoolExecutor(max_workers=1)
        future: Future[tuple[SubplanStatus, dict[str, str]]] = pool.submit(
            self._execute_one_with_retry, status
        )
        try:
            return future.result(timeout=timeout_seconds)
        except TimeoutError:
            future.cancel()
            error_msg = (
                f"SubplanTimeoutError: subplan timed out after {timeout_seconds}s"
            )
            return self._error_status(status, error_msg), {}
        finally:
            pool.shutdown(wait=False, cancel_futures=True)

    def _execute_wave(
        self,
        statuses: list[SubplanStatus],
        timeout: int | None,
    ) -> list[tuple[SubplanStatus, dict[str, str]]]:
        """Execute a wave of independent subplans concurrently.

        Used by :meth:`_execute_dependency_ordered` to run all subplans
        whose dependencies have been satisfied in parallel, respecting
        the ``max_parallel`` limit.

        Args:
            statuses: Independent subplans that can execute in parallel.
            timeout: Optional per-subplan timeout in seconds.

        Returns:
            List of ``(status, output_files)`` tuples.
        """
        if len(statuses) == 1:
            if timeout is not None:
                return [self._execute_with_timeout(statuses[0], timeout)]
            return [self._execute_one_with_retry(statuses[0])]

        max_workers = min(self._config.max_parallel, len(statuses))
        results: list[tuple[SubplanStatus, dict[str, str]]] = []

        with ThreadPoolExecutor(max_workers=max_workers) as pool:
            future_map: dict[
                Future[tuple[SubplanStatus, dict[str, str]]], SubplanStatus
            ] = {}
            for status in statuses:
                if timeout is not None:
                    future: Future[tuple[SubplanStatus, dict[str, str]]] = pool.submit(
                        self._execute_with_timeout, status, timeout
                    )
                else:
                    future = pool.submit(self._execute_one_with_retry, status)
                future_map[future] = status

            for future in as_completed(future_map):
                original_status = future_map[future]
                try:
                    result_status, output = future.result()
                except CancelledError:  # pragma: no cover - defensive
                    result_status = self._cancel_status(original_status)
                    output = {}
                except Exception as exc:  # pragma: no cover - defensive
                    result_status = self._error_status(original_status, str(exc))
                    output = {}
                results.append((result_status, output))

        return results

    # -- helpers -------------------------------------------------------------

    @staticmethod
    def _record_attempt(
        status: SubplanStatus,
        started_at: datetime,
        error: str | None,
    ) -> SubplanStatus:
        """Record a failed attempt in the status's previous_attempts."""
        attempt = SubplanAttempt(
            attempt_number=status.attempt_number,
            started_at=started_at,
            completed_at=datetime.now(tz=UTC),
            error=error,
            was_retried=True,
        )
        new_attempts = [*status.previous_attempts, attempt]
        return status.model_copy(update={"previous_attempts": new_attempts})

    @staticmethod
    def _cancel_status(status: SubplanStatus) -> SubplanStatus:
        """Mark a subplan status as cancelled."""
        return status.model_copy(
            update={
                "status": ProcessingState.CANCELLED,
                "completed_at": datetime.now(tz=UTC),
            }
        )

    @staticmethod
    def _error_status(status: SubplanStatus, error: str) -> SubplanStatus:
        """Mark a subplan status as errored."""
        return status.model_copy(
            update={
                "status": ProcessingState.ERRORED,
                "completed_at": datetime.now(tz=UTC),
                "error": error,
            }
        )

    @staticmethod
    def _topological_sort(
        node_ids: list[str],
        graph: dict[str, list[str]],
    ) -> list[str]:
        """Topological sort of subplan IDs based on dependency graph.

        Args:
            node_ids: All subplan IDs to sort.
            graph: Maps subplan_id to list of subplan_ids it depends on.

        Returns:
            Topologically sorted list of subplan IDs.

        Raises:
            ValueError: If a cycle is detected.
        """
        in_degree: dict[str, int] = {nid: 0 for nid in node_ids}
        adjacency: dict[str, list[str]] = {nid: [] for nid in node_ids}
        node_set = set(node_ids)

        for nid in node_ids:
            deps = graph.get(nid, [])
            for dep in deps:
                if dep in node_set:
                    adjacency[dep].append(nid)
                    in_degree[nid] += 1

        queue: deque[str] = deque(nid for nid in node_ids if in_degree[nid] == 0)
        result: list[str] = []

        while queue:
            node = queue.popleft()
            result.append(node)
            for neighbor in adjacency[node]:
                in_degree[neighbor] -= 1
                if in_degree[neighbor] == 0:
                    queue.append(neighbor)

        if len(result) != len(node_ids):
            raise ValueError("Circular dependency detected in subplan dependency graph")

        return result
