"""Native activation transport over direct gRPC or the existing relay RPC.

This is deliberately a *seam adapter*, not a new relay protocol.  Direct
peers use one generated ``ShardRuntime.Session`` bidi stream for the lifetime
of a Route Session.  A relayed peer uses the relay's existing HTTP-shaped
binary-body contract: each body is exactly a serialized ``SessionRequest`` or
``SessionResponse``.  The relay only routes those bytes and restores its own
request id; it does not deserialize a native frame.

The correlation headers are duplicated outside the opaque frame solely for
the existing tracker/relay observability and billing path.  The authoritative
work, route, epoch, deadline, and cancellation information remains in the
versioned protobuf frame and is validated before it is sent.
"""

from __future__ import annotations

from collections.abc import Callable, Iterator
from dataclasses import dataclass
from queue import Empty, Full, Queue
import threading
import time
from typing import Protocol

from .native_protocol import pb

NATIVE_RELAY_PATH = "/native/session"
NATIVE_FRAME_CONTENT_TYPE = "application/x-protobuf"


class NativeActivationSeamError(RuntimeError):
    """The activation seam cannot safely continue this Route Session."""


class NativeActivationBufferFull(NativeActivationSeamError):
    """The caller exceeded the negotiated local hand-off buffer."""


class NativeActivationDisconnected(NativeActivationSeamError):
    """A direct or relay transport disconnected with an uncertain outcome."""


class RelayRequest(Protocol):
    """The existing ``_RelayHopClient.request`` shape, kept dependency-free."""

    def __call__(
        self, path: str, body: bytes, headers: dict[str, str]
    ) -> tuple[int, dict[str, str], bytes]: ...


@dataclass(frozen=True)
class NativeFrameContext:
    """Correlation owned by Meshnet around one opaque native frame."""

    request_id: str
    node_id: str
    route_session_id: str
    route_epoch: int
    work_id: str = ""
    deadline_unix_nanos: int = 0

    def __post_init__(self) -> None:
        if not self.request_id or not self.node_id or not self.route_session_id:
            raise ValueError("request, node, and Route Session identities are required")
        if self.route_epoch < 0 or self.deadline_unix_nanos < 0:
            raise ValueError("route epoch and deadline must be non-negative")

    def headers(self) -> dict[str, str]:
        """Headers retained by the existing relay/Tracker accounting path."""
        return {
            "Content-Type": NATIVE_FRAME_CONTENT_TYPE,
            "X-Meshnet-Native-Frame": "shard-runtime/v1",
            "X-Meshnet-Request-Id": self.request_id,
            "X-Meshnet-Node-Id": self.node_id,
            "X-Meshnet-Session": self.route_session_id,
            "X-Meshnet-Route-Epoch": str(self.route_epoch),
            "X-Meshnet-Work-Id": self.work_id,
            "X-Meshnet-Deadline-Unix-Nanos": str(self.deadline_unix_nanos),
            # The relay request id is restored on reply and is intentionally
            # distinct from the caller/billing request id above.
            "X-Meshnet-Activation-Id": self.request_id,
        }


@dataclass(frozen=True)
class NativeSeamTelemetry:
    transport: str
    request_id: str
    node_id: str
    work_id: str
    request_bytes: int
    response_bytes: int
    elapsed_seconds: float


TelemetrySink = Callable[[NativeSeamTelemetry], None]


def _request_identity(request: pb.SessionRequest) -> tuple[str, int, str, int]:
    kind = request.WhichOneof("kind")
    if kind == "open":
        return request.open.route_session_id, request.open.route_epoch, "", 0
    if kind == "chunk":
        item = request.chunk.envelope
        return item.route_session_id, item.route_epoch, item.work_id, item.deadline_unix_nanos
    if kind == "decode":
        # DecodeStep relies on the already opened Route Session, while work
        # identity/deadline are carried on every decode frame.
        return "", 0, request.decode.work_id, request.decode.deadline_unix_nanos
    if kind in {"cancel", "release"}:
        item = getattr(request, kind)
        return item.route_session_id, item.route_epoch, item.work_id, 0
    if kind == "flow_control":
        return "", 0, "", 0
    raise NativeActivationSeamError("native SessionRequest has no frame kind")


def _validate_request(request: pb.SessionRequest, context: NativeFrameContext) -> None:
    if request.ByteSize() == 0:
        raise NativeActivationSeamError("empty native SessionRequest is not a versioned frame")
    route_session, epoch, work_id, deadline = _request_identity(request)
    if route_session and route_session != context.route_session_id:
        raise NativeActivationSeamError("native frame Route Session differs from seam context")
    if route_session and epoch != context.route_epoch:
        raise NativeActivationSeamError("native frame route epoch differs from seam context")
    if context.work_id and work_id and work_id != context.work_id:
        raise NativeActivationSeamError("native frame work identity differs from seam context")
    if context.deadline_unix_nanos and deadline and deadline != context.deadline_unix_nanos:
        raise NativeActivationSeamError("native frame deadline differs from seam context")


def _response_work_id(response: pb.SessionResponse) -> str:
    kind = response.WhichOneof("kind")
    if kind == "chunk":
        return response.chunk.envelope.work_id
    if kind == "ack":
        return response.ack.work_id
    if kind == "status":
        return response.status.work_id
    return ""


class NativeActivationSeam:
    """One Route-Session-to-worker seam with bounded direct buffering.

    ``direct_stub`` is the generated ``ShardRuntimeStub`` and is selected when
    it is available.  ``relay_request`` has the exact signature of the
    existing persistent relay client; no relay server or bridge API changes
    are needed.  Relay calls are intentionally not retried: a failed send may
    already have mutated downstream Hot KV state.
    """

    def __init__(
        self,
        context: NativeFrameContext,
        *,
        direct_stub=None,
        relay_request: RelayRequest | None = None,
        max_buffered_frames: int = 8,
        telemetry: TelemetrySink | None = None,
    ) -> None:
        if (direct_stub is None) == (relay_request is None):
            raise ValueError("provide exactly one of direct_stub or relay_request")
        if max_buffered_frames < 1:
            raise ValueError("max_buffered_frames must be positive")
        self.context = context
        self._direct_stub = direct_stub
        self._relay_request = relay_request
        self._telemetry = telemetry
        self._closed = False
        self._failure: BaseException | None = None
        self._responses: Queue[pb.SessionResponse | BaseException] = Queue(maxsize=max_buffered_frames)
        self._requests: Queue[pb.SessionRequest | object] | None = None
        self._thread: threading.Thread | None = None
        self._stop = object()
        if direct_stub is not None:
            self._requests = Queue(maxsize=max_buffered_frames)
            self._thread = threading.Thread(target=self._run_direct, daemon=True, name="native-activation-grpc")
            self._thread.start()

    @property
    def transport(self) -> str:
        return "direct-grpc" if self._direct_stub is not None else "relay"

    def _direct_requests(self) -> Iterator[pb.SessionRequest]:
        assert self._requests is not None
        while True:
            item = self._requests.get()
            if item is self._stop:
                return
            assert isinstance(item, pb.SessionRequest)
            yield item

    def _run_direct(self) -> None:
        try:
            assert self._direct_stub is not None
            for response in self._direct_stub.Session(self._direct_requests()):
                self._put_response(response)
        except BaseException as exc:
            self._failure = exc
            self._put_response(exc)

    def _put_response(self, value: pb.SessionResponse | BaseException) -> None:
        # A worker may finish while a caller is abandoning the session.  Do not
        # let an unconsumed response turn into an unbounded producer queue.
        try:
            self._responses.put(value, timeout=0.1)
        except Full:
            self._failure = NativeActivationBufferFull("native response buffer is full")

    def send(self, request: pb.SessionRequest) -> pb.SessionResponse | None:
        """Send one already-versioned protobuf frame without rewriting it."""
        if self._closed:
            raise NativeActivationDisconnected("native activation seam is closed")
        if self._failure is not None:
            raise NativeActivationDisconnected("native activation stream failed") from self._failure
        _validate_request(request, self.context)
        frame = request.SerializeToString()
        if self._direct_stub is not None:
            assert self._requests is not None
            try:
                self._requests.put_nowait(request)
            except Full as exc:
                raise NativeActivationBufferFull("native direct request buffer is full") from exc
            return None

        assert self._relay_request is not None
        started = time.monotonic()
        try:
            status, _, response_frame = self._relay_request(NATIVE_RELAY_PATH, frame, self.context.headers())
        except Exception as exc:
            self._closed = True
            raise NativeActivationDisconnected("relay outcome is uncertain; refusing replay") from exc
        if status != 200:
            self._closed = True
            raise NativeActivationDisconnected(f"relay native frame returned HTTP {status}")
        response = pb.SessionResponse()
        try:
            response.ParseFromString(response_frame)
        except Exception as exc:
            self._closed = True
            raise NativeActivationSeamError("relay returned a malformed native response frame") from exc
        self._validate_response(response)
        self._record(len(frame), len(response_frame), started)
        return response

    def receive(self, timeout: float | None = None) -> pb.SessionResponse:
        """Receive the next response from the one long-lived direct stream."""
        if self._direct_stub is None:
            raise NativeActivationSeamError("relay sends return their response synchronously")
        try:
            value = self._responses.get(timeout=timeout)
        except Empty as exc:
            raise TimeoutError("timed out waiting for native direct response") from exc
        if isinstance(value, BaseException):
            raise NativeActivationDisconnected("native direct stream disconnected") from value
        self._validate_response(value)
        # gRPC owns its framing, but this records the actual protobuf payload
        # size at the seam for the same telemetry shape as relay.
        self._record(0, len(value.SerializeToString()), time.monotonic())
        return value

    def cancel(self, reason: str = "cancelled") -> pb.SessionResponse | None:
        """Propagate cancellation through the same path and correlation fields."""
        return self.send(pb.SessionRequest(cancel=pb.CancelSignal(
            route_session_id=self.context.route_session_id,
            route_epoch=self.context.route_epoch,
            work_id=self.context.work_id,
            reason=reason,
        )))

    def _validate_response(self, response: pb.SessionResponse) -> None:
        work_id = _response_work_id(response)
        if self.context.work_id and work_id and work_id != self.context.work_id:
            raise NativeActivationSeamError("native response work identity differs from seam context")

    def _record(self, request_bytes: int, response_bytes: int, started: float) -> None:
        if self._telemetry is not None:
            self._telemetry(NativeSeamTelemetry(
                transport=self.transport, request_id=self.context.request_id,
                node_id=self.context.node_id, work_id=self.context.work_id,
                request_bytes=request_bytes, response_bytes=response_bytes,
                elapsed_seconds=max(0.0, time.monotonic() - started),
            ))

    def close(self) -> None:
        if self._closed:
            return
        self._closed = True
        if self._requests is not None:
            try:
                self._requests.put_nowait(self._stop)
            except Full:
                # The bounded queue is intentionally never expanded during
                # shutdown; the worker will observe process/session teardown.
                pass
        if self._thread is not None:
            self._thread.join(timeout=1.0)