"""Atomic file IO — the ONLY way the engine writes files.

Every write goes to a temp sibling, is flushed + fsync'd, then ``os.replace``'d over the target — so a
reader never sees a half-written file and a crashed write leaves the previous version intact. Backs the
per-phase Parquet checkpoints (K13).
"""

from __future__ import annotations

import os
import tempfile
from collections.abc import Callable
from pathlib import Path
from typing import IO

import polars as pl


def _atomic_write(path: str | Path, write_bytes: Callable[[IO[bytes]], object]) -> Path:
    """Write via a temp sibling + ``os.replace`` so the target is updated atomically."""
    target = Path(path)
    target.parent.mkdir(parents=True, exist_ok=True)
    handle_fd, temp_name = tempfile.mkstemp(dir=target.parent, suffix=".tmp")
    try:
        with os.fdopen(handle_fd, "wb") as handle:
            write_bytes(handle)
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(temp_name, target)
    except BaseException:
        Path(temp_name).unlink(missing_ok=True)  # never leave a stray temp file behind
        raise
    return target


def atomic_write_bytes(path: str | Path, payload: bytes) -> Path:
    """Atomically write raw bytes."""
    return _atomic_write(path, lambda handle: handle.write(payload))


def atomic_write_parquet(frame: pl.DataFrame, path: str | Path) -> Path:
    """Atomically write a Polars frame to Parquet."""
    return _atomic_write(path, frame.write_parquet)