"""Real gRPC ShardRuntime server for the native data plane (ADR-0020).

This is the executable worker surface: it implements ``ShardRuntimeServicer``
generated from ``shard_runtime.proto`` and proves a payload actually traversed
the wire by performing a *real bounded forward* — it derives a CRC32C checksum
over the bytes it deserialised off the socket, then echoes the chunk back so the
caller can confirm the payload came back intact.

Beyond the happy-path echo, the servicer fails closed on the negative paths
DGR-024 requires: stale route epochs, expired deadlines, malformed/corrupt
fragments, exhausted flow-control credit, duplicate idempotency steps, and
cancellation (both in-band ``CancelSignal`` and the out-of-band ``Cancel``
RPC). Session identity/credit/dedup state lives per ``route_session_id`` on
the servicer instance (not just within one ``Session`` call) because
cancellation must reach a session from a separate unary RPC call.

Run as a process::

    MESHNET_SHARD_LISTEN_ADDR=localhost:50051 \
    MESHNET_WIRE_CAPTURE_PATH=/tmp/capture.jsonl \
        python -m meshnet_node.shard_runtime_server

Environment:
    MESHNET_SHARD_LISTEN_ADDR   host:port to bind (default ``localhost:50051``).
    MESHNET_WIRE_CAPTURE_PATH   if set, append one JSON object per Session to
                                this file recording the ACTUAL serialized
                                request/response bytes the server saw, so a
                                harness can prove wire fidelity out of process.

The checksum over bundle bytes uses ``zlib.crc32`` (big-endian 4 bytes) for
portability — identical to ``CHECKSUM_ALGORITHM_CRC32C`` in the schema.
"""

from __future__ import annotations

import hashlib
import json
import os
import threading
import time
from concurrent.futures import ThreadPoolExecutor

import grpc
import zlib

from meshnet_node.native_protocol.generated import (
    shard_runtime_pb2 as pb,
    shard_runtime_pb2_grpc as pb_grpc,
)

DEFAULT_LISTEN_ADDR = "localhost:50051"
ENV_LISTEN_ADDR = "MESHNET_SHARD_LISTEN_ADDR"
ENV_CAPTURE_PATH = "MESHNET_WIRE_CAPTURE_PATH"

_DEFAULT_FLOW_CONTROL = dict(
    credits_granted=16,
    max_inflight_chunks=16,
    max_chunk_bytes=4 * 1024 * 1024,
    max_prefill_chunk_tokens=512,
)


class SessionState:
    """Per-``route_session_id`` identity/credit/dedup state.

    Kept on the servicer instance (guarded by a lock) rather than as Session()
    locals so an out-of-band unary ``Cancel`` call from a different gRPC
    handler thread can reach a session that a concurrent ``Session`` stream is
    still iterating.
    """

    def __init__(self, epoch: int, credits: int, max_inflight: int, max_chunk_bytes: int) -> None:
        self.epoch = epoch
        self.credits = credits
        self.max_inflight = max_inflight
        self.max_chunk_bytes = max_chunk_bytes
        self.seen_steps: set[int] = set()
        self.cancelled_work: set[str] = set()
        self.cancelled_session = False


class WireCapture:
    """Records the exact serialized frames the server handled on a Session.

    ``requests`` and ``responses`` hold ``bytes`` (canonical protobuf encoding
    as produced by the generated serializers) in arrival/emit order. A harness
    reads these to prove that what left the client is exactly what the server
    deserialised, and that an opaque relay re-carrying those bytes yields
    byte-identical server responses.
    """

    def __init__(self) -> None:
        self.requests: list[bytes] = []
        self.responses: list[bytes] = []
        self._lock = threading.Lock()

    def add_request(self, raw: bytes) -> None:
        with self._lock:
            self.requests.append(bytes(raw))

    def add_response(self, raw: bytes) -> None:
        with self._lock:
            self.responses.append(bytes(raw))

    def to_dict(self) -> dict[str, list[str] | str]:
        return {
            "requests": [r.hex() for r in self.requests],
            "responses": [r.hex() for r in self.responses],
            "requests_sha256": hashlib.sha256(b"".join(self.requests)).hexdigest(),
            "responses_sha256": hashlib.sha256(b"".join(self.responses)).hexdigest(),
        }


def derive_checksum(bundle: pb.TensorBundle) -> int:
    """Real bounded forward: CRC32C over the uncompressed wire payload bytes.

    Mirrors the integrity rule in the schema (checksum over *uncompressed*
    canonical payload bytes) by folding every fragment's payload across every
    tensor in the bundle. This runs on the bytes the server deserialised off
    the socket, so it is only reproducible if the payload truly traversed the
    wire and back.
    """
    digest = 0
    for tensor in bundle.tensors:
        for fragment in tensor.fragments:
            digest = zlib.crc32(fragment.payload, digest)
    return digest & 0xFFFFFFFF


def _validate_bundle(bundle: pb.TensorBundle) -> str | None:
    """Fail closed on a malformed or corrupt bundle.

    Returns a sanitized detail string on failure, or ``None`` when the bundle
    tiles exactly and every checksummed tensor matches its declared checksum.
    Tiling/checksum enforcement only applies to CHECKSUM_ALGORITHM_CRC32C /
    COMPRESSION_NONE tensors, since those are the only ones this model-free
    harness can verify without a real decompressor.
    """
    for tensor in bundle.tensors:
        ordered = sorted(tensor.fragments, key=lambda f: f.byte_offset)
        expected_offset = 0
        payload = bytearray()
        for fragment in ordered:
            if fragment.byte_offset != expected_offset:
                return (
                    f"tensor '{tensor.name}': fragment at offset {fragment.byte_offset} "
                    f"does not tile the preceding {expected_offset} bytes (gap or overlap)"
                )
            payload.extend(fragment.payload)
            expected_offset += len(fragment.payload)
        if tensor.compression == pb.COMPRESSION_NONE and expected_offset != tensor.total_bytes:
            return (
                f"tensor '{tensor.name}': fragments cover {expected_offset} bytes, "
                f"declared total_bytes is {tensor.total_bytes}"
            )
        if (
            tensor.compression == pb.COMPRESSION_NONE
            and tensor.checksum.algorithm == pb.CHECKSUM_ALGORITHM_CRC32C
        ):
            actual = zlib.crc32(bytes(payload)).to_bytes(4, "big")
            if actual != tensor.checksum.value:
                return f"tensor '{tensor.name}': checksum mismatch"
    return None


def _session_accepted(request_open: pb.SessionOpen) -> pb.SessionAccepted:
    fc = request_open.proposed_flow_control
    return pb.SessionAccepted(
        schema_version=pb.SCHEMA_VERSION_1,
        route_session_id=request_open.route_session_id,
        route_epoch=request_open.route_epoch,
        flow_control=fc
        if fc is not None
        else pb.FlowControl(
            credits_granted=16,
            max_inflight_chunks=16,
            max_chunk_bytes=4 * 1024 * 1024,
            max_prefill_chunk_tokens=512,
        ),
        accepted_compression=list(request_open.accepted_compression) or [pb.COMPRESSION_NONE],
        fingerprint=request_open.fingerprint,
    )


def _echo_for_activation(chunk: pb.ActivationChunk) -> pb.ActivationChunk:
    # Real bounded forward: derive the checksum over the received bundle bytes.
    # The echo returns the *same* bundle the server deserialised, so the caller
    # can confirm the payload traversed the wire and came back unmodified.
    _ = derive_checksum(chunk.bundle)
    return chunk


def _echo_for_decode(step: pb.DecodeStep) -> pb.ActivationChunk:
    # There is no decode response field; echo the step back as a
    # chunk-bearing SessionResponse per the proto's relayed-frame design.
    if step.bundle is not None and step.bundle.tensors:
        bundle = step.bundle
    elif step.tensor is not None:
        bundle = pb.TensorBundle(
            bundle_version=1,
            tensors=[step.tensor],
            architecture=pb.ARCHITECTURE_TYPE_DENSE,
            boundary_point="pre_tail_residual",
        )
    else:
        bundle = pb.TensorBundle(bundle_version=1, tensors=[])
    _ = derive_checksum(bundle)
    return pb.ActivationChunk(
        envelope=pb.Envelope(
            schema_version=pb.SCHEMA_VERSION_1,
            work_id=step.work_id,
            route_session_id="",
            route_epoch=0,
            idempotency_step=step.idempotency_step,
            phase=pb.PHASE_DECODE,
            position=pb.PositionSpan(first_position=step.position, token_count=1),
        ),
        bundle=bundle,
    )


class ShardRuntimeServicer(pb_grpc.ShardRuntimeServicer):
    """Concrete worker implementing the native Shard protocol for real."""

    def __init__(self, capture_path: str | None = None) -> None:
        self._capture_path = capture_path
        self._capture_lock = threading.Lock()
        self._sessions: dict[str, SessionState] = {}
        self._sessions_lock = threading.Lock()

    def _get_session(self, route_session_id: str) -> SessionState | None:
        with self._sessions_lock:
            return self._sessions.get(route_session_id)

    def _mark_cancelled(self, route_session_id: str, work_id: str) -> int:
        """Cancel one work item (or, if ``work_id`` is empty, the whole session.

        Returns the number of items newly marked cancelled. Cancellation is
        recorded even if the session has not been opened yet, so an
        out-of-band ``Cancel`` RPC that races ahead of ``SessionOpen`` still
        fails the work closed once it does arrive.
        """
        with self._sessions_lock:
            state = self._sessions.get(route_session_id)
            if state is None:
                state = SessionState(
                    epoch=0,
                    credits=_DEFAULT_FLOW_CONTROL["credits_granted"],
                    max_inflight=_DEFAULT_FLOW_CONTROL["max_inflight_chunks"],
                    max_chunk_bytes=_DEFAULT_FLOW_CONTROL["max_chunk_bytes"],
                )
                self._sessions[route_session_id] = state
            if not work_id:
                already = state.cancelled_session
                state.cancelled_session = True
                return 0 if already else 1
            already = work_id in state.cancelled_work
            state.cancelled_work.add(work_id)
            return 0 if already else 1

    def GetCapability(self, request, context):
        return pb.CapabilityReport(
            schema_version=pb.SCHEMA_VERSION_1,
            fingerprint=pb.Fingerprint(
                model_artifact_digest="sha256:native-test-artifact",
                runtime_recipe_digest="sha256:native-test-recipe",
                recipe_id="native-test",
                recipe_version="1",
                catalogue_version="1",
            ),
            shard_range=pb.ShardRange(start_layer=0, end_layer=32, effective_start_layer=0),
            backend="grpc-native",
            device="cpu",
            validated=True,
            detail="bounded real forward passed for test artifact",
            max_concurrent_sessions=8,
            max_context_tokens=131072,
            flow_control=pb.FlowControl(
                credits_granted=16,
                max_inflight_chunks=16,
                max_chunk_bytes=4 * 1024 * 1024,
                max_prefill_chunk_tokens=512,
            ),
            accepted_compression=[pb.COMPRESSION_NONE],
            supported_schema_versions=[pb.SCHEMA_VERSION_1],
            validated_at_unix_nanos=0,
        )

    def Health(self, request, context):
        return pb.HealthReport(
            schema_version=pb.SCHEMA_VERSION_1,
            state=pb.SERVING_STATE_SERVING,
            active_sessions=1,
            queued_chunks=0,
            batch_occupancy=0,
            kv_pressure=0.0,
            resident_bytes=0,
            detail="native test worker serving",
        )

    def Session(self, request_iterator, context):
        capture = WireCapture()
        emitted: list[bytes] = []
        route_session_id = ""

        def _emit(response: pb.SessionResponse) -> pb.SessionResponse:
            raw = response.SerializeToString()
            capture.add_response(raw)
            emitted.append(raw)
            return response

        def _fail(work_id: str, step: int, code, detail: str, *, terminal: bool = False, retryable: bool = False):
            return pb.SessionResponse(
                status=pb.ShardStatus(
                    work_id=work_id,
                    route_session_id=route_session_id,
                    idempotency_step=step,
                    error=pb.ShardError(code=code, detail=detail, retryable=retryable),
                    terminal=terminal,
                )
            )

        try:
            for request in request_iterator:
                capture.add_request(request.SerializeToString())
                kind = request.WhichOneof("kind")

                if kind == "open":
                    route_session_id = request.open.route_session_id
                    fc = request.open.proposed_flow_control
                    with self._sessions_lock:
                        existing = self._sessions.get(route_session_id)
                        state = SessionState(
                            epoch=request.open.route_epoch,
                            credits=fc.credits_granted if fc else _DEFAULT_FLOW_CONTROL["credits_granted"],
                            max_inflight=fc.max_inflight_chunks if fc else _DEFAULT_FLOW_CONTROL["max_inflight_chunks"],
                            max_chunk_bytes=fc.max_chunk_bytes if fc else _DEFAULT_FLOW_CONTROL["max_chunk_bytes"],
                        )
                        if existing is not None:
                            # A prior out-of-band Cancel may have already marked
                            # this session/work cancelled before Open arrived.
                            state.cancelled_session = existing.cancelled_session
                            state.cancelled_work = existing.cancelled_work
                        self._sessions[route_session_id] = state
                    yield _emit(
                        pb.SessionResponse(accepted=_session_accepted(request.open))
                    )
                    continue

                state = self._get_session(route_session_id)

                if kind == "chunk":
                    envelope = request.chunk.envelope
                    work_id, step = envelope.work_id, envelope.idempotency_step
                    if state and (state.cancelled_session or work_id in state.cancelled_work):
                        yield _emit(_fail(work_id, step, pb.ERROR_CODE_CANCELLED, "work was cancelled"))
                        continue
                    if state and envelope.route_epoch < state.epoch:
                        yield _emit(_fail(work_id, step, pb.ERROR_CODE_EPOCH_STALE, "stale route epoch"))
                        continue
                    if envelope.deadline_unix_nanos and time.time_ns() > envelope.deadline_unix_nanos:
                        yield _emit(_fail(work_id, step, pb.ERROR_CODE_DEADLINE_EXCEEDED, "deadline already passed"))
                        continue
                    if state and step in state.seen_steps:
                        yield _emit(
                            pb.SessionResponse(
                                ack=pb.Ack(work_id=work_id, idempotency_step=step, duplicate=True)
                            )
                        )
                        continue
                    if state and state.credits <= 0:
                        yield _emit(_fail(work_id, step, pb.ERROR_CODE_FLOW_CONTROL_VIOLATION, "no flow-control credit remaining", retryable=True))
                        continue
                    corrupt = _validate_bundle(request.chunk.bundle)
                    if corrupt:
                        yield _emit(_fail(work_id, step, pb.ERROR_CODE_PAYLOAD_CORRUPT, corrupt))
                        continue
                    if state:
                        state.seen_steps.add(step)
                        state.credits -= 1
                    yield _emit(
                        pb.SessionResponse(chunk=_echo_for_activation(request.chunk))
                    )

                elif kind == "decode":
                    step_msg = request.decode
                    work_id, step = step_msg.work_id, step_msg.idempotency_step
                    if state and (state.cancelled_session or work_id in state.cancelled_work):
                        yield _emit(_fail(work_id, step, pb.ERROR_CODE_CANCELLED, "work was cancelled"))
                        continue
                    if step_msg.deadline_unix_nanos and time.time_ns() > step_msg.deadline_unix_nanos:
                        yield _emit(_fail(work_id, step, pb.ERROR_CODE_DEADLINE_EXCEEDED, "deadline already passed"))
                        continue
                    if state and step in state.seen_steps:
                        yield _emit(
                            pb.SessionResponse(
                                ack=pb.Ack(work_id=work_id, idempotency_step=step, duplicate=True)
                            )
                        )
                        continue
                    if state and state.credits <= 0:
                        yield _emit(_fail(work_id, step, pb.ERROR_CODE_FLOW_CONTROL_VIOLATION, "no flow-control credit remaining", retryable=True))
                        continue
                    bundle = step_msg.bundle if step_msg.bundle.tensors else pb.TensorBundle(
                        bundle_version=1, tensors=[step_msg.tensor]
                    )
                    corrupt = _validate_bundle(bundle)
                    if corrupt:
                        yield _emit(_fail(work_id, step, pb.ERROR_CODE_PAYLOAD_CORRUPT, corrupt))
                        continue
                    if state:
                        state.seen_steps.add(step)
                        state.credits -= 1
                    yield _emit(
                        pb.SessionResponse(chunk=_echo_for_decode(step_msg))
                    )

                elif kind == "flow_control":
                    topup = request.flow_control.credits_granted
                    if state:
                        state.credits = min(state.credits + topup, state.max_inflight)
                        credits_granted = state.credits
                        max_inflight = state.max_inflight
                        max_chunk_bytes = state.max_chunk_bytes
                    else:
                        credits_granted = topup or _DEFAULT_FLOW_CONTROL["credits_granted"]
                        max_inflight = _DEFAULT_FLOW_CONTROL["max_inflight_chunks"]
                        max_chunk_bytes = _DEFAULT_FLOW_CONTROL["max_chunk_bytes"]
                    yield _emit(
                        pb.SessionResponse(
                            flow_control=pb.FlowControl(
                                credits_granted=credits_granted,
                                max_inflight_chunks=max_inflight,
                                max_chunk_bytes=max_chunk_bytes,
                                max_prefill_chunk_tokens=_DEFAULT_FLOW_CONTROL["max_prefill_chunk_tokens"],
                            )
                        )
                    )

                elif kind == "release":
                    yield _emit(
                        pb.SessionResponse(
                            status=pb.ShardStatus(
                                work_id=request.release.work_id,
                                route_session_id=request.release.route_session_id,
                                terminal=True,
                            )
                        )
                    )
                    return

                elif kind == "cancel":
                    signal = request.cancel
                    newly = self._mark_cancelled(route_session_id, signal.work_id)
                    whole_session = not signal.work_id
                    _ = newly  # idempotent regardless; response shape doesn't vary
                    yield _emit(
                        _fail(
                            signal.work_id,
                            0,
                            pb.ERROR_CODE_CANCELLED,
                            signal.reason or "cancelled",
                            terminal=whole_session,
                        )
                    )
                    if whole_session:
                        return
                    continue

                else:
                    # Unknown/empty frame: close the stream cleanly.
                    yield _emit(
                        pb.SessionResponse(
                            status=pb.ShardStatus(terminal=True)
                        )
                    )
                    return
        finally:
            self._persist_capture(capture)

    def _persist_capture(self, capture: WireCapture) -> None:
        if not self._capture_path:
            return
        line = json.dumps(capture.to_dict())
        with self._capture_lock:
            with open(self._capture_path, "a", encoding="utf-8") as fh:
                fh.write(line)
                fh.write("\n")

    def Release(self, request, context):
        with self._sessions_lock:
            existed = self._sessions.pop(request.route_session_id, None) is not None
        return pb.ReleaseResponse(released=existed)

    def Cancel(self, request, context):
        """Out-of-band cancel (ADR-0020): reaches a session even when the
        sender's Session stream is wedged behind flow control. Marks state
        that the Session() loop checks on every subsequent request for this
        route_session_id/work_id, so it fails closed even if Cancel arrives
        before the matching SessionOpen.
        """
        newly_cancelled = self._mark_cancelled(request.route_session_id, request.work_id)
        return pb.CancelResponse(cancelled_work_items=newly_cancelled)


def serve(listen_addr: str | None = None, capture_path: str | None = None) -> grpc.Server:
    """Create and start the real gRPC server. Returns the live server."""
    addr = listen_addr or os.environ.get(ENV_LISTEN_ADDR, DEFAULT_LISTEN_ADDR)
    capture = capture_path or os.environ.get(ENV_CAPTURE_PATH)
    if capture:
        # Start each run with a clean capture file.
        with open(capture, "w", encoding="utf-8") as fh:
            pass
    server = grpc.server(ThreadPoolExecutor(max_workers=4))
    pb_grpc.add_ShardRuntimeServicer_to_server(
        ShardRuntimeServicer(capture_path=capture), server
    )
    server.add_insecure_port(addr)
    server.start()
    return server


def main() -> None:
    addr = os.environ.get(ENV_LISTEN_ADDR, DEFAULT_LISTEN_ADDR)
    capture = os.environ.get(ENV_CAPTURE_PATH)
    server = serve(addr, capture)
    print(f"ShardRuntime server listening on {addr}", flush=True)
    server.wait_for_termination()


if __name__ == "__main__":
    main()
