"""Race-safe immutable checkpoint management.

Checkpoint directories are written under a temporary name and atomically
renamed once complete.  ``best`` and ``latest`` are atomically replaced symlinks.
No checkpoint is deleted while a run is active; pruning, if ever requested,
must happen only after all consumers have stopped.
"""
from __future__ import annotations

import json
import os
import uuid
from pathlib import Path
from typing import Any, Callable


class AtomicCheckpointManager:
    def __init__(self, root: Path):
        self.root = root.resolve()
        self.checkpoints = self.root / "checkpoints"
        self.checkpoints.mkdir(parents=True, exist_ok=True)

    def _atomic_link(self, link_name: str, target: Path) -> None:
        link = self.root / link_name
        temp = self.root / f".{link_name}.tmp.{os.getpid()}.{uuid.uuid4().hex}"
        relative_target = os.path.relpath(target, self.root)
        os.symlink(relative_target, temp)
        os.replace(temp, link)

    def save(
        self,
        tag: str,
        writer: Callable[[Path], None],
        metadata: dict[str, Any],
        *,
        is_best: bool,
    ) -> Path:
        final = self.checkpoints / tag
        if final.exists():
            raise FileExistsError(f"immutable checkpoint already exists: {final}")
        temp = self.checkpoints / f".{tag}.tmp.{os.getpid()}.{uuid.uuid4().hex}"
        temp.mkdir(parents=False)
        try:
            writer(temp)
            with (temp / "checkpoint_metadata.json").open("w", encoding="utf-8") as f:
                json.dump(metadata, f, ensure_ascii=False, indent=2)
                f.flush()
                os.fsync(f.fileno())
            os.replace(temp, final)
        except Exception:
            # Leave a failed temporary directory for diagnosis; never remove a
            # possibly-active directory from another process.
            raise

        self._atomic_link("latest", final)
        if is_best:
            self._atomic_link("best", final)
        return final