from __future__ import annotations

from typing import Any

import cloudpickle

from .backend import StateCodecCompatibilityError


DEFAULT_COMPATIBILITY_VERSION = 1
DEFAULT_SCHEMA_VERSION = 1


def encode_state(
    value: Any,
    *,
    compatibility_version: int = DEFAULT_COMPATIBILITY_VERSION,
    schema_version: int = DEFAULT_SCHEMA_VERSION,
) -> bytes:
    payload = cloudpickle.dumps(value)
    envelope = {
        "compatibility_version": int(compatibility_version),
        "schema_version": int(schema_version),
        "payload": payload,
    }
    return cloudpickle.dumps(envelope)


def decode_state(
    blob: bytes,
    *,
    expected_compatibility_version: int = DEFAULT_COMPATIBILITY_VERSION,
) -> Any:
    envelope = cloudpickle.loads(blob)
    if not isinstance(envelope, dict):
        raise StateCodecCompatibilityError("State envelope must be a dict.")
    compatibility_version = envelope.get("compatibility_version")
    if int(compatibility_version) != int(expected_compatibility_version):
        raise StateCodecCompatibilityError(
            "State compatibility_version mismatch: "
            f"expected={expected_compatibility_version}, got={compatibility_version}"
        )
    payload = envelope.get("payload")
    if not isinstance(payload, (bytes, bytearray)):
        raise StateCodecCompatibilityError("State payload must be bytes.")
    return cloudpickle.loads(bytes(payload))


__all__ = [
    "DEFAULT_COMPATIBILITY_VERSION",
    "DEFAULT_SCHEMA_VERSION",
    "encode_state",
    "decode_state",
]
