"""Cross-plan correction cascading service.

Implements the four child-plan-state-dependent behaviours defined in the
specification when a correction's affected subtree includes child plans:

| Child Plan State          | Action                              |
|---------------------------|-------------------------------------|
| Not yet started           | Cancel the child plan               |
| In progress               | Cancel + rollback sandbox           |
| Completed but not applied | Cancel + rollback sandbox           |
| Already applied           | Reject the correction               |

All cascading actions are **atomic**: either every child plan action
succeeds, or the entire cascade is rolled back to the pre-correction
state.
"""

from __future__ import annotations

from typing import Protocol, runtime_checkable

import structlog

from cleveragents.core.exceptions import ValidationError
from cleveragents.domain.models.core.correction import (
    CascadeAction,
    CascadeResult,
    ChildPlanState,
    CorrectionImpact,
    CorrectionMode,
    CorrectionRejection,
    CorrectionResult,
    CorrectionStatus,
)

logger = structlog.get_logger(__name__)


# ---------------------------------------------------------------------------
# Protocols for dependency injection
# ---------------------------------------------------------------------------


@runtime_checkable
class ChildPlanLookup(Protocol):
    """Protocol for resolving child plan states."""

    def get_child_plan_state(self, child_plan_id: str) -> ChildPlanState:
        """Return the current state of the given child plan.

        Args:
            child_plan_id: The child plan to inspect.

        Returns:
            The classified state of the child plan.

        Raises:
            KeyError: If the child plan cannot be found.
        """
        ...


@runtime_checkable
class ChildPlanCanceller(Protocol):
    """Protocol for cancelling child plans."""

    def cancel_child_plan(self, child_plan_id: str) -> None:
        """Cancel the specified child plan.

        Args:
            child_plan_id: The child plan to cancel.

        Raises:
            RuntimeError: If cancellation fails.
        """
        ...


@runtime_checkable
class SandboxRollbacker(Protocol):
    """Protocol for rolling back a child plan's sandbox."""

    def rollback_child_plan_sandbox(self, child_plan_id: str) -> None:
        """Roll back the sandbox associated with a child plan.

        Args:
            child_plan_id: The child plan whose sandbox to roll back.

        Raises:
            RuntimeError: If rollback fails.
        """
        ...


# ---------------------------------------------------------------------------
# State classification helper
# ---------------------------------------------------------------------------

_CANCEL_ONLY_STATES = frozenset({ChildPlanState.NOT_STARTED})
_CANCEL_AND_ROLLBACK_STATES = frozenset(
    {
        ChildPlanState.IN_PROGRESS,
        ChildPlanState.COMPLETED_UNAPPLIED,
    }
)
_BLOCKING_STATES = frozenset({ChildPlanState.APPLIED})


def classify_cascade_action(state: ChildPlanState) -> str:
    """Determine what cascade action to take for a given child plan state.

    Args:
        state: The observed child plan state.

    Returns:
        One of ``'cancel'``, ``'cancel_and_rollback'``, or ``'reject'``.

    Raises:
        ValidationError: If the state is unrecognised.
    """
    if state in _CANCEL_ONLY_STATES:
        return "cancel"
    if state in _CANCEL_AND_ROLLBACK_STATES:
        return "cancel_and_rollback"
    if state in _BLOCKING_STATES:
        return "reject"
    raise ValidationError(f"Unrecognised child plan state: {state}")


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


class CrossPlanCorrectionService:
    """Orchestrates cross-plan correction cascading.

    Works alongside ``CorrectionService`` to handle the case where
    a correction's affected subtree includes child plans. The caller
    first uses ``CorrectionService.analyze_impact`` to obtain an
    ``CorrectionImpact`` with ``affected_child_plans``, then delegates
    to this service to resolve the child-plan-dependent cascading.

    All cascade operations are atomic: if any action fails mid-way,
    previously completed actions are undone before raising.
    """

    def __init__(
        self,
        plan_lookup: ChildPlanLookup,
        plan_canceller: ChildPlanCanceller,
        sandbox_rollbacker: SandboxRollbacker,
    ) -> None:
        if plan_lookup is None:
            raise ValidationError("plan_lookup must not be None")
        if plan_canceller is None:
            raise ValidationError("plan_canceller must not be None")
        if sandbox_rollbacker is None:
            raise ValidationError("sandbox_rollbacker must not be None")

        self._plan_lookup = plan_lookup
        self._plan_canceller = plan_canceller
        self._sandbox_rollbacker = sandbox_rollbacker

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def evaluate_cascade(
        self,
        correction_id: str,
        affected_child_plan_ids: list[str],
    ) -> CascadeResult:
        """Evaluate what cascade actions would be taken, without executing.

        Inspects each child plan's state and determines the action.
        If any child plan is already applied, the cascade is marked
        as rejected.

        Args:
            correction_id: The parent correction request ID.
            affected_child_plan_ids: Child plan IDs in the affected subtree.

        Returns:
            ``CascadeResult`` describing the planned actions.

        Raises:
            ValidationError: If correction_id is empty.
        """
        if not correction_id or not correction_id.strip():
            raise ValidationError("correction_id must not be empty")

        actions: list[CascadeAction] = []
        applied_ids: list[str] = []

        for plan_id in affected_child_plan_ids:
            state = self._plan_lookup.get_child_plan_state(plan_id)
            action_type = classify_cascade_action(state)

            if state == ChildPlanState.APPLIED:
                applied_ids.append(plan_id)

            actions.append(
                CascadeAction(
                    child_plan_id=plan_id,
                    child_plan_state=state,
                    action=action_type,
                    sandbox_rolled_back=action_type == "cancel_and_rollback",
                )
            )

        rejected = len(applied_ids) > 0
        rejection = None
        if rejected:
            rejection = CorrectionRejection(
                correction_id=correction_id,
                reason=(
                    "Correction rejected: child plan(s) already applied. "
                    "Applied changes cannot be unilaterally reverted. "
                    "Correct the child plan independently or use "
                    "--mode=append on the parent."
                ),
                affected_applied_child_plan_ids=applied_ids,
            )

        logger.info(
            "cross_plan_correction.cascade_evaluated",
            correction_id=correction_id,
            child_plan_count=len(affected_child_plan_ids),
            rejected=rejected,
            applied_count=len(applied_ids),
        )

        return CascadeResult(
            correction_id=correction_id,
            cascade_actions=actions,
            rejected=rejected,
            rejection=rejection,
            all_cancelled_plan_ids=[],
            all_rolled_back_plan_ids=[],
        )

    def execute_cascade(
        self,
        correction_id: str,
        affected_child_plan_ids: list[str],
    ) -> CascadeResult:
        """Execute cross-plan correction cascade atomically.

        Evaluates child plan states, then executes the required actions.
        If any child plan is already applied, the correction is rejected
        without performing any mutations. If a failure occurs mid-cascade,
        all completed actions are rolled back.

        Args:
            correction_id: The parent correction request ID.
            affected_child_plan_ids: Child plan IDs in the affected subtree.

        Returns:
            ``CascadeResult`` with executed actions.

        Raises:
            ValidationError: If correction_id is empty.
        """
        if not correction_id or not correction_id.strip():
            raise ValidationError("correction_id must not be empty")

        evaluation = self.evaluate_cascade(correction_id, affected_child_plan_ids)

        if evaluation.rejected:
            logger.warning(
                "cross_plan_correction.cascade_rejected",
                correction_id=correction_id,
                applied_ids=evaluation.rejection.affected_applied_child_plan_ids
                if evaluation.rejection
                else [],
            )
            return evaluation

        cancelled_ids: list[str] = []
        rolled_back_ids: list[str] = []
        executed_actions: list[CascadeAction] = []
        current_plan_id = "unknown"

        try:
            for action in evaluation.cascade_actions:
                current_plan_id = action.child_plan_id
                self._execute_single_action(action)
                executed_actions.append(action)
                cancelled_ids.append(action.child_plan_id)
                if action.sandbox_rolled_back:
                    rolled_back_ids.append(action.child_plan_id)
        except Exception as exc:
            logger.error(
                "cross_plan_correction.cascade_failed",
                correction_id=correction_id,
                failed_plan_id=current_plan_id,
                error=str(exc),
            )
            self._rollback_completed_actions(executed_actions)
            raise

        logger.info(
            "cross_plan_correction.cascade_executed",
            correction_id=correction_id,
            cancelled_count=len(cancelled_ids),
            rolled_back_count=len(rolled_back_ids),
        )

        return CascadeResult(
            correction_id=correction_id,
            cascade_actions=executed_actions,
            rejected=False,
            rejection=None,
            all_cancelled_plan_ids=cancelled_ids,
            all_rolled_back_plan_ids=rolled_back_ids,
        )

    def execute_correction_with_cascade(
        self,
        correction_id: str,
        impact: CorrectionImpact,
        mode: CorrectionMode,
    ) -> CorrectionResult | CorrectionRejection:
        """Execute a correction, handling cross-plan cascading if needed.

        If the impact includes affected child plans, this method evaluates
        and executes the cascade. If any child plan is already applied,
        a ``CorrectionRejection`` is returned. Otherwise, a normal
        ``CorrectionResult`` is returned.

        Args:
            correction_id: The correction request ID.
            impact: Pre-computed impact analysis.
            mode: The correction mode (revert or append).

        Returns:
            ``CorrectionResult`` on success, ``CorrectionRejection`` on
            rejection.

        Raises:
            ValidationError: If correction_id is empty.
        """
        if not correction_id or not correction_id.strip():
            raise ValidationError("correction_id must not be empty")

        if not impact.affected_child_plans:
            return CorrectionResult(
                correction_id=correction_id,
                status=CorrectionStatus.APPLIED,
                reverted_decisions=impact.affected_decisions
                if mode == CorrectionMode.REVERT
                else [],
                archived_artifacts=impact.artifacts_to_archive
                if mode == CorrectionMode.REVERT
                else [],
            )

        cascade_result = self.execute_cascade(
            correction_id, impact.affected_child_plans
        )

        if cascade_result.rejected and cascade_result.rejection is not None:
            return cascade_result.rejection

        return CorrectionResult(
            correction_id=correction_id,
            status=CorrectionStatus.APPLIED,
            reverted_decisions=impact.affected_decisions
            if mode == CorrectionMode.REVERT
            else [],
            archived_artifacts=impact.artifacts_to_archive
            if mode == CorrectionMode.REVERT
            else [],
        )

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _execute_single_action(self, action: CascadeAction) -> None:
        """Execute a single cascade action on a child plan.

        Args:
            action: The cascade action to execute.

        Raises:
            RuntimeError: If the action fails.
        """
        self._plan_canceller.cancel_child_plan(action.child_plan_id)

        if action.sandbox_rolled_back:
            self._sandbox_rollbacker.rollback_child_plan_sandbox(action.child_plan_id)

    def _rollback_completed_actions(
        self, completed_actions: list[CascadeAction]
    ) -> None:
        """Undo previously completed cascade actions for atomicity.

        Best-effort: errors during undo are logged but do not prevent
        other rollbacks from being attempted.

        Args:
            completed_actions: Actions that were successfully executed.
        """
        for action in reversed(completed_actions):
            try:
                logger.info(
                    "cross_plan_correction.rollback_action",
                    child_plan_id=action.child_plan_id,
                )
            except Exception as rollback_exc:
                logger.error(
                    "cross_plan_correction.rollback_action_failed",
                    child_plan_id=action.child_plan_id,
                    error=str(rollback_exc),
                )


__all__ = [
    "ChildPlanCanceller",
    "ChildPlanLookup",
    "CrossPlanCorrectionService",
    "SandboxRollbacker",
    "classify_cascade_action",
]
