"""Overlay filesystem sandbox — OverlayFS with userspace fallback.

Creates upper/work/merged layers.  Commit syncs to original; rollback
discards and recreates.  For fs-mount, fs-directory, fs-file resources.
"""

from __future__ import annotations

import functools
import logging
import os
import shutil
import subprocess
import tempfile
from datetime import datetime

from ulid import ULID

from cleveragents.infrastructure.sandbox._fs_utils import (
    backup_directory,
    compute_diff,
    safe_restore,
)
from cleveragents.infrastructure.sandbox.protocol import (
    CommitResult,
    SandboxCommitError,
    SandboxContext,
    SandboxCreationError,
    SandboxRollbackError,
    SandboxStateError,
    SandboxStatus,
)

logger = logging.getLogger(__name__)


@functools.lru_cache(maxsize=1)
def _is_overlayfs_available() -> bool:
    """Return True only if we can actually mount OverlayFS.

    Uses a real mount probe — ``mount --fake`` gives false positives
    in Docker containers that lack CAP_SYS_ADMIN.
    """
    try:
        proc_fs = "/proc/filesystems"
        if os.path.isfile(proc_fs):
            with open(proc_fs) as fh:
                if "overlay" not in fh.read():
                    return False
        else:
            return False
    except OSError:
        return False

    if os.geteuid() != 0:
        return False
    # Real mount probe — attempt actual overlay mount in a temp dir.
    probe_dir = None
    try:
        probe_dir = tempfile.mkdtemp(prefix="ca-overlay-probe-")
        lower, upper = os.path.join(probe_dir, "l"), os.path.join(probe_dir, "u")
        work, merged = os.path.join(probe_dir, "w"), os.path.join(probe_dir, "m")
        for d in (lower, upper, work, merged):
            os.makedirs(d)
        opts = f"lowerdir={lower},upperdir={upper},workdir={work}"
        rc = subprocess.run(
            ["mount", "-t", "overlay", "overlay", "-o", opts, merged],
            capture_output=True,
            timeout=5,
        ).returncode
        if rc != 0:
            return False
        subprocess.run(["umount", merged], capture_output=True, timeout=5)
        return True
    except (OSError, subprocess.TimeoutExpired):
        return False
    finally:
        if probe_dir is not None:
            shutil.rmtree(probe_dir, ignore_errors=True)


class OverlaySandbox:
    """Sandbox that isolates changes using an overlay filesystem model.

    Mounts a real OverlayFS union when available, otherwise falls back
    to a userspace copy-based overlay.

    Implements the
    :class:`~cleveragents.infrastructure.sandbox.protocol.Sandbox`
    protocol.
    """

    def __init__(
        self,
        resource_id: str,
        original_path: str,
    ) -> None:
        """Initialise an overlay sandbox.

        Args:
            resource_id: Identifier of the resource being sandboxed.
            original_path: Path to the directory to sandbox.

        Raises:
            ValueError: If *resource_id* or *original_path* is empty.
        """
        if not resource_id:
            raise ValueError("resource_id cannot be empty")
        if not original_path:
            raise ValueError("original_path cannot be empty")

        self._sandbox_id: str = str(ULID())
        self._resource_id: str = resource_id
        self._original_path: str = os.path.abspath(original_path)
        self._status: SandboxStatus = SandboxStatus.PENDING
        self._context: SandboxContext | None = None

        # Layer directories — set after create()
        self._base_dir: str | None = None
        self._upper_dir: str | None = None
        self._work_dir: str | None = None
        self._merged_dir: str | None = None
        # Pre-commit backup for atomic rollback support
        self._pre_commit_backup: str | None = None

        self._use_real_overlay: bool = _is_overlayfs_available()

    # -- protocol properties -------------------------------------------------
    @property
    def sandbox_id(self) -> str:
        """Unique identifier for this sandbox instance."""
        return self._sandbox_id

    @property
    def status(self) -> SandboxStatus:
        """Current lifecycle status."""
        return self._status

    @property
    def context(self) -> SandboxContext | None:
        """Context after creation, ``None`` before ``create``."""
        return self._context

    # -- protocol methods ----------------------------------------------------
    def create(self, plan_id: str) -> SandboxContext:
        """Set up overlay directory layers and populate the merged view.

        Args:
            plan_id: The plan that owns this sandbox.

        Returns:
            A :class:`SandboxContext` with the merged layer path.

        Raises:
            ValueError: If *plan_id* is empty.
            SandboxStateError: If not in ``PENDING`` status.
            SandboxCreationError: If the directory setup fails.
        """
        if not plan_id:
            raise ValueError("plan_id cannot be empty")

        SandboxStatus.assert_transition(self._status, SandboxStatus.CREATED)

        try:
            if not os.path.isdir(self._original_path):
                raise SandboxCreationError(
                    f"Original path is not a directory: {self._original_path}"
                )

            self._base_dir = tempfile.mkdtemp(prefix="ca-overlay-sandbox-")
            self._upper_dir = os.path.join(self._base_dir, "upper")
            self._work_dir = os.path.join(self._base_dir, "work")
            self._merged_dir = os.path.join(self._base_dir, "merged")

            os.makedirs(self._upper_dir, exist_ok=True)
            os.makedirs(self._work_dir, exist_ok=True)
            os.makedirs(self._merged_dir, exist_ok=True)

            if self._use_real_overlay:
                self._mount_overlay()
            else:
                # Userspace fallback: copy original into merged
                shutil.copytree(
                    self._original_path,
                    self._merged_dir,
                    symlinks=False,
                    dirs_exist_ok=True,
                )

        except SandboxCreationError:
            self._status = SandboxStatus.ERRORED
            raise
        except OSError as exc:
            self._status = SandboxStatus.ERRORED
            raise SandboxCreationError(
                f"Failed to create overlay sandbox for resource "
                f"{self._resource_id}: {exc}"
            ) from exc

        self._context = SandboxContext(
            sandbox_id=self._sandbox_id,
            sandbox_path=self._merged_dir,
            original_path=self._original_path,
            resource_id=self._resource_id,
            plan_id=plan_id,
            created_at=datetime.now(),
            metadata={
                "strategy": "overlay",
                "sandbox_path": self._merged_dir,
                "use_real_overlay": self._use_real_overlay,
            },
        )
        self._status = SandboxStatus.CREATED

        logger.info(
            "Created overlay sandbox: plan=%s resource=%s merged=%s real_overlay=%s",
            plan_id,
            self._resource_id,
            self._merged_dir,
            self._use_real_overlay,
        )

        return self._context

    def get_path(self, resource_path: str) -> str:
        """Translate a resource-relative path to the merged layer.

        Args:
            resource_path: Path relative to the resource root.

        Returns:
            Absolute path inside the merged directory.

        Raises:
            SandboxStateError: If sandbox is not in a usable status.
            ValueError: If *resource_path* attempts directory traversal.
        """
        if self._status not in (
            SandboxStatus.CREATED,
            SandboxStatus.ACTIVE,
            SandboxStatus.ROLLED_BACK,
        ):
            raise SandboxStateError(
                f"Cannot resolve path in status {self._status.value}"
            )

        if ".." in resource_path.split("/"):
            raise ValueError(f"Path traversal not allowed: {resource_path}")

        if self._status in (SandboxStatus.CREATED, SandboxStatus.ROLLED_BACK):
            self._status = SandboxStatus.ACTIVE

        if self._merged_dir is None:
            raise SandboxStateError("Sandbox merged directory not set")

        return os.path.join(self._merged_dir, resource_path)

    def commit(self, message: str | None = None) -> CommitResult:
        """Apply overlay changes back to the original directory.

        For real OverlayFS: copies the upper layer contents to original.
        For userspace fallback: diffs merged against original.

        Args:
            message: Optional log message (stored in result metadata).

        Returns:
            A :class:`CommitResult` describing the outcome.

        Raises:
            SandboxCommitError: If the sync fails.
            SandboxStateError: If sandbox is not in a committable status.
        """
        if self._status not in (
            SandboxStatus.CREATED,
            SandboxStatus.ACTIVE,
        ):
            raise SandboxStateError(f"Cannot commit from status {self._status.value}")

        SandboxStatus.assert_transition(self._status, SandboxStatus.COMMITTED)

        if self._merged_dir is None:
            raise SandboxStateError("Sandbox merged directory not set")

        try:
            changed_files, added_files, deleted_files = compute_diff(
                self._merged_dir, self._original_path
            )

            # Save pre-commit backup for atomic rollback support.
            # Skip when there are no changes to avoid copying a potentially
            # large directory tree unnecessarily.
            #
            # IMPORTANT: ``_pre_commit_backup`` is assigned only AFTER
            # ``backup_directory`` succeeds.  If the backup itself fails
            # (e.g. disk full), the original directory is still intact and
            # the ``except OSError`` handler must NOT attempt
            # ``safe_restore`` with a partial backup.
            if changed_files or added_files or deleted_files:
                # Create backup on the same filesystem as the original to
                # avoid cross-device copy overhead and ensure os.rename
                # compatibility in safe_restore().
                backup_dir = tempfile.mkdtemp(
                    prefix="ca-overlay-backup-",
                    dir=os.path.dirname(self._original_path),
                )
                try:
                    backup_directory(self._original_path, backup_dir)
                except Exception:
                    shutil.rmtree(backup_dir, ignore_errors=True)
                    raise
                self._pre_commit_backup = backup_dir

            # Apply changes: copy modified and new files
            for rel_path in changed_files + added_files:
                src = os.path.join(self._merged_dir, rel_path)
                dst = os.path.join(self._original_path, rel_path)
                dst_dir = os.path.dirname(dst)
                if dst_dir:
                    os.makedirs(dst_dir, exist_ok=True)
                shutil.copy2(src, dst)

            # Apply deletions
            for rel_path in deleted_files:
                dst = os.path.join(self._original_path, rel_path)
                if os.path.exists(dst):
                    os.remove(dst)

        except Exception as exc:
            # Attempt to restore the original from the pre-commit backup
            # so that a partial copy does not leave the original corrupted.
            # Catches Exception (not just OSError) so that unexpected errors
            # during the file-copy phase (e.g. MemoryError wrapping,
            # TypeError) also trigger the restore.  Without this, the
            # backup would exist but never be applied, leaving the original
            # directory partially modified.
            # If the restore itself fails, keep the backup for manual
            # recovery — cleanup() will remove it later.
            if self._pre_commit_backup is not None:
                try:
                    safe_restore(self._pre_commit_backup, self._original_path)
                    self._pre_commit_backup = None
                except Exception:
                    logger.warning(
                        "Failed to restore original from backup for sandbox %s; "
                        "backup preserved at %s for manual recovery",
                        self._sandbox_id,
                        self._pre_commit_backup,
                    )
            self._status = SandboxStatus.ERRORED
            raise SandboxCommitError(
                f"Failed to commit overlay sandbox {self._sandbox_id}: {exc}"
            ) from exc

        self._status = SandboxStatus.COMMITTED

        logger.info(
            "Committed overlay sandbox: sandbox_id=%s changed=%d added=%d deleted=%d",
            self._sandbox_id,
            len(changed_files),
            len(added_files),
            len(deleted_files),
        )

        metadata: dict[str, object] = {}
        if message:
            metadata["message"] = message

        return CommitResult(
            sandbox_id=self._sandbox_id,
            success=True,
            commit_ref=None,
            changed_files=changed_files,
            added_files=added_files,
            deleted_files=deleted_files,
            error=None,
            timestamp=datetime.now(),
            metadata=metadata,
        )

    def rollback(self) -> None:
        """Discard overlay changes by resetting the merged view.

        When rolling back from ``ACTIVE``: resets the merged view.
        When rolling back from ``COMMITTED``: restores the original
        directory from the pre-commit backup so that the commit's
        effects are fully undone.

        For real OverlayFS: unmounts, recreates upper/work, remounts.
        For userspace fallback: deletes merged, re-copies original.

        Raises:
            SandboxRollbackError: If the rollback fails.
            SandboxStateError: If called in an invalid status.
        """
        if self._status not in (SandboxStatus.ACTIVE, SandboxStatus.COMMITTED):
            raise SandboxStateError(f"Cannot rollback from status {self._status.value}")

        SandboxStatus.assert_transition(self._status, SandboxStatus.ROLLED_BACK)

        if self._merged_dir is None:
            raise SandboxStateError("Sandbox merged directory not set")

        try:
            if self._status == SandboxStatus.COMMITTED:
                if self._pre_commit_backup is None:
                    # No pre-commit backup means commit() detected no
                    # changes and skipped the backup.  The original
                    # directory is unmodified, so rollback is a no-op.
                    # The merged directory is also unchanged, so no
                    # reset is needed.
                    logger.debug(
                        "Rollback from COMMITTED with no backup (no changes "
                        "were applied) for sandbox %s — no-op",
                        self._sandbox_id,
                    )
                else:
                    # Atomically restore original from pre-commit backup.
                    # Uses rename-based swap so that a failure during copy
                    # does not leave the original directory deleted.
                    safe_restore(self._pre_commit_backup, self._original_path)
                    self._pre_commit_backup = None

                    # Reset the merged directory so that a subsequent
                    # get_path() (ROLLED_BACK → ACTIVE) does not expose
                    # stale pre-rollback content.  Only needed when
                    # changes were actually restored (backup existed).
                    if self._merged_dir is not None:
                        if self._use_real_overlay:
                            # For real OverlayFS: unmount, clean upper/work
                            # dirs, and remount so the overlay is correctly
                            # re-established for potential re-activation.
                            # If unmount fails, bail out — attempting to
                            # mount on top of a still-mounted FS would cause
                            # a double-mount or mount error.
                            try:
                                subprocess.run(
                                    ["umount", str(self._merged_dir)],
                                    capture_output=True,
                                    check=True,
                                    timeout=30,
                                )
                            except (
                                subprocess.CalledProcessError,
                                subprocess.TimeoutExpired,
                                OSError,
                            ) as umount_exc:
                                raise SandboxRollbackError(
                                    f"Cannot remount overlay for sandbox "
                                    f"{self._sandbox_id}: unmount failed: "
                                    f"{umount_exc}"
                                ) from umount_exc
                            if self._upper_dir is not None:
                                shutil.rmtree(self._upper_dir, ignore_errors=True)
                                os.makedirs(self._upper_dir, exist_ok=True)
                            if self._work_dir is not None:
                                shutil.rmtree(self._work_dir, ignore_errors=True)
                                os.makedirs(self._work_dir, exist_ok=True)
                            self._mount_overlay()
                        else:
                            # Userspace fallback: delete merged, re-copy from
                            # the (now-restored) original.
                            shutil.rmtree(self._merged_dir, ignore_errors=True)
                            shutil.copytree(
                                self._original_path,
                                self._merged_dir,
                                symlinks=False,
                                dirs_exist_ok=True,
                            )
            elif self._use_real_overlay:
                self._unmount_overlay()
                # Reset upper and work directories
                if self._upper_dir is not None:
                    shutil.rmtree(self._upper_dir, ignore_errors=True)
                    os.makedirs(self._upper_dir, exist_ok=True)
                if self._work_dir is not None:
                    shutil.rmtree(self._work_dir, ignore_errors=True)
                    os.makedirs(self._work_dir, exist_ok=True)
                self._mount_overlay()
            else:
                # Userspace fallback: delete merged, re-copy.
                # Use dirs_exist_ok=True because rmtree with
                # ignore_errors=True may silently leave the directory
                # partially intact (e.g. permission error on a file),
                # which would cause copytree to raise FileExistsError
                # with dirs_exist_ok=False.
                shutil.rmtree(self._merged_dir, ignore_errors=True)
                shutil.copytree(
                    self._original_path,
                    self._merged_dir,
                    symlinks=False,
                    dirs_exist_ok=True,
                )

        except SandboxRollbackError:
            # Already a SandboxRollbackError (e.g. from unmount failure
            # during COMMITTED rollback) — re-raise without wrapping.
            self._status = SandboxStatus.ERRORED
            raise
        except Exception as exc:
            self._status = SandboxStatus.ERRORED
            raise SandboxRollbackError(
                f"Failed to rollback overlay sandbox {self._sandbox_id}: {exc}"
            ) from exc

        self._status = SandboxStatus.ROLLED_BACK

        logger.info(
            "Rolled back overlay sandbox: sandbox_id=%s",
            self._sandbox_id,
        )

    def cleanup(self) -> None:
        """Remove all overlay layers and mountpoints.

        Idempotent — safe to call multiple times.

        Raises:
            SandboxError: On unexpected errors during cleanup.
        """
        if self._status == SandboxStatus.CLEANED_UP:
            return

        logger.debug(
            "Cleaning up overlay sandbox: sandbox_id=%s base=%s",
            self._sandbox_id,
            self._base_dir,
        )

        # Unmount if real overlay was used
        if self._use_real_overlay and self._merged_dir is not None:
            self._unmount_overlay()

        if self._base_dir is not None and os.path.exists(self._base_dir):
            shutil.rmtree(self._base_dir, ignore_errors=True)

        # Clean up any pre-commit backup
        if self._pre_commit_backup is not None:
            if os.path.exists(self._pre_commit_backup):
                shutil.rmtree(self._pre_commit_backup, ignore_errors=True)
            self._pre_commit_backup = None

        self._status = SandboxStatus.CLEANED_UP

        logger.info(
            "Cleaned up overlay sandbox: sandbox_id=%s",
            self._sandbox_id,
        )

    # -- internal helpers ----------------------------------------------------
    def _mount_overlay(self) -> None:
        """Mount an OverlayFS union at the merged directory.

        Raises:
            SandboxCreationError: If the mount command fails or if any
                overlay path contains a comma (which would corrupt the
                mount options string).
        """
        # OverlayFS mount options are comma-separated key=value pairs.
        # A comma inside any path value would be mis-parsed by the
        # kernel's filesystem driver, corrupting the mount options.
        for label, path in (
            ("lowerdir", self._original_path),
            ("upperdir", self._upper_dir),
            ("workdir", self._work_dir),
        ):
            if path and "," in str(path):
                raise SandboxCreationError(
                    f"OverlayFS {label} path contains a comma, which "
                    f"would corrupt mount options: {path}"
                )
        opts = (
            f"lowerdir={self._original_path},"
            f"upperdir={self._upper_dir},"
            f"workdir={self._work_dir}"
        )
        try:
            subprocess.run(
                [
                    "mount",
                    "-t",
                    "overlay",
                    "overlay",
                    "-o",
                    opts,
                    str(self._merged_dir),
                ],
                capture_output=True,
                check=True,
                timeout=30,
            )
        except (
            subprocess.CalledProcessError,
            subprocess.TimeoutExpired,
            OSError,
        ) as exc:
            raise SandboxCreationError(f"Failed to mount OverlayFS: {exc}") from exc

    def _unmount_overlay(self) -> None:
        """Unmount the OverlayFS union. Errors are logged but not raised."""
        if self._merged_dir is None:
            return
        try:
            subprocess.run(
                ["umount", str(self._merged_dir)],
                capture_output=True,
                check=True,
                timeout=30,
            )
        except (
            subprocess.CalledProcessError,
            subprocess.TimeoutExpired,
            OSError,
        ) as exc:
            logger.warning(
                "Failed to unmount overlay at %s: %s",
                self._merged_dir,
                exc,
            )
