"""Bounded, private, atomic JSON output for assessment evidence."""

from __future__ import annotations

import json
import os
from pathlib import Path
import stat
import tempfile
from typing import Any


class AtomicJsonWriteError(RuntimeError):
    """Raised when JSON output cannot be written with the required guarantees."""


def atomic_write_json(
    path: str | os.PathLike[str],
    payload: Any,
    *,
    max_bytes: int,
) -> None:
    """Write bounded JSON through a private same-directory temporary file.

    A successful return means the file contents and the containing directory
    entry have both been flushed. If flushing the directory fails after the
    atomic replacement, the replacement may be visible but its crash
    durability is unknown, so the function reports failure. Every parent must
    already exist as a real directory; symbolic-link components are rejected.
    """

    if isinstance(max_bytes, bool) or not isinstance(max_bytes, int) or max_bytes <= 0:
        raise ValueError("max_bytes must be a positive integer")

    encoded = _encode_bounded_json(payload, max_bytes=max_bytes)
    target = Path(os.path.abspath(os.fspath(path)))
    parent = target.parent
    descriptor = -1
    verification_descriptor = -1
    temporary_path: Path | None = None
    temporary_identity: tuple[int, int] | None = None

    try:
        _assert_real_parent(parent)
        _assert_replaceable_target(target)
        try:
            descriptor, temporary_name = tempfile.mkstemp(
                dir=parent,
                prefix=".govoplan-json-",
                suffix=".tmp",
            )
            temporary_path = Path(temporary_name)
            os.fchmod(descriptor, 0o600)
            metadata = os.fstat(descriptor)
            if (
                not stat.S_ISREG(metadata.st_mode)
                or stat.S_IMODE(metadata.st_mode) != 0o600
            ):
                raise AtomicJsonWriteError(
                    "Atomic JSON temporary file could not be secured."
                )
            temporary_identity = (metadata.st_dev, metadata.st_ino)
            verification_descriptor = os.dup(descriptor)

            handle = os.fdopen(descriptor, "wb")
            descriptor = -1
            with handle:
                handle.write(encoded)
                handle.flush()
                os.fsync(handle.fileno())

            # Recheck after creating and flushing the temporary file so an
            # unsafe target discovered before replacement is never followed.
            _assert_real_parent(parent)
            _assert_replaceable_target(target)
            os.replace(temporary_path, target)
            temporary_path = None
            _assert_installed_target(
                target,
                expected_identity=temporary_identity,
                recovery_descriptor=verification_descriptor,
            )
            _fsync_directory(parent)
        finally:
            try:
                if descriptor >= 0:
                    os.close(descriptor)
            finally:
                try:
                    if verification_descriptor >= 0:
                        os.close(verification_descriptor)
                finally:
                    if temporary_path is not None:
                        try:
                            temporary_path.unlink()
                        except FileNotFoundError:
                            pass
    except AtomicJsonWriteError:
        raise
    except OSError as exc:
        raise AtomicJsonWriteError(
            "JSON output could not be written atomically."
        ) from exc


def _encode_bounded_json(payload: Any, *, max_bytes: int) -> bytes:
    encoder = json.JSONEncoder(
        allow_nan=False,
        ensure_ascii=False,
        indent=2,
        sort_keys=True,
    )
    chunks: list[bytes] = []
    encoded_size = 1  # The document always ends with one newline.
    for chunk in encoder.iterencode(payload):
        encoded_chunk = chunk.encode("utf-8")
        encoded_size += len(encoded_chunk)
        if encoded_size > max_bytes:
            raise AtomicJsonWriteError("JSON output exceeds its size limit.")
        chunks.append(encoded_chunk)
    return b"".join(chunks) + b"\n"


def _assert_replaceable_target(path: Path) -> None:
    try:
        metadata = path.lstat()
    except FileNotFoundError:
        return
    except OSError as exc:
        raise AtomicJsonWriteError(
            "JSON output target could not be inspected safely."
        ) from exc
    if stat.S_ISLNK(metadata.st_mode):
        raise AtomicJsonWriteError("JSON output target must not be a symbolic link.")
    if not stat.S_ISREG(metadata.st_mode):
        raise AtomicJsonWriteError("JSON output target must be a regular file.")


def _assert_real_parent(path: Path) -> None:
    current = Path(path.anchor)
    for part in path.parts[1:]:
        current /= part
        try:
            metadata = current.lstat()
        except FileNotFoundError as exc:
            raise AtomicJsonWriteError(
                "JSON output parent directory must already exist."
            ) from exc
        except OSError as exc:
            raise AtomicJsonWriteError(
                "JSON output parent directory could not be inspected safely."
            ) from exc
        if stat.S_ISLNK(metadata.st_mode):
            raise AtomicJsonWriteError(
                "JSON output parent path must not contain symbolic links."
            )
        if not stat.S_ISDIR(metadata.st_mode):
            raise AtomicJsonWriteError(
                "JSON output parent path must contain only directories."
            )


def _assert_installed_target(
    path: Path,
    *,
    expected_identity: tuple[int, int] | None,
    recovery_descriptor: int,
) -> None:
    try:
        metadata = path.lstat()
    except OSError as exc:
        raise AtomicJsonWriteError(
            "Atomic JSON replacement could not be verified."
        ) from exc
    observed_identity = (metadata.st_dev, metadata.st_ino)
    if (
        expected_identity is None
        or observed_identity != expected_identity
        or not stat.S_ISREG(metadata.st_mode)
    ):
        raise AtomicJsonWriteError(
            "Atomic JSON replacement did not preserve the secured regular file."
        )
    if stat.S_IMODE(metadata.st_mode) == 0o600:
        return
    try:
        recovery_metadata = os.fstat(recovery_descriptor)
        if (
            not stat.S_ISREG(recovery_metadata.st_mode)
            or (recovery_metadata.st_dev, recovery_metadata.st_ino) != expected_identity
        ):
            raise AtomicJsonWriteError(
                "Atomic JSON replacement recovery descriptor is invalid."
            )
        os.fchmod(recovery_descriptor, 0o600)
        os.fsync(recovery_descriptor)
        recovered_metadata = os.fstat(recovery_descriptor)
        installed_metadata = path.lstat()
    except OSError as exc:
        raise AtomicJsonWriteError(
            "Atomic JSON replacement permissions could not be restored."
        ) from exc
    if (
        stat.S_IMODE(recovered_metadata.st_mode) != 0o600
        or (installed_metadata.st_dev, installed_metadata.st_ino) != expected_identity
        or not stat.S_ISREG(installed_metadata.st_mode)
        or stat.S_IMODE(installed_metadata.st_mode) != 0o600
    ):
        raise AtomicJsonWriteError(
            "Atomic JSON replacement permissions could not be restored."
        )
    _fsync_directory(path.parent)
    raise AtomicJsonWriteError(
        "Atomic JSON replacement permissions changed and were restored."
    )


def _fsync_directory(path: Path) -> None:
    flags = os.O_RDONLY
    if hasattr(os, "O_DIRECTORY"):
        flags |= os.O_DIRECTORY
    if hasattr(os, "O_NOFOLLOW"):
        flags |= os.O_NOFOLLOW
    descriptor = os.open(path, flags)
    try:
        if not stat.S_ISDIR(os.fstat(descriptor).st_mode):
            raise AtomicJsonWriteError("JSON output parent path must be a directory.")
        os.fsync(descriptor)
    finally:
        os.close(descriptor)