"""DGR-033 integration tests for the standalone native C++ Shard worker.

These tests spawn the *real* compiled ``shard_worker`` executable as a separate
OS process, connect to its real localhost socket with the committed generated
``ShardRuntimeStub`` stubs, and drive the complete lifecycle/stream contract.
There is no in-memory channel, no Python servicer, and no fake transport: the
server under test is the C++ binary DGR-033 builds.

The worker binary is located via ``MESHNET_SHARD_WORKER_BIN`` or the default
out-of-tree build path ``build/native/shard_worker``. When it has not been
built (a default developer/CI checkout without the pinned gRPC C++ toolchain),
every test here is skipped rather than failed — the same ``requires_cmake``
gating pattern DGR-029/DGR-030 use for native-build-dependent tests. The
session that implemented DGR-033 built the binary and ran these for real; see
``evidence/DGR-033/README.md`` for the exact commands and results.
"""

from __future__ import annotations

import os
import signal
import socket
import subprocess
import time
import zlib

import grpc
import pytest

REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

_PYTHONPATH = os.pathsep.join(
    [os.path.join(REPO_ROOT, "packages", "node"), os.path.join(REPO_ROOT, "packages", "tracker")]
)

from meshnet_node.native_protocol.generated import (  # noqa: E402
    shard_runtime_pb2 as pb,
    shard_runtime_pb2_grpc as pb_grpc,
)


def _worker_binary() -> str | None:
    explicit = os.environ.get("MESHNET_SHARD_WORKER_BIN")
    if explicit and os.path.exists(explicit):
        return explicit
    default = os.path.join(REPO_ROOT, "build", "native", "shard_worker")
    if os.path.exists(default):
        return default
    return None


_WORKER_BIN = _worker_binary()
pytestmark = pytest.mark.skipif(
    _WORKER_BIN is None,
    reason=(
        "native shard_worker binary not built; build packages/node/native with the "
        "pinned gRPC C++ toolchain or set MESHNET_SHARD_WORKER_BIN"
    ),
)


def _free_port() -> int:
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.bind(("127.0.0.1", 0))
    port = s.getsockname()[1]
    s.close()
    return port


def _start_worker(listen_addr: str, extra_env: dict[str, str] | None = None) -> subprocess.Popen:
    env = dict(os.environ)
    env["PYTHONPATH"] = _PYTHONPATH
    if extra_env:
        env.update(extra_env)
    proc = subprocess.Popen(
        [_WORKER_BIN, listen_addr],
        cwd=REPO_ROOT,
        env=env,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
    )
    deadline = time.time() + 30.0
    while time.time() < deadline:
        line = proc.stdout.readline()
        if not line:
            if proc.poll() is not None:
                out, _ = proc.communicate()
                raise RuntimeError(f"worker exited early:\n{out}")
            continue
        if "listening on" in line:
            return proc
    raise RuntimeError("worker did not start listening in time")


class _Worker:
    """A spawned worker plus a ready channel; also captures stdout on close."""

    def __init__(self, extra_env: dict[str, str] | None = None) -> None:
        self.port = _free_port()
        self.addr = f"127.0.0.1:{self.port}"
        self.proc = _start_worker(self.addr, extra_env)
        self.channel = grpc.insecure_channel(self.addr)
        grpc.channel_ready_future(self.channel).result(timeout=15.0)

    def stub(self) -> pb_grpc.ShardRuntimeStub:
        return pb_grpc.ShardRuntimeStub(self.channel)

    def session(self, requests):
        call = self.channel.stream_stream(
            "/meshnet.shard.v1.ShardRuntime/Session",
            request_serializer=lambda m: m.SerializeToString(),
            response_deserializer=pb.SessionResponse.FromString,
        )
        return list(call(iter(requests)))

    def close(self, *, sig: int = signal.SIGTERM) -> str:
        self.channel.close()
        self.proc.send_signal(sig)
        try:
            out, _ = self.proc.communicate(timeout=10)
        except subprocess.TimeoutExpired:
            self.proc.kill()
            out, _ = self.proc.communicate()
        return out or ""


@pytest.fixture()
def worker():
    w = _Worker()
    try:
        yield w
    finally:
        if w.proc.poll() is None:
            w.close()


def _crc32c(payload: bytes) -> bytes:
    return zlib.crc32(payload).to_bytes(4, "big")


def _open(*, route_session_id="rs-1", route_epoch=7, credits_granted=16) -> pb.SessionRequest:
    return pb.SessionRequest(
        open=pb.SessionOpen(
            schema_version=pb.SCHEMA_VERSION_1,
            route_session_id=route_session_id,
            route_epoch=route_epoch,
            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),
            proposed_flow_control=pb.FlowControl(
                credits_granted=credits_granted,
                max_inflight_chunks=16,
                max_chunk_bytes=4 * 1024 * 1024,
                max_prefill_chunk_tokens=512,
            ),
            accepted_compression=[pb.COMPRESSION_NONE],
        )
    )


def _chunk(
    work_id,
    payload: bytes,
    step,
    *,
    route_session_id="rs-1",
    route_epoch=7,
    deadline_unix_nanos=0,
    fragments=1,
    total_bytes=None,
) -> pb.SessionRequest:
    total = len(payload) if total_bytes is None else total_bytes
    frags = []
    if fragments == 1:
        frags = [pb.TensorFragment(fragment_index=0, fragment_count=1, byte_offset=0, payload=payload)]
    else:
        # Split into ``fragments`` tiling pieces.
        size = max(1, len(payload) // fragments)
        offset = 0
        idx = 0
        while offset < len(payload):
            piece = payload[offset : offset + size] if idx < fragments - 1 else payload[offset:]
            frags.append(
                pb.TensorFragment(
                    fragment_index=idx, fragment_count=fragments, byte_offset=offset, payload=piece
                )
            )
            offset += len(piece)
            idx += 1
    tensor = pb.NamedTensor(
        name="hidden_states",
        shape=[1, 1, 4096],
        dtype=pb.DTYPE_BFLOAT16,
        byte_order=pb.BYTE_ORDER_LITTLE_ENDIAN,
        total_bytes=total,
        compression=pb.COMPRESSION_NONE,
        checksum=pb.Checksum(algorithm=pb.CHECKSUM_ALGORITHM_CRC32C, value=_crc32c(payload)),
        fragments=frags,
    )
    bundle = pb.TensorBundle(
        bundle_version=1,
        tensors=[tensor],
        architecture=pb.ARCHITECTURE_TYPE_DENSE,
        boundary_point="pre_tail_residual",
    )
    envelope = pb.Envelope(
        schema_version=pb.SCHEMA_VERSION_1,
        work_id=work_id,
        route_session_id=route_session_id,
        route_epoch=route_epoch,
        idempotency_step=step,
        phase=pb.PHASE_PREFILL,
        position=pb.PositionSpan(first_position=0, token_count=1),
        deadline_unix_nanos=deadline_unix_nanos,
    )
    return pb.SessionRequest(chunk=pb.ActivationChunk(envelope=envelope, bundle=bundle))


def _decode(work_id, payload: bytes, step, position) -> pb.SessionRequest:
    tensor = pb.NamedTensor(
        name="hidden_states",
        shape=[1, 1, 4096],
        dtype=pb.DTYPE_BFLOAT16,
        byte_order=pb.BYTE_ORDER_LITTLE_ENDIAN,
        total_bytes=len(payload),
        compression=pb.COMPRESSION_NONE,
        checksum=pb.Checksum(algorithm=pb.CHECKSUM_ALGORITHM_CRC32C, value=_crc32c(payload)),
        fragments=[pb.TensorFragment(fragment_index=0, fragment_count=1, byte_offset=0, payload=payload)],
    )
    return pb.SessionRequest(
        decode=pb.DecodeStep(
            idempotency_step=step,
            position=position,
            expected_past_len=position,
            work_id=work_id,
            bundle=pb.TensorBundle(bundle_version=1, tensors=[tensor], architecture=pb.ARCHITECTURE_TYPE_DENSE),
        )
    )


def _release() -> pb.SessionRequest:
    return pb.SessionRequest(
        release=pb.ReleaseSignal(route_session_id="rs-1", route_epoch=7, work_id="work-final")
    )


def _cancel(*, route_session_id="rs-1", work_id="", reason="test cancel") -> pb.SessionRequest:
    return pb.SessionRequest(
        cancel=pb.CancelSignal(route_session_id=route_session_id, route_epoch=7, work_id=work_id, reason=reason)
    )


# --- startup / health / capability ----------------------------------------


def test_worker_startup_and_health(worker):
    health = worker.stub().Health(pb.HealthRequest(schema_version=pb.SCHEMA_VERSION_1))
    assert health.state == pb.SERVING_STATE_SERVING
    assert health.schema_version == pb.SCHEMA_VERSION_1


def test_worker_capability(worker):
    cap = worker.stub().GetCapability(pb.CapabilityRequest(schema_version=pb.SCHEMA_VERSION_1))
    assert cap.validated is True
    assert cap.schema_version == pb.SCHEMA_VERSION_1
    assert cap.shard_range.end_layer == 32
    assert pb.SCHEMA_VERSION_1 in cap.supported_schema_versions


# --- fragmented prefill / decode / release ---------------------------------


def test_fragmented_prefill_echoes_reassembled_payload(worker):
    payload = b"REAL_ACTIVATION_BYTES_prefill_across_three_fragments_1234567890"
    responses = worker.session([_open(), _chunk("w1", payload, step=1, fragments=3), _release()])
    assert responses[0].WhichOneof("kind") == "accepted"
    echoed = responses[1]
    assert echoed.WhichOneof("kind") == "chunk"
    got = b"".join(f.payload for f in echoed.chunk.bundle.tensors[0].fragments)
    assert got == payload
    assert echoed.chunk.bundle.tensors[0].checksum.value == _crc32c(payload)
    assert responses[2].status.terminal is True


def test_decode_step_is_served(worker):
    payload = b"REAL_ACTIVATION_BYTES_decode_step"
    responses = worker.session([_open(), _decode("w2", payload, step=1, position=1)])
    echoed = responses[1]
    assert echoed.WhichOneof("kind") == "chunk"
    assert echoed.chunk.envelope.phase == pb.PHASE_DECODE
    assert echoed.chunk.bundle.tensors[0].fragments[0].payload == payload


def test_release_is_terminal(worker):
    responses = worker.session([_open(), _release()])
    assert responses[0].WhichOneof("kind") == "accepted"
    assert responses[1].status.terminal is True


# --- deadlines / flow control / bounded messages ---------------------------


def test_expired_deadline_is_rejected(worker):
    responses = worker.session([_open(), _chunk("w-late", b"payload", step=1, deadline_unix_nanos=1)])
    assert responses[1].status.error.code == pb.ERROR_CODE_DEADLINE_EXCEEDED


def test_flow_control_violation_and_topup(worker):
    responses = worker.session(
        [
            _open(credits_granted=1),
            _chunk("w-a", b"payload-a", step=1),
            _chunk("w-b", b"payload-b", step=2),
            pb.SessionRequest(flow_control=pb.FlowControl(credits_granted=5)),
            _chunk("w-c", b"payload-c", step=3),
        ]
    )
    assert responses[1].WhichOneof("kind") == "chunk"
    assert responses[2].status.error.code == pb.ERROR_CODE_FLOW_CONTROL_VIOLATION
    assert responses[2].status.error.retryable is True
    assert responses[3].WhichOneof("kind") == "flow_control"
    assert responses[3].flow_control.credits_granted >= 5
    assert responses[4].WhichOneof("kind") == "chunk"


def test_bounded_message_is_rejected():
    """A tensor whose declared payload exceeds the negotiated ceiling is refused."""
    w = _Worker(extra_env={"MESHNET_MAX_CHUNK_BYTES": "64"})
    try:
        big = b"x" * 128
        responses = w.session([_open(), _chunk("w-big", big, step=1, total_bytes=128)])
        status = responses[1].status
        assert status.error.code == pb.ERROR_CODE_RESOURCE_EXHAUSTED
        assert "max_chunk_bytes" in status.error.detail
    finally:
        if w.proc.poll() is None:
            w.close()


def test_malformed_fragment_tiling_is_rejected(worker):
    # A fragment at a non-zero offset with no predecessor cannot tile.
    tensor = pb.NamedTensor(
        name="hidden_states",
        shape=[1, 1, 4096],
        dtype=pb.DTYPE_BFLOAT16,
        byte_order=pb.BYTE_ORDER_LITTLE_ENDIAN,
        total_bytes=7,
        compression=pb.COMPRESSION_NONE,
        checksum=pb.Checksum(algorithm=pb.CHECKSUM_ALGORITHM_CRC32C, value=_crc32c(b"payload")),
        fragments=[pb.TensorFragment(fragment_index=0, fragment_count=1, byte_offset=5, payload=b"payload")],
    )
    bad = pb.SessionRequest(
        chunk=pb.ActivationChunk(
            envelope=pb.Envelope(
                schema_version=pb.SCHEMA_VERSION_1,
                work_id="w-gap",
                route_session_id="rs-1",
                route_epoch=7,
                idempotency_step=1,
            ),
            bundle=pb.TensorBundle(bundle_version=1, tensors=[tensor]),
        )
    )
    responses = worker.session([_open(), bad])
    assert responses[1].status.error.code == pb.ERROR_CODE_PAYLOAD_CORRUPT
    assert "tile" in responses[1].status.error.detail


def test_stale_route_epoch_is_rejected(worker):
    responses = worker.session([_open(route_epoch=7), _chunk("w-stale", b"payload", step=1, route_epoch=5)])
    assert responses[1].status.error.code == pb.ERROR_CODE_EPOCH_STALE


def test_duplicate_idempotency_step_is_acked(worker):
    chunk = _chunk("w-dup", b"payload", step=1)
    responses = worker.session([_open(), chunk, chunk])
    assert responses[1].WhichOneof("kind") == "chunk"
    assert responses[2].WhichOneof("kind") == "ack"
    assert responses[2].ack.duplicate is True


# --- cancellation ----------------------------------------------------------


def test_in_band_cancel_of_single_work_item_does_not_end_stream(worker):
    responses = worker.session(
        [
            _open(),
            _cancel(work_id="work-x"),
            _chunk("work-x", b"payload", step=1),
            _chunk("work-y", b"payload", step=2),
            _release(),
        ]
    )
    assert responses[1].status.error.code == pb.ERROR_CODE_CANCELLED
    assert responses[1].status.terminal is False
    assert responses[2].status.error.code == pb.ERROR_CODE_CANCELLED
    assert responses[3].WhichOneof("kind") == "chunk"
    assert responses[4].status.terminal is True


def test_in_band_cancel_of_whole_session_is_terminal(worker):
    responses = worker.session([_open(), _cancel(work_id="")])
    assert responses[1].status.error.code == pb.ERROR_CODE_CANCELLED
    assert responses[1].status.terminal is True


def test_out_of_band_cancel_rpc_races_ahead_of_open(worker):
    stub = worker.stub()
    resp = stub.Cancel(
        pb.CancelRequest(
            schema_version=pb.SCHEMA_VERSION_1,
            route_session_id="rs-precancel",
            route_epoch=1,
            work_id="work-precancelled",
            reason="operator abort",
        )
    )
    assert resp.cancelled_work_items == 1
    responses = worker.session(
        [
            _open(route_session_id="rs-precancel"),
            _chunk("work-precancelled", b"payload", step=1, route_session_id="rs-precancel"),
        ]
    )
    assert responses[1].status.error.code == pb.ERROR_CODE_CANCELLED


def test_release_rpc_is_idempotent(worker):
    stub = worker.stub()
    # Open a session so state exists, then release it out of band twice.
    worker.session([_open(route_session_id="rs-rel"), _release()])
    # (release signal in-stream does not erase state; the unary Release RPC does)
    first = stub.Release(pb.ReleaseRequest(schema_version=pb.SCHEMA_VERSION_1, route_session_id="rs-rel", route_epoch=7))
    second = stub.Release(pb.ReleaseRequest(schema_version=pb.SCHEMA_VERSION_1, route_session_id="rs-rel", route_epoch=7))
    assert first.released is True
    assert second.released is False  # idempotent: nothing left to drop


def test_independent_session_cancellation(worker):
    # Cancel the whole of session A; session B must remain fully serviceable.
    a = worker.session([_open(route_session_id="sess-A"), _cancel(route_session_id="sess-A", work_id="")])
    assert a[1].status.terminal is True
    b = worker.session(
        [
            _open(route_session_id="sess-B"),
            _chunk("work-b", b"payload-b", step=1, route_session_id="sess-B"),
            _release(),
        ]
    )
    assert b[1].WhichOneof("kind") == "chunk", "cancelling session A must not affect session B"


# --- graceful shutdown -----------------------------------------------------


def test_graceful_shutdown_on_sigterm():
    w = _Worker()
    # Confirm it is serving, then send SIGTERM and require a clean drain/exit.
    assert w.stub().Health(pb.HealthRequest(schema_version=pb.SCHEMA_VERSION_1)).state == pb.SERVING_STATE_SERVING
    out = w.close(sig=signal.SIGTERM)
    assert w.proc.returncode == 0, f"worker did not exit cleanly on SIGTERM:\n{out}"
    assert "shut down cleanly" in out


# --- direct vs opaque relay byte identity ----------------------------------


def test_direct_and_opaque_relay_yield_identical_responses(worker):
    """A direct hop and an opaque relay of the exact captured request bytes must
    produce byte-identical server responses (relays carry frames verbatim)."""
    payload = b"RELAY_ACTIVATION_BYTES"
    requests = [_open(), _chunk("w1", payload, step=1), _release()]

    direct_call = worker.channel.stream_stream(
        "/meshnet.shard.v1.ShardRuntime/Session",
        request_serializer=lambda m: m.SerializeToString(),
        response_deserializer=lambda b: b,
    )
    direct_resp = list(direct_call(iter(requests)))
    captured = [m.SerializeToString() for m in requests]

    relay_call = worker.channel.stream_stream(
        "/meshnet.shard.v1.ShardRuntime/Session",
        request_serializer=lambda b: b,  # raw captured bytes, no reinterpretation
        response_deserializer=lambda b: b,
    )
    relay_resp = list(relay_call(iter(captured)))

    assert len(direct_resp) == len(relay_resp) == 3
    for i, (d, r) in enumerate(zip(direct_resp, relay_resp)):
        assert d == r, f"response #{i} differs between direct and opaque relay"
