"""Helpers for strict JSON-compatible checkpoint payloads."""
from __future__ import annotations

import json
import math
from typing import Any, TypeAlias, cast


JSONScalar: TypeAlias = None | bool | int | float | str
JSONValue: TypeAlias = JSONScalar | list["JSONValue"] | dict[str, "JSONValue"]
JSONObject: TypeAlias = dict[str, JSONValue]


def ensure_json_object(value: Any, *, path: str = "$") -> JSONObject:
    """Validate that a value is a strict JSON object tree."""
    validated = _validate_json_value(value, path=path)
    if not isinstance(validated, dict):
        raise ValueError(f"{path} must be a JSON object, got {type(value).__name__}")
    return validated


def clone_json_object(value: JSONObject, *, path: str = "$") -> JSONObject:
    """Deep-copy a validated JSON object using the JSON codec."""
    validated = ensure_json_object(value, path=path)
    return cast(
        JSONObject,
        json.loads(json.dumps(validated, ensure_ascii=False, allow_nan=False)),
    )


def _validate_json_value(value: Any, *, path: str) -> JSONValue:
    if value is None or isinstance(value, (str, bool, int)):
        return cast(JSONValue, value)
    if isinstance(value, float):
        if not math.isfinite(value):
            raise ValueError(f"{path} must be a finite JSON number, got {value!r}")
        return value
    if isinstance(value, list):
        return [
            _validate_json_value(item, path=f"{path}[{index}]")
            for index, item in enumerate(value)
        ]
    if isinstance(value, dict):
        normalized: JSONObject = {}
        for key, item in value.items():
            if not isinstance(key, str):
                raise ValueError(
                    f"{path} must use string keys, got key of type {type(key).__name__}"
                )
            child_path = f"{path}.{key}" if path else key
            normalized[key] = _validate_json_value(item, path=child_path)
        return normalized
    raise ValueError(
        f"{path} must be JSON-compatible, got {type(value).__name__}"
    )


__all__ = [
    "JSONObject",
    "JSONValue",
    "clone_json_object",
    "ensure_json_object",
]
