"""Sandbox lifecycle manager for CleverAgents.

Manages the creation, tracking, and cleanup of sandbox instances across
plan executions.  Thread-safe via an internal reentrant lock.

Stage B3.7 of the implementation plan.
Updated in M6 to support sandbox boundary algebra (#548):
  - ``get_or_create_sandbox_for_resource`` keys by ``(plan_id, boundary_id)``
  - ``resolve_sandbox_key`` computes the boundary key for a resource
  - Boundary cache per plan execution
"""

from __future__ import annotations

import atexit
import contextlib
import logging
import threading
from collections.abc import Mapping
from datetime import datetime
from typing import cast

from cleveragents.domain.models.core.resource import Resource
from cleveragents.infrastructure.sandbox.boundary import (
    BoundaryCache,
    NoSandboxBoundaryError,
)
from cleveragents.infrastructure.sandbox.factory import (
    SandboxFactory,
    SandboxStrategyStr,
)
from cleveragents.infrastructure.sandbox.no_sandbox import NoSandbox
from cleveragents.infrastructure.sandbox.protocol import (
    AtomicCommitError,
    CommitResult,
    Sandbox,
    SandboxError,
    SandboxStatus,
)
from cleveragents.infrastructure.sandbox.transaction_sandbox import (
    TransactionSandbox,
)

logger = logging.getLogger(__name__)


class SandboxManager:
    """Manages sandbox lifecycles across plans.

    Provides lazy sandbox creation (only created when first needed),
    batch commit/rollback/cleanup operations, and automatic cleanup
    on process exit.

    Thread-safe: all mutable state is protected by ``_lock``.

    Usage::

        manager = SandboxManager(factory=SandboxFactory())
        sandbox = manager.get_or_create_sandbox(
            plan_id="01ARZ3...",
            resource_id="res-001",
            original_path="/path/to/repo",
            sandbox_strategy="git_worktree",
        )
        # ... use sandbox ...
        results = manager.commit_all("01ARZ3...")
        manager.cleanup_all("01ARZ3...")
    """

    def __init__(
        self,
        factory: SandboxFactory,
        cleanup_on_exit: bool = True,
    ) -> None:
        """Initialise the sandbox manager.

        Args:
            factory: Factory used to create new sandbox instances.
            cleanup_on_exit: If ``True`` (default), register an
                :func:`atexit` handler that cleans up all tracked
                sandboxes when the process exits.

        Raises:
            ValueError: If *factory* is ``None``.
        """
        if factory is None:
            raise ValueError("factory cannot be None")

        self._factory: SandboxFactory = factory
        self._active_sandboxes: dict[str, dict[str, Sandbox]] = {}
        self._lock: threading.RLock = threading.RLock()
        self._cleanup_on_exit: bool = cleanup_on_exit
        self._boundary_cache: BoundaryCache = BoundaryCache()

        if cleanup_on_exit:
            atexit.register(self._cleanup_on_exit_handler)

    # -- primary API ---------------------------------------------------------

    def get_or_create_sandbox(
        self,
        plan_id: str,
        resource_id: str,
        original_path: str,
        sandbox_strategy: SandboxStrategyStr,
    ) -> Sandbox:
        """Return an existing sandbox or create a new one (lazy pattern).

        If a sandbox already exists for the given ``(plan_id, resource_id)``
        pair and is in a usable status (``CREATED``, ``ACTIVE``, or
        ``ROLLED_BACK``), it is returned directly.  Otherwise a new sandbox
        is created via the factory and initialised.

        Args:
            plan_id: Identifier of the plan requesting the sandbox.
            resource_id: Identifier of the resource to sandbox.
            original_path: Filesystem path or location of the resource.
            sandbox_strategy: Which sandboxing strategy to use.

        Returns:
            A :class:`Sandbox` instance ready for use.

        Raises:
            ValueError: If any required argument is empty.
            SandboxError: If sandbox creation or initialisation fails.
        """
        if not plan_id:
            raise ValueError("plan_id cannot be empty")
        if not resource_id:
            raise ValueError("resource_id cannot be empty")
        if not original_path:
            raise ValueError("original_path cannot be empty")

        with self._lock:
            # Check for existing usable sandbox
            plan_sandboxes = self._active_sandboxes.get(plan_id, {})
            existing = plan_sandboxes.get(resource_id)

            if existing is not None:
                if existing.status in (
                    SandboxStatus.CREATED,
                    SandboxStatus.ACTIVE,
                    SandboxStatus.ROLLED_BACK,
                ):
                    return existing

                # Existing sandbox is in a terminal/unusable state -- remove
                logger.debug(
                    "Removing stale sandbox for plan=%s resource=%s (status=%s)",
                    plan_id,
                    resource_id,
                    existing.status.value,
                )
                plan_sandboxes.pop(resource_id, None)

            # Create new sandbox via factory
            sandbox = self._factory.create_sandbox(
                resource_id=resource_id,
                original_path=original_path,
                sandbox_strategy=sandbox_strategy,
            )

            # Initialise sandbox
            sandbox.create(plan_id)

            # Track it
            if plan_id not in self._active_sandboxes:
                self._active_sandboxes[plan_id] = {}
            self._active_sandboxes[plan_id][resource_id] = sandbox

            logger.info(
                "Created sandbox for plan=%s resource=%s strategy=%s (sandbox_id=%s)",
                plan_id,
                resource_id,
                sandbox_strategy,
                sandbox.sandbox_id,
            )

            return sandbox

    def get_sandbox(self, plan_id: str, resource_id: str) -> Sandbox | None:
        """Look up an existing sandbox without creating a new one.

        Args:
            plan_id: Plan identifier.
            resource_id: Resource identifier.

        Returns:
            The tracked sandbox, or ``None`` if none exists.
        """
        with self._lock:
            return self._active_sandboxes.get(plan_id, {}).get(resource_id)

    def list_sandboxes(self, plan_id: str) -> list[Sandbox]:
        """Return all sandboxes tracked for a plan.

        Args:
            plan_id: Plan identifier.

        Returns:
            List of sandboxes (may include cleaned-up or errored ones).
        """
        with self._lock:
            return list(self._active_sandboxes.get(plan_id, {}).values())

    # -- batch operations ----------------------------------------------------

    def commit_all(self, plan_id: str) -> list[CommitResult]:
        """Atomically commit all active sandboxes for a plan.

        The Apply phase is an atomic operation — either all sandbox
        changes are committed successfully, or none are.  If any single
        sandbox commit fails, all previously-committed sandboxes in the
        same batch are rolled back.

        .. note:: Specification contradiction

           Line 45938 requires atomicity ("all or none"), while line
           19193 states "partial apply" is possible.  This method
           implements the atomicity requirement from line 45938 per
           issue #925.  A spec update for line 19193 may be needed.

        Args:
            plan_id: Plan identifier.

        Returns:
            On success, a list of :class:`CommitResult` (one per committed
            sandbox, all with ``success=True``).

            On failure, a single-element list containing a
            :class:`CommitResult` with ``success=False`` whose ``error``
            field describes which sandbox failed and which sandboxes
            were rolled back.  The ``metadata`` dict includes a
            ``"rolled_back"`` key listing the IDs of rolled-back
            sandboxes.

        Raises:
            ValueError: If *plan_id* is empty.
            AtomicCommitError: If a sandbox ``commit()`` raises a non-
                :class:`SandboxError` exception.  The atomic rollback is
                performed and the original exception is wrapped in an
                :class:`AtomicCommitError` (chained as ``__cause__``)
                that carries ``rolled_back_ids`` and
                ``failed_rollback_ids`` attributes.  The caller receives
                the :class:`AtomicCommitError`, not the raw exception.

        .. warning:: Thread safety

           This method is **not safe** for concurrent calls on the same
           *plan_id*.  The internal lock serializes access to the
           sandbox registry, but individual sandbox ``commit()`` and
           ``rollback()`` calls run outside the lock to avoid blocking
           all manager operations during potentially slow I/O.
           Callers must ensure that ``commit_all`` is not invoked
           concurrently for the same plan.
        """
        if not plan_id:
            raise ValueError("plan_id cannot be empty")

        with self._lock:
            sandboxes = list(self._active_sandboxes.get(plan_id, {}).values())

        committable = [
            sb
            for sb in sandboxes
            if sb.status in (SandboxStatus.CREATED, SandboxStatus.ACTIVE)
        ]

        if not committable:
            return []

        # Warn about sandboxes that cannot be rolled back (e.g. NoSandbox,
        # TransactionSandbox after COMMIT).  These break the all-or-nothing
        # guarantee if a later commit in the batch fails.
        non_rollbackable: list[Sandbox] = []
        rollbackable: list[Sandbox] = []
        for sb in committable:
            if isinstance(sb, NoSandbox):
                logger.warning(
                    "Sandbox %s (resource strategy 'none') cannot be rolled "
                    "back after commit.  Atomicity is broken for plan %s "
                    "because changes are applied in-place immediately.",
                    sb.sandbox_id,
                    plan_id,
                )
                non_rollbackable.append(sb)
            elif isinstance(sb, TransactionSandbox):
                logger.warning(
                    "Sandbox %s (transaction_rollback strategy) cannot be "
                    "rolled back after database COMMIT.  Atomicity is "
                    "broken for plan %s because database changes are "
                    "permanent once committed.",
                    sb.sandbox_id,
                    plan_id,
                )
                non_rollbackable.append(sb)
            else:
                rollbackable.append(sb)

        # Commit rollbackable sandboxes first so that if any fail, we
        # can undo them.  Non-rollbackable sandboxes commit last — they
        # only run after all rollbackable sandboxes succeed.
        ordered = rollbackable + non_rollbackable

        committed: list[tuple[Sandbox, CommitResult]] = []

        for sandbox in ordered:
            try:
                result = sandbox.commit()
                committed.append((sandbox, result))
            except Exception as exc:
                # Atomic rollback: undo all previously-committed sandboxes.
                # Catches Exception (not just SandboxError) so that
                # unexpected errors cannot bypass the rollback and leave
                # already-committed sandboxes in an inconsistent state.
                failed_id = sandbox.sandbox_id
                rolled_back_ids, failed_rollback_ids = self._rollback_committed(
                    committed, plan_id
                )

                error_msg = f"Atomic commit failed at sandbox {failed_id}: {exc}"
                if rolled_back_ids:
                    error_msg += f"; rolled back sandboxes: {rolled_back_ids}"
                if failed_rollback_ids:
                    error_msg += (
                        f"; FAILED to roll back sandboxes: {failed_rollback_ids}"
                    )

                logger.error(
                    "Atomic commit_all failed for plan %s: %s",
                    plan_id,
                    error_msg,
                )

                # For non-SandboxError exceptions, wrap in
                # AtomicCommitError carrying rollback metadata so the
                # caller can determine which sandboxes were rolled back.
                # The original exception is chained as __cause__.
                if not isinstance(exc, SandboxError):
                    raise AtomicCommitError(
                        error_msg,
                        rolled_back_ids=rolled_back_ids,
                        failed_rollback_ids=failed_rollback_ids,
                    ) from exc

                return [
                    CommitResult(
                        sandbox_id=failed_id,
                        success=False,
                        error=error_msg,
                        timestamp=datetime.now(),
                        metadata={
                            "rolled_back": rolled_back_ids,
                            "rollback_failed": failed_rollback_ids,
                        },
                    )
                ]

        return [result for _, result in committed]

    def _rollback_committed(
        self,
        committed: list[tuple[Sandbox, CommitResult]],
        plan_id: str,
    ) -> tuple[list[str], list[str]]:
        """Roll back sandboxes that were already committed in this batch.

        Called when an atomic ``commit_all`` encounters a failure
        partway through.  Each previously-committed sandbox is rolled
        back to undo its changes.  Errors during individual rollbacks
        are logged but do not prevent other rollbacks from being
        attempted.

        Rollback proceeds in **LIFO** (reverse-of-commit) order,
        following the standard transaction-log undo pattern: the most
        recently committed sandbox is rolled back first.

        .. note::

           The specification's "top-down" rollback ordering (line 24632)
           refers to resource **DAG traversal** (parent before child)
           within a single sandbox domain.  LIFO ordering here applies
           to the *batch undo* of independently committed sandboxes,
           which is a different context.

        Args:
            committed: List of ``(sandbox, commit_result)`` tuples for
                sandboxes that committed successfully before the failure.
            plan_id: Plan identifier (used for log messages).

        Returns:
            A tuple of ``(rolled_back_ids, failed_rollback_ids)`` where
            *rolled_back_ids* lists sandboxes that were successfully
            rolled back and *failed_rollback_ids* lists sandboxes where
            the rollback raised an exception.
        """
        rolled_back_ids: list[str] = []
        failed_rollback_ids: list[str] = []

        for sandbox, _result in reversed(committed):
            try:
                sandbox.rollback()
                rolled_back_ids.append(sandbox.sandbox_id)
                logger.info(
                    "Rolled back sandbox %s for plan %s during atomic commit recovery",
                    sandbox.sandbox_id,
                    plan_id,
                )
            except Exception as rb_exc:
                # Catches Exception (not just SandboxError) so that
                # unexpected rollback errors do not prevent remaining
                # rollbacks from being attempted.
                failed_rollback_ids.append(sandbox.sandbox_id)
                logger.error(
                    "Failed to rollback sandbox %s for plan %s "
                    "during atomic commit recovery: %s",
                    sandbox.sandbox_id,
                    plan_id,
                    rb_exc,
                )

        return rolled_back_ids, failed_rollback_ids

    def rollback_all(self, plan_id: str) -> None:
        """Roll back all active sandboxes for a plan.

        Errors during individual rollbacks are logged but do not prevent
        other sandboxes from being rolled back.

        Args:
            plan_id: Plan identifier.
        """
        if not plan_id:
            raise ValueError("plan_id cannot be empty")

        with self._lock:
            sandboxes = list(self._active_sandboxes.get(plan_id, {}).values())

        for sandbox in sandboxes:
            if sandbox.status not in (SandboxStatus.ACTIVE, SandboxStatus.COMMITTED):
                continue

            try:
                sandbox.rollback()
            except Exception as exc:
                # Catches Exception (not just SandboxError) so that
                # unexpected rollback errors do not prevent remaining
                # sandboxes from being rolled back, consistent with
                # _rollback_committed().
                logger.error(
                    "Failed to rollback sandbox %s for plan %s: %s",
                    sandbox.sandbox_id,
                    plan_id,
                    exc,
                )

    def cleanup_all(self, plan_id: str) -> None:
        """Clean up all sandboxes for a plan and remove tracking.

        Args:
            plan_id: Plan identifier.
        """
        if not plan_id:
            raise ValueError("plan_id cannot be empty")

        with self._lock:
            sandboxes = list(self._active_sandboxes.get(plan_id, {}).values())

        for sandbox in sandboxes:
            try:
                sandbox.cleanup()
            except Exception as exc:
                # Catches Exception (not just SandboxError) so that
                # unexpected errors (e.g. raw OSError, PermissionError)
                # do not abort the loop and prevent remaining sandboxes
                # from being cleaned up, consistent with
                # _rollback_committed() and rollback_all().
                logger.error(
                    "Failed to cleanup sandbox %s for plan %s: %s",
                    sandbox.sandbox_id,
                    plan_id,
                    exc,
                )

        with self._lock:
            self._active_sandboxes.pop(plan_id, None)

    def cleanup_abandoned(self) -> int:
        """Clean up sandboxes in terminal/errored states.

        Iterates over all tracked sandboxes and cleans up any that are in
        ``ERRORED``, ``COMMITTED``, or ``ROLLED_BACK`` status (i.e.
        usable work is done but artefacts remain).

        Returns:
            Count of sandboxes that were cleaned up.
        """
        cleaned = 0

        with self._lock:
            plan_ids = list(self._active_sandboxes.keys())

        for plan_id in plan_ids:
            with self._lock:
                sandboxes = dict(self._active_sandboxes.get(plan_id, {}))

            for _resource_id, sandbox in sandboxes.items():
                if sandbox.status in (
                    SandboxStatus.ERRORED,
                    SandboxStatus.COMMITTED,
                    SandboxStatus.ROLLED_BACK,
                ):
                    try:
                        sandbox.cleanup()
                        cleaned += 1
                    except Exception as exc:
                        # Catches Exception (not just SandboxError) so
                        # that unexpected errors (e.g. raw OSError,
                        # PermissionError) do not crash the loop and
                        # prevent remaining abandoned sandboxes from
                        # being cleaned up, consistent with cleanup_all,
                        # rollback_all, and _rollback_committed.
                        logger.error(
                            "Failed to cleanup abandoned sandbox %s: %s",
                            sandbox.sandbox_id,
                            exc,
                        )

            # Remove plan entry if all sandboxes are cleaned up
            with self._lock:
                remaining = self._active_sandboxes.get(plan_id, {})
                all_cleaned = all(
                    s.status == SandboxStatus.CLEANED_UP for s in remaining.values()
                )
                if all_cleaned and remaining:
                    self._active_sandboxes.pop(plan_id, None)

        if cleaned > 0:
            logger.info("Cleaned up %d abandoned sandbox(es)", cleaned)

        return cleaned

    # -- boundary-aware API (M6 #548) ----------------------------------------

    def resolve_sandbox_key(
        self,
        resource: Resource,
        resource_registry: Mapping[str, Resource],
    ) -> str:
        """Compute the sandbox key for a resource via boundary algebra.

        The key is the ``resource_id`` of the resource's sandbox boundary.
        If the resource itself is a boundary, its own ID is returned.

        Args:
            resource: The resource to resolve.
            resource_registry: Full resource registry for DAG traversal.

        Returns:
            The ``resource_id`` of the sandbox boundary.

        Raises:
            ValueError: If *resource* or *resource_registry* is ``None``.
            NoSandboxBoundaryError: If no sandbox boundary exists
                in the resource's ancestor chain.
        """
        if resource is None:
            raise ValueError("resource cannot be None")
        if resource_registry is None:
            raise ValueError("resource_registry cannot be None")

        boundary = self._boundary_cache.get_boundary(resource, resource_registry)
        if boundary is None:
            raise NoSandboxBoundaryError(
                f"Resource {resource.resource_id} has no sandboxable ancestor"
            )
        return boundary.resource_id

    def get_or_create_sandbox_for_resource(
        self,
        plan_id: str,
        resource: Resource,
        resource_registry: Mapping[str, Resource],
    ) -> Sandbox:
        """Return a sandbox for the resource, keyed by sandbox boundary.

        Uses the boundary algebra to determine which sandbox boundary
        governs this resource, then delegates to
        :meth:`get_or_create_sandbox` with the boundary's ID and
        properties.  Multiple resources sharing the same boundary will
        share the same sandbox instance.

        Args:
            plan_id: Identifier of the plan requesting the sandbox.
            resource: The resource that needs sandboxing.
            resource_registry: Full resource registry for DAG traversal.

        Returns:
            A :class:`Sandbox` instance for the resource's boundary.

        Raises:
            ValueError: If any required argument is empty/``None``.
            NoSandboxBoundaryError: If no sandbox boundary exists.
            SandboxError: If sandbox creation or initialisation fails.
        """
        if not plan_id:
            raise ValueError("plan_id cannot be empty")
        if resource is None:
            raise ValueError("resource cannot be None")
        if resource_registry is None:
            raise ValueError("resource_registry cannot be None")

        boundary_id = self.resolve_sandbox_key(resource, resource_registry)
        boundary_resource = resource_registry[boundary_id]

        # Determine strategy from the boundary resource.
        strategy: SandboxStrategyStr = "none"
        if boundary_resource.sandbox_strategy is not None:
            strategy = cast("SandboxStrategyStr", boundary_resource.sandbox_strategy)

        # Determine location from the boundary resource.
        original_path = boundary_resource.location or ""
        if not original_path:
            raise ValueError(f"Boundary resource {boundary_id} has no location")

        return self.get_or_create_sandbox(
            plan_id=plan_id,
            resource_id=boundary_id,
            original_path=original_path,
            sandbox_strategy=strategy,
        )

    def clear_boundary_cache(self) -> None:
        """Clear the boundary cache.

        Should be called at the start of each plan execution or when
        the resource DAG is modified.
        """
        self._boundary_cache.clear()

    @property
    def boundary_cache_size(self) -> int:
        """Return the number of cached boundary lookups."""
        return self._boundary_cache.size

    # -- internal ------------------------------------------------------------

    def _cleanup_on_exit_handler(self) -> None:
        """Atexit handler: clean up all tracked sandboxes."""
        with self._lock:
            plan_ids = list(self._active_sandboxes.keys())

        for plan_id in plan_ids:
            with contextlib.suppress(Exception):
                self.cleanup_all(plan_id)
