"""Atomic cross-process updates for JSON data files."""

import fcntl
import json
import os
import shutil
import tempfile
import threading
from collections.abc import Callable
from contextlib import suppress
from pathlib import Path
from typing import TypeAlias, TypeVar

JsonValue: TypeAlias = (
    None | bool | int | float | str | list["JsonValue"] | dict[str, "JsonValue"]
)
JsonContainer: TypeAlias = list[JsonValue] | dict[str, JsonValue]
JsonContainerT = TypeVar("JsonContainerT", bound=JsonContainer)

_update_lock = threading.Lock()


def _load_json(path: Path, default_factory: Callable[[], JsonContainerT]) -> JsonContainerT:
    try:
        with path.open(encoding="utf-8") as file:
            return json.load(file)
    except FileNotFoundError:
        return default_factory()
    except json.JSONDecodeError:
        shutil.copy2(path, f"{path}.corrupt")
        # Recover from the last good copy: update_json is about to overwrite it.
        return _load_json(Path(f"{path}.bak"), default_factory)


def _atomic_write(path: Path, data: JsonContainer) -> None:
    descriptor, temporary_path = tempfile.mkstemp(
        dir=path.parent,
        prefix=f".{path.name}.",
    )
    try:
        with os.fdopen(descriptor, "w", encoding="utf-8") as file:
            json.dump(data, file, ensure_ascii=False, indent=2)
            file.flush()
            os.fsync(file.fileno())
        os.replace(temporary_path, path)
    finally:
        with suppress(FileNotFoundError):
            os.unlink(temporary_path)


def update_json(
    path: str | Path,
    default_factory: Callable[[], JsonContainerT],
    update: Callable[[JsonContainerT], None],
) -> JsonContainerT:
    """Lock, update, and atomically replace a JSON file."""
    json_path = Path(path)
    json_path.parent.mkdir(parents=True, exist_ok=True)
    lock_path = Path(f"{json_path}.lock")

    # ponytail: one process-local lock is enough for these low-volume data files.
    with _update_lock, lock_path.open("a") as lock_file:
        fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
        data = _load_json(json_path, default_factory)
        update(data)
        if json_path.exists():
            shutil.copy2(json_path, f"{json_path}.bak")
        _atomic_write(json_path, data)
        return data