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)