Serialization

def _deserialize_value(val: Any) -> Any:
    """Recursively restore a value from JSON persistence.

    Recognizes type-tagged dicts (`__llmeter_class__`), bytes markers (`__llmeter_bytes__`),
    ISO-8601 datetime strings, and recursively processes nested dicts and lists.
    """
    match val:
        case None | bool() | int() | float():
            return val
        case str() if _DATETIME_RE.fullmatch(val):
            return str_to_datetime(val)
        case str():
            return val
        case {"__llmeter_class__": _, "__llmeter_state__": _}:
            return load_object(val)
        case {"__llmeter_bytes__": b64} if len(val) == 1:
            return base64.b64decode(b64)
        case dict():
            return {k: _deserialize_value(v) for k, v in val.items()}
        case list() | tuple():
            return [_deserialize_value(item) for item in val]
        case _:
            return val

def _get_type_args(tp) -> tuple:
    """Return the members of a union type (e.g. `datetime | None` -> (datetime, NoneType))."""
    if isinstance(tp, _types.UnionType):
        return tp.__args__
    origin = getattr(tp, "__origin__", None)
    if origin is _types.UnionType:
        return tp.__args__
    return (tp,) if isinstance(tp, type) else ()

def _serialize_value(val: Any) -> Any:
    """Recursively prepare a value for JSON persistence.

    Handles primitives, known types (bytes, datetime, PathLike), nested
    [`Serializable`][llmeter.serialization.Serializable] objects (via
    [`dump_object`][llmeter.serialization.dump_object]), dicts, and lists/tuples.

    Raises:
        TypeError: for objects it cannot serialize.
    """
    if val is None or isinstance(val, (str, int, float, bool)):
        return val
    for types, fn in _SERIALIZERS:
        if isinstance(val, types):
            return fn(val)
    if hasattr(val, "_get_llmeter_state"):
        return dump_object(val)
    if isinstance(val, dict):
        return {k: _serialize_value(v) for k, v in val.items()}
    if isinstance(val, (list, tuple)):
        return [_serialize_value(item) for item in val]
    raise TypeError(
        f"Cannot serialize {type(val).__name__!r} object: {val!r}. "
        "Only JSON primitives, bytes, datetime, PathLike, dicts, lists/tuples, and "
        "LLMeter-serializable objects are supported. To make a custom object "
        "serializable (e.g. a custom cost dimension or callback), have its class "
        "inherit from llmeter.serialization.Serializable."
    )

def bytes_decoder(dct: dict) -> dict | bytes:
    """Decode `__llmeter_bytes__` marker objects back to Python `bytes`.

    Intended for use as the `object_hook` argument to `json.load` or `json.loads`. Marker objects
    produced by [`json_default`][llmeter.serialization.json_default] are detected and converted
    back to `bytes`; all other dicts pass through unchanged.

    Args:
        dct: A dictionary produced by the JSON parser.

    Returns:
        The original `bytes` if *dct* is an LLMeter bytes marker object, otherwise *dct* unchanged.
    """
    if "__llmeter_bytes__" in dct and len(dct) == 1:
        return base64.b64decode(dct["__llmeter_bytes__"])
    return dct

def datetime_to_str(dt: datetime) -> str:
    """Convert a datetime to a UTC ISO-8601 string with ``Z`` suffix.

    Timezone-aware datetimes are converted to UTC first. Naive datetimes are serialized as-is
    (assumed UTC).
    """
    if dt.tzinfo is not None:
        dt = dt.astimezone(timezone.utc)
    return dt.isoformat(timespec="seconds").replace("+00:00", "Z")

def dump_object(obj: Any) -> dict:
    """Serialize an object to a type-tagged dict for round-trip persistence.

    The returned envelope has the form
    `{"__llmeter_class__": "module.Class", "__llmeter_state__": {...}}`.

    Serialization strategy (checked in order):

    1. If the object implements the LLMeter state protocol (as in
        [`Serializable._get_llmeter_state`][llmeter.serialization.Serializable._get_llmeter_state])
        this will call it to obtain the state dict.
    2. If the object is a dataclass, uses `dataclasses.asdict`.
    3. Otherwise, takes all public (non-underscore-prefixed) entries from `__dict__`.

    Args:
        obj: The object to serialize.

    Returns:
        dct: A JSON-serializable dict that [`load_object`][llmeter.serialization.load_object] can
            reconstruct.
    """
    class_path = f"{obj.__class__.__module__}.{obj.__class__.__qualname__}"
    if hasattr(obj, "_get_llmeter_state"):
        state = obj._get_llmeter_state()
    elif is_dataclass(obj) and not isinstance(obj, type):
        state = asdict(obj)
    elif hasattr(obj, "__dict__"):
        state = {k: v for k, v in obj.__dict__.items() if not k.startswith("_")}
    else:
        state = {}
    return {"__llmeter_class__": class_path, "__llmeter_state__": state}

def json_default(obj: Any) -> Any:
    """Serialize a single non-natively-JSON-serializable object.

    Intended for use as the `default` argument to `json.dump` or `json.dumps`. This **does not**
    handle full recursive LLMeter serialization - see
    [`dump_object`][llmeter.serialization.dump_object] instead.

    Type handling (checked in order):

    * `bytes` — wrapped in a `{"__llmeter_bytes__": "<base64>"}` marker so that
        [`bytes_decoder`][llmeter.serialization.bytes_decoder] can restore them on the way back.
    * `datetime` — converted to a UTC ISO-8601 string with a `Z` suffix.
    * `date` / `time` — converted via `.isoformat()`.
    * `os.PathLike` — converted to a POSIX path string.
    * Anything else — `str()` fallback (returns `None` if that also fails).

    Args:
        obj: The object that the default JSON encoder could not handle.

    Returns:
        A JSON-serializable representation of *obj*.
    """
    if isinstance(obj, bytes):
        return {"__llmeter_bytes__": base64.b64encode(obj).decode("utf-8")}
    if isinstance(obj, datetime):
        return datetime_to_str(obj)
    if isinstance(obj, (date, time)):
        return obj.isoformat()
    if isinstance(obj, (os.PathLike, Path)):
        return Path(obj).as_posix()
    try:
        return str(obj)
    except Exception:
        return None

def load_object(data: dict) -> Any:
    """Restore an object from a type-tagged dict produced by `dump_object`.

    !!! warning
        This method imports and instantiates class paths specified by the input data, which (like
        pickle) can enable arbitrary running arbitrary code. Do not run it on data from unstrusted
        sources!

    Imports the module identified by `__llmeter_class__`, instantiates the class (bypassing
    `__init__` via `__new__`), and restores its state via
    [`Serializable._set_llmeter_state`][llmeter.serialization.Serializable._set_llmeter_state].
    Objects that do not implement the protocol are restored by assigning the (deserialized) state
    onto their `__dict__`.

    Args:
        data: A dict with `__llmeter_class__` and `__llmeter_state__` keys, as
            produced by [`dump_object`][llmeter.serialization.dump_object].

    Returns:
        The reconstructed object instance.
    """
    class_path = data["__llmeter_class__"]
    module_path, class_name = class_path.rsplit(".", 1)
    module = importlib.import_module(module_path)
    cls = getattr(module, class_name)

    state = data["__llmeter_state__"]
    obj = cls.__new__(cls)
    if hasattr(obj, "_set_llmeter_state"):
        obj._set_llmeter_state(state)
    else:
        # Fallback for non-Serializable classes (state came from dump_object's dataclass /
        # __dict__ branches). We assign attributes directly rather than calling __setstate__:
        # `object` has no __setstate__ (so it would raise for a plain class), and if the class
        # *does* define one it's the pickle hook, which may expect a different state shape than
        # our JSON dict. Direct assignment (with per-value deserialization) is the safe generic
        # restore, and mirrors pickle's own "restore __dict__ without calling __init__" default.
        for key, value in state.items():
            setattr(obj, key, _deserialize_value(value))
    return obj

def restore_dataclass_types(cls: type, data: dict) -> None:
    """Restore typed fields in a dict destined for a dataclass constructor.

    Introspects `cls` (a dataclass) and converts JSON-native values back to their annotated Python
    types. Currently handles:

    * `datetime` fields — parses ISO-8601 strings via `str_to_datetime`.
    * `bytes` fields — decodes `__llmeter_bytes__` markers via base64.

    Only fields declared on `cls` are touched — nested user payloads (e.g. `input_payload`) are
    left unchanged. Mutates *data* in place.

    Args:
        cls: A dataclass type to introspect for field type annotations.
        data: A dictionary of field values (e.g. from `json.load`) to coerce.
    """
    for f in fields(cls):
        val = data.get(f.name)
        if val is None:
            continue
        type_args = _get_type_args(f.type)
        match val:
            case str() if datetime in type_args:
                try:
                    data[f.name] = str_to_datetime(val)
                except ValueError:
                    pass
            case {"__llmeter_bytes__": b64} if bytes in type_args and len(val) == 1:
                data[f.name] = base64.b64decode(b64)

def str_to_datetime(s: str) -> datetime:
    """Parse an ISO-8601 string (with optional ``Z`` suffix) to a datetime."""
    return datetime.fromisoformat(s.replace("Z", "+00:00"))