from __future__ import annotations

import shutil
import uuid
from contextlib import contextmanager
from pathlib import Path


def _remove_path(path: Path) -> None:
    if path.is_symlink() or path.is_file():
        path.unlink()
    elif path.exists():
        shutil.rmtree(path)


@contextmanager
def staged_directory(destination: str | Path):
    """Build a directory beside its destination and replace on success."""
    target = Path(destination)
    target.parent.mkdir(parents=True, exist_ok=True)
    token = uuid.uuid4().hex
    staging = target.with_name(f".{target.name}.{token}.tmp")
    backup = target.with_name(f".{target.name}.{token}.backup")
    staging.mkdir()
    committed = False
    try:
        yield staging
        if target.exists():
            target.replace(backup)
        try:
            staging.replace(target)
            committed = True
        except Exception:
            if backup.exists() and not target.exists():
                backup.replace(target)
            raise
        if backup.exists():
            _remove_path(backup)
    finally:
        if staging.exists():
            _remove_path(staging)
        if not committed and backup.exists() and not target.exists():
            backup.replace(target)


@contextmanager
def staged_files(destination: str | Path):
    """Stage a set of files and merge them into a directory with rollback."""
    target = Path(destination)
    target.parent.mkdir(parents=True, exist_ok=True)
    token = uuid.uuid4().hex
    staging = target.parent / f".{target.name}.{token}.files.tmp"
    backup = target.parent / f".{target.name}.{token}.files.backup"
    staging.mkdir()
    committed: list[Path] = []
    try:
        yield staging
        staged = sorted(path for path in staging.iterdir() if path.is_file())
        target.mkdir(parents=True, exist_ok=True)
        backup.mkdir()
        try:
            for source in staged:
                destination_path = target / source.name
                if destination_path.exists() or destination_path.is_symlink():
                    destination_path.replace(backup / source.name)
                source.replace(destination_path)
                committed.append(destination_path)
        except Exception:
            for destination_path in reversed(committed):
                _remove_path(destination_path)
            for saved in backup.iterdir():
                saved.replace(target / saved.name)
            raise
        _remove_path(backup)
    finally:
        if staging.exists():
            _remove_path(staging)
        if backup.exists():
            for saved in backup.iterdir():
                destination_path = target / saved.name
                if not destination_path.exists():
                    saved.replace(destination_path)
            _remove_path(backup)