# SPDX-License-Identifier: Apache-2.0
"""Explicit tensor/value continuation format; never deserialize Python code."""

import json
import math

import torch
from safetensors.torch import load, save

from ..utils.device_telemetry import emit_device_telemetry
from .input_cache import InputCache


class ContinuationCodec:
    VERSION = 2

    @staticmethod
    def dumps(state) -> bytes:
        tensors = {}
        active = set()

        def encode(value):
            if value is None or type(value) in (bool, int, str):
                return ["scalar", value]
            if type(value) is float and math.isfinite(value):
                return ["scalar", value]
            if isinstance(value, torch.Tensor):
                if (
                    value.device.type not in {"cpu", "cuda"}
                    or value.layout != torch.strided
                ):
                    raise TypeError("continuation requires materialized dense tensors")
                name = str(len(tensors))
                tensors[name] = value.detach().to(device="cpu").contiguous().clone()
                emit_device_telemetry(
                    "checkpoint_tensor_captured",
                    tensor_id=name,
                    source_device=str(value.device),
                    storage_device="cpu",
                    shape=list(value.shape),
                    dtype=str(value.dtype),
                )
                return ["tensor", {"name": name, "device": str(value.device)}]
            if type(value) is torch.device:
                device = value
                if device.type not in {"cpu", "cuda"}:
                    raise TypeError("unsupported continuation metadata device")
                if device.type == "cuda" and device.index is None:
                    device = torch.device("cuda", torch.cuda.current_device())
                emit_device_telemetry(
                    "checkpoint_device_metadata_captured", device=str(device)
                )
                return ["device", str(device)]
            if id(value) in active:
                raise TypeError("cyclic continuation state is unsupported")
            active.add(id(value))
            try:
                if type(value) is InputCache:
                    return ["input_cache", encode(vars(value))]
                if type(value) is dict:
                    if any(type(key) not in (str, int, bool) for key in value):
                        raise TypeError(
                            "continuation keys must be strings, integers, or booleans"
                        )
                    return [
                        "dict",
                        [[encode(key), encode(item)] for key, item in value.items()],
                    ]
                if type(value) in (list, tuple):
                    return [
                        "tuple" if type(value) is tuple else "list",
                        [encode(item) for item in value],
                    ]
            finally:
                active.remove(id(value))
            raise TypeError(f"unsupported continuation value: {type(value).__name__}")

        tree = encode(state)
        # Metadata and tensor payload are one immutable object, not independently
        # published files. Cloning removes safetensors shared-storage ambiguity.
        return save(
            tensors,
            metadata={
                "continuation": json.dumps(
                    {"version": ContinuationCodec.VERSION, "tree": tree},
                    allow_nan=False,
                )
            },
        )

    @staticmethod
    def loads(data: bytes):
        # safetensors exposes metadata through safe_open for files, but this
        # codec operates on verified bytes, so parse its length-prefixed header.
        header_size = int.from_bytes(data[:8], "little")
        if header_size > len(data) - 8 or header_size < 2:
            raise ValueError("invalid continuation header")
        header = json.loads(data[8 : 8 + header_size])
        document = json.loads(header["__metadata__"]["continuation"])
        if document["version"] != ContinuationCodec.VERSION:
            raise ValueError("unsupported continuation version")
        tensors = load(data)

        def decode(node):
            kind, value = node
            if kind == "scalar" and (
                value is None
                or type(value) in (str, bool, int)
                or type(value) is float
                and math.isfinite(value)
            ):
                return value
            if kind == "tensor":
                device = torch.device(value["device"])
                if (
                    device.type not in {"cpu", "cuda"}
                    or device.type == "cuda"
                    and device.index is None
                ):
                    raise ValueError(
                        "unsupported or unindexed continuation tensor device"
                    )
                stored = tensors[value["name"]]
                expected_shape, expected_dtype = stored.shape, stored.dtype
                # Keep an independent reference: a faulty in-place transfer must
                # not mutate both the restored tensor and its verification input.
                expected_bits = (
                    stored.contiguous().reshape(-1).view(torch.uint8).clone()
                )
                tensor = stored.to(device=device)
                placement_matched = tensor.device == device
                state_matched = (
                    tensor.dtype == expected_dtype
                    and tensor.shape == expected_shape
                    and torch.equal(
                        tensor.detach()
                        .cpu()
                        .contiguous()
                        .reshape(-1)
                        .view(torch.uint8),
                        expected_bits,
                    )
                )
                matched = placement_matched and state_matched
                emit_device_telemetry(
                    "checkpoint_tensor_restored",
                    tensor_id=value["name"],
                    expected_device=str(device),
                    actual_device=str(tensor.device),
                    matched=matched,
                    placement_matched=placement_matched,
                    state_matched=state_matched,
                    shape=list(tensor.shape),
                    dtype=str(tensor.dtype),
                )
                if not matched:
                    raise ValueError(
                        "checkpoint tensor state or device placement restore mismatch"
                    )
                return tensor
            if kind == "device":
                device = torch.device(value)
                if (
                    device.type not in {"cpu", "cuda"}
                    or device.type == "cuda"
                    and device.index is None
                ):
                    raise ValueError(
                        "unsupported or unindexed continuation metadata device"
                    )
                emit_device_telemetry(
                    "checkpoint_device_metadata_restored",
                    expected_device=value,
                    actual_device=str(device),
                    matched=str(device) == value,
                )
                if str(device) != value:
                    raise ValueError("checkpoint device metadata restore mismatch")
                return device
            if kind == "input_cache":
                fields = decode(value)
                if set(fields) != set(InputCache.__dataclass_fields__):
                    raise ValueError("unsupported InputCache schema")
                return InputCache(**fields)
            if kind == "dict":
                result = {}
                for key, item in value:
                    key = decode(key)
                    if type(key) not in (str, int, bool) or key in result:
                        raise ValueError("invalid continuation dictionary key")
                    result[key] = decode(item)
                return result
            if kind in ("list", "tuple"):
                items = [decode(item) for item in value]
                return tuple(items) if kind == "tuple" else items
            raise ValueError(f"unsupported continuation tag: {kind}")

        return decode(document["tree"])