"""Canonical conformance vectors for the native Shard protocol.

Two independently-written codecs that each round-trip their own output prove
nothing about each other. These vectors are the shared reference: Python builds
the canonical message, the bytes are committed under
`packages/node/native/testdata/`, and the C++ test parses those exact bytes and
asserts the same field values. A change that alters the wire meaning of a field
breaks the vector in both languages instead of drifting silently in one.

The vector deliberately exercises every field group the protocol promises to
carry — identity, epoch, fingerprint, range, phase, position, idempotency,
cache expectation, deadline, chunking, compression, checksum and a multi-
fragment named tensor — so it doubles as an executable inventory of the
contract.
"""

from __future__ import annotations

import pathlib

from . import codec
from .generated import shard_runtime_pb2 as pb

# Committed vectors live beside the schema, under `packages/node/native/`.
# parents[2] is `packages/node`: native_protocol -> meshnet_node -> node.
TESTDATA_DIR = pathlib.Path(__file__).resolve().parents[2] / "native/testdata"

GOLDEN_SESSION_REQUEST = "session_request_golden.binpb"
GOLDEN_CAPABILITY_REPORT = "capability_report_golden.binpb"
GOLDEN_DECODE_STEP = "decode_step_golden.binpb"

# Written by the C++ conformance test into its build tree; the Python test picks
# it up when present to prove the two languages agree byte-for-byte.
CPP_ROUNDTRIP = "cpp_roundtrip.binpb"

# Fixed, non-default values. Every one is chosen to be distinguishable from a
# proto3 default so an unset field can never masquerade as a correct one.
WORK_ID = "work-7f3a"
ROUTE_SESSION_ID = "rs-2b91"
ROUTE_EPOCH = 7
IDEMPOTENCY_STEP = 42
FIRST_POSITION = 256
TOKEN_COUNT = 128
EXPECTED_PAST_LEN = 256
DEADLINE_UNIX_NANOS = 1_800_000_000_000_000_000
MODEL_ARTIFACT_DIGEST = "sha256:1f0c9d2e"
RUNTIME_RECIPE_DIGEST = "sha256:ab77e410"
RECIPE_ID = "llama-gguf-q4km-rocm"
RECIPE_VERSION = "3"
CATALOGUE_VERSION = "2026.07.1"
START_LAYER = 12
END_LAYER = 24
EFFECTIVE_START_LAYER = 16
HIDDEN_SIZE = 8

# A payload big enough to force more than one fragment at the bound below, so
# the vector actually exercises reassembly rather than the one-fragment path.
FRAGMENT_BYTES = 64
TENSOR_SHAPE = [1, TOKEN_COUNT, HIDDEN_SIZE]


def canonical_payload() -> bytes:
    """Deterministic bfloat16-sized payload for the canonical tensor."""
    total = codec.expected_bytes(TENSOR_SHAPE, pb.DTYPE_BFLOAT16)
    return bytes((i * 7 + 11) % 256 for i in range(total))


def canonical_session_request() -> pb.SessionRequest:
    """The canonical prefill chunk carried on a session stream."""
    tensor = codec.encode_tensor(
        codec.HIDDEN_STATES,
        canonical_payload(),
        TENSOR_SHAPE,
        pb.DTYPE_BFLOAT16,
        max_fragment_bytes=FRAGMENT_BYTES,
    )
    envelope = pb.Envelope(
        schema_version=pb.SCHEMA_VERSION_1,
        work_id=WORK_ID,
        route_session_id=ROUTE_SESSION_ID,
        route_epoch=ROUTE_EPOCH,
        fingerprint=pb.Fingerprint(
            model_artifact_digest=MODEL_ARTIFACT_DIGEST,
            runtime_recipe_digest=RUNTIME_RECIPE_DIGEST,
            recipe_id=RECIPE_ID,
            recipe_version=RECIPE_VERSION,
            catalogue_version=CATALOGUE_VERSION,
        ),
        shard_range=pb.ShardRange(
            start_layer=START_LAYER,
            end_layer=END_LAYER,
            effective_start_layer=EFFECTIVE_START_LAYER,
        ),
        phase=pb.PHASE_PREFILL,
        position=pb.PositionSpan(
            first_position=FIRST_POSITION, token_count=TOKEN_COUNT
        ),
        idempotency_step=IDEMPOTENCY_STEP,
        cache_expectation=pb.CacheExpectation(
            mode=pb.CACHE_MODE_PREFILL, expected_past_len=EXPECTED_PAST_LEN
        ),
        deadline_unix_nanos=DEADLINE_UNIX_NANOS,
        chunk=pb.ChunkInfo(chunk_index=1, chunk_count=3, final_chunk=False),
    )
    return pb.SessionRequest(
        chunk=pb.ActivationChunk(
            envelope=envelope, bundle=codec.encode_bundle([tensor])
        )
    )


def canonical_capability_report() -> pb.CapabilityReport:
    """The canonical capability report a worker answers admission with."""
    return pb.CapabilityReport(
        schema_version=pb.SCHEMA_VERSION_1,
        fingerprint=pb.Fingerprint(
            model_artifact_digest=MODEL_ARTIFACT_DIGEST,
            runtime_recipe_digest=RUNTIME_RECIPE_DIGEST,
            recipe_id=RECIPE_ID,
            recipe_version=RECIPE_VERSION,
            catalogue_version=CATALOGUE_VERSION,
        ),
        shard_range=pb.ShardRange(
            start_layer=START_LAYER,
            end_layer=END_LAYER,
            effective_start_layer=EFFECTIVE_START_LAYER,
        ),
        backend="rocm",
        device="gfx1151",
        validated=True,
        max_concurrent_sessions=4,
        max_context_tokens=8192,
        flow_control=codec.default_flow_control(),
        accepted_compression=[pb.COMPRESSION_NONE, pb.COMPRESSION_ZSTD],
        supported_schema_versions=[pb.SCHEMA_VERSION_1],
        validated_at_unix_nanos=DEADLINE_UNIX_NANOS,
    )


def canonical_decode_step() -> pb.SessionRequest:
    """The DGR-006 multi-tensor decode boundary vector."""
    hidden = codec.encode_tensor(
        codec.HIDDEN_STATES, bytes(range(16)), [1, 1, 4], pb.DTYPE_FLOAT32
    )
    index_topk = codec.encode_tensor(
        "index_topk", (3).to_bytes(4, "little"), [1], pb.DTYPE_INT32
    )
    return pb.SessionRequest(
        decode=pb.DecodeStep(
            idempotency_step=43,
            position=384,
            expected_past_len=384,
            work_id="decode-7f3a",
            deadline_unix_nanos=DEADLINE_UNIX_NANOS,
            bundle=codec.encode_bundle(
                [hidden, index_topk],
                architecture=pb.ARCHITECTURE_TYPE_MLA,
                boundary_point="pre_tail_residual",
            ),
        )
    )


def serialize(message) -> bytes:
    """Serialize deterministically, so committed golden bytes are stable."""
    return message.SerializeToString(deterministic=True)