from __future__ import annotations

import shutil
import uuid
from collections.abc import Iterable
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,
    *,
    remove: Iterable[str] = (),
):
    """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())
        staged_names = {path.name for path in staged}
        removed_names = set(remove) - staged_names
        target.mkdir(parents=True, exist_ok=True)
        backup.mkdir()
        try:
            for name in sorted(removed_names):
                destination_path = target / name
                if destination_path.is_file() or destination_path.is_symlink():
                    destination_path.replace(backup / name)
            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)


def publish_file_sets(file_sets: dict[Path, tuple[Path, Iterable[str]]]) -> None:
    """Commit several staged file sets together, rolling every set back on error.

    Staging and destination directories must be on the same filesystem. Unlike
    nesting staged_files, failure in a later directory restores earlier ones too.
    """
    token = uuid.uuid4().hex
    backups: list[tuple[Path, Path]] = []
    committed: list[Path] = []
    try:
        for target, (staging, remove) in file_sets.items():
            target.mkdir(parents=True, exist_ok=True)
            backup = target / f".publish-{token}.backup"
            backup.mkdir()
            backups.append((target, backup))
            staged = sorted(path for path in staging.iterdir() if path.is_file())
            names = {path.name for path in staged} | set(remove)
            for name in sorted(names):
                if Path(name).name != name:
                    raise ValueError(f"Invalid output filename: {name}")
                destination = target / name
                if destination.exists() or destination.is_symlink():
                    destination.replace(backup / name)
            for source in staged:
                destination = target / source.name
                source.replace(destination)
                committed.append(destination)
    except Exception:
        for destination in reversed(committed):
            _remove_path(destination)
        for target, backup in reversed(backups):
            for saved in backup.iterdir():
                saved.replace(target / saved.name)
        raise
    finally:
        for _, backup in backups:
            if backup.exists() and not any(backup.iterdir()):
                backup.rmdir()
    for _, backup in backups:
        _remove_path(backup)