"""Versioned activation-stream envelope for shard hops.

The transport still moves raw bytes over HTTP/WebSocket, but the payload now has
a stable, extensible envelope that names each tensor, preserves unknown fields,
and can round-trip deterministically across direct and relayed hops.
"""

from __future__ import annotations

from dataclasses import dataclass, field
import base64
import hashlib
import json
from typing import Any

SCHEMA_NAME = "meshnet.activation-stream"
SCHEMA_VERSION = 1
DEFAULT_FRAGMENT_BYTES = 64 * 1024


def _canonical_json(data: Any) -> bytes:
    return json.dumps(data, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode("utf-8")


def _sha256_hex(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def _normalize_shape(shape: list[int] | tuple[int, ...]) -> list[int]:
    normalized = [int(dim) for dim in shape]
    if not normalized or any(dim <= 0 for dim in normalized):
        raise ValueError("shape must be a non-empty list of positive integers")
    return normalized


def _fragment_bytes(body: bytes, max_fragment_bytes: int) -> tuple[bytes, ...]:
    if max_fragment_bytes <= 0:
        raise ValueError("max_fragment_bytes must be positive")
    if not body:
        return (b"",)
    return tuple(body[offset : offset + max_fragment_bytes] for offset in range(0, len(body), max_fragment_bytes))


@dataclass(frozen=True)
class TensorFragment:
    """One bounded chunk of a named tensor."""

    offset: int
    body: bytes
    checksum: str
    compression: str = "identity"
    extensions: dict[str, Any] = field(default_factory=dict)

    @classmethod
    def from_bytes(
        cls,
        body: bytes,
        *,
        offset: int,
        compression: str = "identity",
        extensions: dict[str, Any] | None = None,
    ) -> "TensorFragment":
        return cls(
            offset=int(offset),
            body=bytes(body),
            checksum=_sha256_hex(body),
            compression=compression,
            extensions=dict(extensions or {}),
        )

    def to_dict(self) -> dict[str, Any]:
        data = {
            "offset": self.offset,
            "compression": self.compression,
            "checksum": self.checksum,
            "body_base64": base64.b64encode(self.body).decode("ascii"),
        }
        data.update(self.extensions)
        return data

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "TensorFragment":
        known = {"offset", "compression", "checksum", "body_base64"}
        body = base64.b64decode(data.get("body_base64", ""))
        fragment = cls(
            offset=int(data["offset"]),
            body=body,
            checksum=str(data.get("checksum") or _sha256_hex(body)),
            compression=str(data.get("compression") or "identity"),
            extensions={k: v for k, v in data.items() if k not in known},
        )
        fragment.validate()
        return fragment

    def validate(self) -> None:
        if self.checksum != _sha256_hex(self.body):
            raise ValueError("fragment checksum mismatch")
        if self.offset < 0:
            raise ValueError("fragment offset must be non-negative")


@dataclass(frozen=True)
class NamedTensor:
    """A tensor named within a versioned activation envelope."""

    name: str
    shape: list[int]
    dtype: str
    byte_order: str
    checksum: str
    fragments: tuple[TensorFragment, ...]
    compression: str = "identity"
    extensions: dict[str, Any] = field(default_factory=dict)

    @classmethod
    def from_bytes(
        cls,
        *,
        name: str,
        body: bytes,
        shape: list[int] | tuple[int, ...],
        dtype: str,
        byte_order: str = "little",
        compression: str = "identity",
        max_fragment_bytes: int = DEFAULT_FRAGMENT_BYTES,
        extensions: dict[str, Any] | None = None,
    ) -> "NamedTensor":
        normalized_shape = _normalize_shape(shape)
        fragments = tuple(
            TensorFragment.from_bytes(fragment, offset=offset, compression=compression)
            for offset, fragment in enumerate(_fragment_bytes(body, max_fragment_bytes))
            for offset in (offset * max_fragment_bytes,)
        )
        return cls(
            name=str(name),
            shape=normalized_shape,
            dtype=str(dtype),
            byte_order=str(byte_order),
            checksum=_sha256_hex(body),
            fragments=fragments,
            compression=compression,
            extensions=dict(extensions or {}),
        )

    def body(self) -> bytes:
        ordered = sorted(self.fragments, key=lambda frag: frag.offset)
        body = b"".join(fragment.body for fragment in ordered)
        if _sha256_hex(body) != self.checksum:
            raise ValueError(f"tensor {self.name!r} checksum mismatch")
        return body

    def validate(self) -> None:
        for fragment in self.fragments:
            fragment.validate()
        self.body()

    def to_dict(self) -> dict[str, Any]:
        data = {
            "name": self.name,
            "shape": list(self.shape),
            "dtype": self.dtype,
            "byte_order": self.byte_order,
            "compression": self.compression,
            "checksum": self.checksum,
            "fragments": [fragment.to_dict() for fragment in self.fragments],
        }
        data.update(self.extensions)
        return data

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "NamedTensor":
        known = {
            "name",
            "shape",
            "dtype",
            "byte_order",
            "compression",
            "checksum",
            "fragments",
        }
        tensor = cls(
            name=str(data["name"]),
            shape=_normalize_shape(list(data["shape"])),
            dtype=str(data["dtype"]),
            byte_order=str(data.get("byte_order", "little")),
            compression=str(data.get("compression", "identity")),
            checksum=str(data["checksum"]),
            fragments=tuple(TensorFragment.from_dict(fragment) for fragment in data.get("fragments", [])),
            extensions={k: v for k, v in data.items() if k not in known},
        )
        tensor.validate()
        return tensor


@dataclass(frozen=True)
class ActivationEnvelope:
    """Versioned envelope for shard activation traffic."""

    request_id: str
    work_id: str
    route_session: str
    route_epoch: int
    shard_start: int
    effective_start: int
    phase: str
    position: int
    idempotency_step: int
    tensors: tuple[NamedTensor, ...]
    version: int = SCHEMA_VERSION
    schema: str = SCHEMA_NAME
    token_id_sideband: list[int] | None = None
    architecture_state: dict[str, Any] | None = None
    recurrent_state: dict[str, Any] | None = None
    mtp: dict[str, Any] | None = None
    extensions: dict[str, Any] = field(default_factory=dict)

    @classmethod
    def from_tensor_payload(
        cls,
        *,
        payload: Any,
        name: str,
        request_id: str,
        work_id: str,
        route_session: str,
        route_epoch: int,
        shard_start: int,
        effective_start: int,
        phase: str,
        position: int,
        idempotency_step: int,
        byte_order: str = "little",
        compression: str = "identity",
        max_fragment_bytes: int = DEFAULT_FRAGMENT_BYTES,
        token_id_sideband: list[int] | None = None,
        architecture_state: dict[str, Any] | None = None,
        recurrent_state: dict[str, Any] | None = None,
        mtp: dict[str, Any] | None = None,
        extensions: dict[str, Any] | None = None,
    ) -> "ActivationEnvelope":
        tensor = NamedTensor.from_bytes(
            name=name,
            body=payload.body,
            shape=payload.shape,
            dtype="bfloat16",
            byte_order=byte_order,
            compression=compression,
            max_fragment_bytes=max_fragment_bytes,
            extensions={
                "attention_mask_header": payload.attention_mask_header,
                "position_ids_header": payload.position_ids_header,
                **({"past_len": payload.past_len} if payload.past_len is not None else {}),
            },
        )
        return cls(
            request_id=request_id,
            work_id=work_id,
            route_session=route_session,
            route_epoch=int(route_epoch),
            shard_start=int(shard_start),
            effective_start=int(effective_start),
            phase=str(phase),
            position=int(position),
            idempotency_step=int(idempotency_step),
            tensors=(tensor,),
            token_id_sideband=list(token_id_sideband) if token_id_sideband is not None else None,
            architecture_state=architecture_state,
            recurrent_state=recurrent_state,
            mtp=mtp,
            extensions=dict(extensions or {}),
        )

    def to_tensor_payload(self, *, tensor_name: str = "activations") -> Any:
        from .model_backend import TensorPayload

        tensor = self.tensor(tensor_name)
        return TensorPayload(
            body=tensor.body(),
            shape=list(tensor.shape),
            attention_mask_header=tensor.extensions.get("attention_mask_header"),
            position_ids_header=tensor.extensions.get("position_ids_header"),
            past_len=tensor.extensions.get("past_len"),
        )

    def tensor(self, name: str = "activations") -> NamedTensor:
        for tensor in self.tensors:
            if tensor.name == name:
                return tensor
        raise KeyError(name)

    def to_dict(self) -> dict[str, Any]:
        data = {
            "schema": self.schema,
            "version": self.version,
            "request_id": self.request_id,
            "work_id": self.work_id,
            "route_session": self.route_session,
            "route_epoch": self.route_epoch,
            "shard_start": self.shard_start,
            "effective_start": self.effective_start,
            "phase": self.phase,
            "position": self.position,
            "idempotency_step": self.idempotency_step,
            "tensors": [tensor.to_dict() for tensor in self.tensors],
        }
        if self.token_id_sideband is not None:
            data["token_id_sideband"] = list(self.token_id_sideband)
        if self.architecture_state is not None:
            data["architecture_state"] = self.architecture_state
        if self.recurrent_state is not None:
            data["recurrent_state"] = self.recurrent_state
        if self.mtp is not None:
            data["mtp"] = self.mtp
        data.update(self.extensions)
        return data

    def to_bytes(self, *, max_bytes: int | None = None) -> bytes:
        raw = _canonical_json(self.to_dict())
        if max_bytes is not None and len(raw) > max_bytes:
            raise ValueError("activation envelope exceeds the size limit")
        return raw

    @classmethod
    def from_bytes(cls, data: bytes) -> "ActivationEnvelope":
        payload = json.loads(data)
        if not isinstance(payload, dict):
            raise ValueError("activation envelope must be a JSON object")
        known = {
            "schema",
            "version",
            "request_id",
            "work_id",
            "route_session",
            "route_epoch",
            "shard_start",
            "effective_start",
            "phase",
            "position",
            "idempotency_step",
            "tensors",
            "token_id_sideband",
            "architecture_state",
            "recurrent_state",
            "mtp",
        }
        envelope = cls(
            schema=str(payload.get("schema", SCHEMA_NAME)),
            version=int(payload.get("version", SCHEMA_VERSION)),
            request_id=str(payload["request_id"]),
            work_id=str(payload["work_id"]),
            route_session=str(payload["route_session"]),
            route_epoch=int(payload["route_epoch"]),
            shard_start=int(payload["shard_start"]),
            effective_start=int(payload["effective_start"]),
            phase=str(payload["phase"]),
            position=int(payload["position"]),
            idempotency_step=int(payload["idempotency_step"]),
            tensors=tuple(NamedTensor.from_dict(item) for item in payload.get("tensors", [])),
            token_id_sideband=payload.get("token_id_sideband"),
            architecture_state=payload.get("architecture_state"),
            recurrent_state=payload.get("recurrent_state"),
            mtp=payload.get("mtp"),
            extensions={k: v for k, v in payload.items() if k not in known},
        )
        envelope.validate()
        return envelope

    def validate(self) -> None:
        if self.version != SCHEMA_VERSION:
            raise ValueError("unsupported activation envelope version")
        if self.schema != SCHEMA_NAME:
            raise ValueError("unsupported activation envelope schema")
        if self.phase not in {"prefill", "decode"}:
            raise ValueError("phase must be prefill or decode")
        for tensor in self.tensors:
            tensor.validate()