"""Encode and decode the native Shard protocol's named-tensor bundles.

The generated stubs give us message *structure*; they cannot enforce the
invariants that keep a distributed forward correct. A bundle whose declared
shape disagrees with its byte count, whose fragments leave a hole, or whose
checksum does not match is not a slightly-wrong activation — it is silently
wrong tokens for the rest of the generation. So decoding is validating: every
path into a tensor's bytes goes through :func:`decode_tensor`, which refuses a
payload it cannot fully account for.

Compression is a transport optimisation and is decided by the same policy layer
the existing HTTP seam already uses (``activation_compression``), so a node's
tuned thresholds apply to both transports.
"""

from __future__ import annotations

from dataclasses import dataclass
import struct
from typing import Iterable, Sequence

from ..activation_compression import (
    CompressionPolicy,
    compress_activation,
    decompress_activation,
)
from .generated import shard_runtime_pb2 as pb

# The schema generation this build speaks. A peer offering something else is
# rejected at the handshake rather than being half-understood.
SCHEMA_VERSION = pb.SCHEMA_VERSION_1

# Generation of the tensor-bundle layout, versioned independently of the
# protocol so a boundary payload can evolve without a protocol bump.
BUNDLE_VERSION = 1

# Token-aligned prefill chunk bound. 128 tokens is the size ADR-0008 already
# uses on the HTTP seam; keeping it identical means seam bytes stay comparable
# across transports.
DEFAULT_MAX_PREFILL_CHUNK_TOKENS = 128

# gRPC's default maximum receive size. Fragmenting below it keeps us inside the
# default limits of any conformant peer instead of requiring every client to
# raise its window.
DEFAULT_MAX_CHUNK_BYTES = 4 * 1024 * 1024

# Leave room for envelope and framing overhead inside one chunk message.
DEFAULT_MAX_FRAGMENT_BYTES = 1024 * 1024

DEFAULT_MAX_INFLIGHT_CHUNKS = 8
DEFAULT_MAX_FRAGMENTS_PER_TENSOR = 64
DEFAULT_MAX_TENSORS_PER_BUNDLE = 64
DEFAULT_MAX_TENSOR_RANK = 8
DEFAULT_MAX_TENSOR_DIMENSION = (1 << 31) - 1

# Canonical boundary tensor name for a dense transformer hidden state.
HIDDEN_STATES = "hidden_states"

_DTYPE_ITEMSIZE: dict[int, int] = {
    pb.DTYPE_BFLOAT16: 2,
    pb.DTYPE_FLOAT16: 2,
    pb.DTYPE_FLOAT32: 4,
    pb.DTYPE_INT32: 4,
    pb.DTYPE_INT64: 8,
    pb.DTYPE_UINT8: 1,
    pb.DTYPE_INT8: 1,
    pb.DTYPE_BOOL: 1,
}


class ProtocolError(Exception):
    """A peer sent something this build cannot safely interpret."""


class PayloadCorrupt(ProtocolError):
    """A tensor payload failed validation: size, coverage, or checksum."""


def itemsize(dtype: int) -> int:
    try:
        return _DTYPE_ITEMSIZE[dtype]
    except KeyError:
        raise ProtocolError(f"unsupported dtype {dtype}") from None


def expected_bytes(
    shape: Sequence[int],
    dtype: int,
    *,
    max_rank: int = DEFAULT_MAX_TENSOR_RANK,
    max_dimension: int = DEFAULT_MAX_TENSOR_DIMENSION,
    max_bytes: int | None = None,
) -> int:
    """Byte count a tensor of `shape` and `dtype` must occupy."""
    if len(shape) > max_rank:
        raise ProtocolError(f"tensor rank {len(shape)} exceeds limit {max_rank}")
    if any(dim < 0 or dim > max_dimension for dim in shape):
        raise ProtocolError(
            f"dimension outside 0..{max_dimension} in shape {list(shape)}"
        )
    size = itemsize(dtype)
    count = 1
    for dim in shape:
        count *= dim
        if max_bytes is not None and count * size > max_bytes:
            raise ProtocolError(f"tensor shape {list(shape)} exceeds byte limit {max_bytes}")
    return count * size


# --- CRC32C ----------------------------------------------------------------
#
# CRC32C (Castagnoli), not zlib's CRC32: it is the checksum gRPC, and the
# storage systems these payloads pass through, already use, and hardware
# implements it. `google_crc32c` is used when present; the table fallback keeps
# the default test suite dependency-free.

_CRC32C_POLY = 0x82F63B78
_CRC32C_TABLE: list[int] = []
for _i in range(256):
    _c = _i
    for _ in range(8):
        _c = (_c >> 1) ^ (_CRC32C_POLY if _c & 1 else 0)
    _CRC32C_TABLE.append(_c)

try:  # pragma: no cover - depends on an optional native package
    from google_crc32c import value as _fast_crc32c
except ImportError:  # pragma: no cover
    _fast_crc32c = None


def crc32c(data: bytes) -> int:
    if _fast_crc32c is not None:  # pragma: no cover - optional fast path
        return _fast_crc32c(data)
    crc = 0xFFFFFFFF
    for byte in data:
        crc = (crc >> 8) ^ _CRC32C_TABLE[(crc ^ byte) & 0xFF]
    return crc ^ 0xFFFFFFFF


def checksum_of(data: bytes) -> pb.Checksum:
    return pb.Checksum(
        algorithm=pb.CHECKSUM_ALGORITHM_CRC32C,
        value=struct.pack(">I", crc32c(data)),
    )


# --- Tensors ---------------------------------------------------------------


def encode_tensor(
    name: str,
    data: bytes,
    shape: Sequence[int],
    dtype: int = pb.DTYPE_BFLOAT16,
    *,
    policy: CompressionPolicy | None = None,
    max_chunk_bytes: int = DEFAULT_MAX_CHUNK_BYTES,
    max_fragment_bytes: int = DEFAULT_MAX_FRAGMENT_BYTES,
    max_fragments: int = DEFAULT_MAX_FRAGMENTS_PER_TENSOR,
) -> pb.NamedTensor:
    """Build a NamedTensor, compressing and fragmenting as needed.

    `data` is the uncompressed little-endian payload. The checksum is taken over
    it *before* compression so it stays valid whichever framing a hop chooses.
    """
    if max_chunk_bytes <= 0 or max_fragment_bytes <= 0 or max_fragments <= 0:
        raise ProtocolError("tensor byte/count bounds must be positive")
    declared = expected_bytes(shape, dtype, max_bytes=max_chunk_bytes)
    if len(data) != declared:
        raise ProtocolError(
            f"tensor {name!r} declares shape {list(shape)} ({declared} bytes) "
            f"but carries {len(data)} bytes"
        )

    body = data
    compression = pb.COMPRESSION_NONE
    if policy is not None:
        result = compress_activation(data, policy)
        if result.compressed:
            body = result.body
            compression = pb.COMPRESSION_ZSTD

    tensor = pb.NamedTensor(
        name=name,
        shape=list(shape),
        dtype=dtype,
        byte_order=pb.BYTE_ORDER_LITTLE_ENDIAN,
        total_bytes=len(data),
        compression=compression,
        checksum=checksum_of(data),
    )

    # Fragment the wire body (compressed if we compressed). Offsets walk the
    # wire body so a receiver can verify coverage without assuming arrival
    # order; a zstd frame is not decodable per fragment, so reassembly comes
    # first and decompression happens once, in decode_tensor.
    slices = [body[i : i + max_fragment_bytes] for i in range(0, len(body), max_fragment_bytes)]
    if not slices:
        # A zero-element tensor is legal (e.g. an empty mask) and still needs a
        # fragment, so coverage checks have something to verify.
        slices = [b""]
    if len(slices) > max_fragments:
        raise ProtocolError(
            f"tensor {name!r} needs {len(slices)} fragments, exceeding limit {max_fragments}"
        )

    offset = 0
    for index, piece in enumerate(slices):
        tensor.fragments.append(
            pb.TensorFragment(
                fragment_index=index,
                fragment_count=len(slices),
                byte_offset=offset,
                payload=piece,
            )
        )
        offset += len(piece)
    return tensor


def decode_tensor(
    tensor: pb.NamedTensor,
    *,
    max_chunk_bytes: int = DEFAULT_MAX_CHUNK_BYTES,
    max_fragment_bytes: int = DEFAULT_MAX_FRAGMENT_BYTES,
    max_fragments: int = DEFAULT_MAX_FRAGMENTS_PER_TENSOR,
) -> bytes:
    """Reassemble, decompress and validate a NamedTensor's payload.

    Raises PayloadCorrupt rather than returning a payload it cannot fully
    account for: a hole in the fragments or a bad checksum means the activation
    is not what the sender computed, and continuing would corrupt the route.
    """
    if max_chunk_bytes <= 0 or max_fragment_bytes <= 0 or max_fragments <= 0:
        raise ProtocolError("negotiated byte/count bounds must be positive")
    if tensor.total_bytes > max_chunk_bytes:
        raise ProtocolError(
            f"tensor {tensor.name!r} declares {tensor.total_bytes} bytes, exceeding "
            f"the {max_chunk_bytes}-byte negotiated chunk bound"
        )
    if tensor.byte_order == pb.BYTE_ORDER_BIG_ENDIAN:
        raise ProtocolError(f"tensor {tensor.name!r} is big-endian; wire order is little-endian")
    if tensor.byte_order != pb.BYTE_ORDER_LITTLE_ENDIAN:
        raise ProtocolError(f"tensor {tensor.name!r} declares no byte order")
    declared = expected_bytes(
        tensor.shape, tensor.dtype, max_bytes=max_chunk_bytes
    )
    if declared != tensor.total_bytes:
        raise PayloadCorrupt(
            f"tensor {tensor.name!r} shape {list(tensor.shape)} implies {declared} bytes "
            f"but declares {tensor.total_bytes}"
        )

    if not tensor.fragments:
        raise PayloadCorrupt(f"tensor {tensor.name!r} carries no fragments")
    if len(tensor.fragments) > max_fragments:
        raise PayloadCorrupt(
            f"tensor {tensor.name!r} carries {len(tensor.fragments)} fragments, "
            f"exceeding limit {max_fragments}"
        )
    if any(len(fragment.payload) > max_fragment_bytes for fragment in tensor.fragments):
        raise PayloadCorrupt(
            f"tensor {tensor.name!r} carries a fragment larger than "
            f"{max_fragment_bytes} bytes"
        )
    wire_bytes = sum(len(fragment.payload) for fragment in tensor.fragments)
    if wire_bytes > max_chunk_bytes:
        raise PayloadCorrupt(
            f"tensor {tensor.name!r} wire body exceeds the "
            f"{max_chunk_bytes}-byte negotiated chunk bound"
        )

    fragments = sorted(tensor.fragments, key=lambda f: f.byte_offset)
    count = fragments[0].fragment_count
    if any(f.fragment_count != count for f in fragments):
        raise PayloadCorrupt(f"tensor {tensor.name!r} has inconsistent fragment_count")
    if len(fragments) != count:
        raise PayloadCorrupt(
            f"tensor {tensor.name!r} expects {count} fragments but carries {len(fragments)}"
        )
    if {f.fragment_index for f in fragments} != set(range(count)):
        raise PayloadCorrupt(f"tensor {tensor.name!r} has duplicate or missing fragment indices")

    # Contiguity: offsets must tile the body exactly, with no hole and no overlap.
    body = bytearray()
    for fragment in fragments:
        if fragment.byte_offset != len(body):
            raise PayloadCorrupt(
                f"tensor {tensor.name!r} fragment {fragment.fragment_index} starts at "
                f"{fragment.byte_offset}, expected {len(body)}"
            )
        body.extend(fragment.payload)

    if tensor.compression == pb.COMPRESSION_ZSTD:
        try:
            data = decompress_activation(
                bytes(body), "zstd", max_output_bytes=tensor.total_bytes
            ).body
        except ValueError as exc:
            raise PayloadCorrupt(
                f"tensor {tensor.name!r} has invalid bounded zstd payload"
            ) from exc
    elif tensor.compression == pb.COMPRESSION_NONE:
        data = bytes(body)
    else:
        raise ProtocolError(
            f"tensor {tensor.name!r} uses unspecified or unsupported compression"
        )

    if len(data) != tensor.total_bytes:
        raise PayloadCorrupt(
            f"tensor {tensor.name!r} declares {tensor.total_bytes} bytes but "
            f"reassembled {len(data)}"
        )

    algorithm = tensor.checksum.algorithm
    if algorithm == pb.CHECKSUM_ALGORITHM_CRC32C:
        if tensor.checksum.value != struct.pack(">I", crc32c(data)):
            raise PayloadCorrupt(f"tensor {tensor.name!r} failed its CRC32C check")
    elif algorithm != pb.CHECKSUM_ALGORITHM_NONE:
        raise ProtocolError(
            f"tensor {tensor.name!r} uses unspecified or unsupported checksum algorithm"
        )

    return data


def encode_bundle(
    tensors: Iterable[pb.NamedTensor],
    *,
    max_chunk_bytes: int = DEFAULT_MAX_CHUNK_BYTES,
    max_tensors: int = DEFAULT_MAX_TENSORS_PER_BUNDLE,
) -> pb.TensorBundle:
    if max_chunk_bytes <= 0 or max_tensors <= 0:
        raise ProtocolError("bundle byte/count bounds must be positive")
    tensor_list = list(tensors)
    if len(tensor_list) > max_tensors:
        raise ProtocolError(
            f"bundle carries {len(tensor_list)} tensors, exceeding limit {max_tensors}"
        )
    bundle = pb.TensorBundle(bundle_version=BUNDLE_VERSION, tensors=tensor_list)
    if bundle.ByteSize() > max_chunk_bytes:
        raise ProtocolError(
            f"serialized tensor bundle exceeds the {max_chunk_bytes}-byte "
            "negotiated chunk bound"
        )
    return bundle


def decode_bundle(
    bundle: pb.TensorBundle,
    *,
    max_chunk_bytes: int = DEFAULT_MAX_CHUNK_BYTES,
    max_fragment_bytes: int = DEFAULT_MAX_FRAGMENT_BYTES,
    max_fragments: int = DEFAULT_MAX_FRAGMENTS_PER_TENSOR,
    max_tensors: int = DEFAULT_MAX_TENSORS_PER_BUNDLE,
) -> dict[str, bytes]:
    """Validate every tensor in a bundle and return name -> payload."""
    if bundle.bundle_version != BUNDLE_VERSION:
        raise ProtocolError(
            f"bundle version {bundle.bundle_version} is not supported by this build "
            f"({BUNDLE_VERSION})"
        )
    if (
        max_chunk_bytes <= 0
        or max_fragment_bytes <= 0
        or max_fragments <= 0
        or max_tensors <= 0
    ):
        raise ProtocolError("negotiated byte/count bounds must be positive")
    if len(bundle.tensors) > max_tensors:
        raise ProtocolError(
            f"bundle carries {len(bundle.tensors)} tensors, exceeding limit {max_tensors}"
        )
    if bundle.ByteSize() > max_chunk_bytes:
        raise ProtocolError(
            f"serialized tensor bundle exceeds the {max_chunk_bytes}-byte "
            "negotiated chunk bound"
        )
    payloads: dict[str, bytes] = {}
    for tensor in bundle.tensors:
        if not tensor.name:
            raise ProtocolError("bundle carries an unnamed tensor")
        if tensor.name in payloads:
            raise ProtocolError(f"bundle carries duplicate tensor {tensor.name!r}")
        payloads[tensor.name] = decode_tensor(
            tensor,
            max_chunk_bytes=max_chunk_bytes,
            max_fragment_bytes=max_fragment_bytes,
            max_fragments=max_fragments,
        )
    return payloads


def validate_session_message_size(
    message: pb.SessionRequest | pb.SessionResponse,
    *,
    max_chunk_bytes: int = DEFAULT_MAX_CHUNK_BYTES,
) -> int:
    """Reject an oversized complete stream frame, including protobuf overhead.

    Bundle validation alone is insufficient because the envelope and oneof
    framing are part of the same gRPC message. Senders call this immediately
    before writing; receivers configure gRPC's receive limit to the same value
    and call it again before semantic decoding.
    """
    if max_chunk_bytes <= 0:
        raise ProtocolError("max_chunk_bytes must be positive")
    if not isinstance(message, (pb.SessionRequest, pb.SessionResponse)):
        raise ProtocolError("size validation requires a session request or response")
    size = message.ByteSize()
    if size > max_chunk_bytes:
        raise ProtocolError(
            f"serialized session message is {size} bytes, exceeding the "
            f"{max_chunk_bytes}-byte negotiated chunk bound"
        )
    return size


# --- Bounded prefill chunking ----------------------------------------------


@dataclass(frozen=True)
class PrefillChunk:
    """One token-aligned slice of a prefill."""

    chunk_index: int
    chunk_count: int
    first_position: int
    token_count: int

    @property
    def final_chunk(self) -> bool:
        return self.chunk_index == self.chunk_count - 1

    def chunk_info(self) -> pb.ChunkInfo:
        return pb.ChunkInfo(
            chunk_index=self.chunk_index,
            chunk_count=self.chunk_count,
            final_chunk=self.final_chunk,
        )

    def position(self) -> pb.PositionSpan:
        return pb.PositionSpan(
            first_position=self.first_position, token_count=self.token_count
        )


def plan_prefill_chunks(
    total_tokens: int,
    *,
    first_position: int = 0,
    max_tokens: int = DEFAULT_MAX_PREFILL_CHUNK_TOKENS,
) -> list[PrefillChunk]:
    """Split a prefill into bounded, token-aligned chunks.

    Splits fall on token boundaries only (ADR-0008): a fragment of a token's
    hidden state is not a thing a receiver can execute.
    """
    if total_tokens <= 0:
        raise ProtocolError("a prefill must carry at least one token")
    if max_tokens <= 0:
        raise ProtocolError("max_tokens must be positive")

    count = (total_tokens + max_tokens - 1) // max_tokens
    chunks = []
    for index in range(count):
        offset = index * max_tokens
        chunks.append(
            PrefillChunk(
                chunk_index=index,
                chunk_count=count,
                first_position=first_position + offset,
                token_count=min(max_tokens, total_tokens - offset),
            )
        )
    return chunks


def default_flow_control() -> pb.FlowControl:
    return pb.FlowControl(
        credits_granted=DEFAULT_MAX_INFLIGHT_CHUNKS,
        max_inflight_chunks=DEFAULT_MAX_INFLIGHT_CHUNKS,
        max_chunk_bytes=DEFAULT_MAX_CHUNK_BYTES,
        max_prefill_chunk_tokens=DEFAULT_MAX_PREFILL_CHUNK_TOKENS,
    )


def negotiate_flow_control(
    proposed: pb.FlowControl, limits: pb.FlowControl
) -> pb.FlowControl:
    """Settle a stream's limits: the strictest bound of either peer wins.

    Taking the minimum means neither peer can raise the other's ceiling, so a
    misconfigured — or hostile — sender cannot talk a worker into unbounded
    queues by proposing a large window.
    """

    def _min(a: int, b: int, fallback: int) -> int:
        candidates = [v for v in (a, b) if v > 0]
        return min(candidates) if candidates else fallback

    max_inflight_chunks = _min(
        proposed.max_inflight_chunks,
        limits.max_inflight_chunks,
        DEFAULT_MAX_INFLIGHT_CHUNKS,
    )
    credits_granted = min(
        _min(
            proposed.credits_granted,
            limits.credits_granted,
            DEFAULT_MAX_INFLIGHT_CHUNKS,
        ),
        max_inflight_chunks,
    )
    return pb.FlowControl(
        credits_granted=credits_granted,
        max_inflight_chunks=max_inflight_chunks,
        max_chunk_bytes=_min(
            proposed.max_chunk_bytes, limits.max_chunk_bytes, DEFAULT_MAX_CHUNK_BYTES
        ),
        max_prefill_chunk_tokens=_min(
            proposed.max_prefill_chunk_tokens,
            limits.max_prefill_chunk_tokens,
            DEFAULT_MAX_PREFILL_CHUNK_TOKENS,
        ),
    )