"""gRPC ServiceBridgeWorker servicer — receives Handle and DeliverMessage from the server."""

from __future__ import annotations

import asyncio
import inspect
import json
import logging
from typing import TYPE_CHECKING, Any

import grpc

from service_bridge.pb import servicebridge_pb2 as pb
from service_bridge.pb import servicebridge_pb2_grpc as pb_grpc

if TYPE_CHECKING:
    from service_bridge.servicebridge import ServiceBridge

logger = logging.getLogger(__name__)


class WorkerServicer(pb_grpc.ServiceBridgeWorkerServicer):
    """Implements the ServiceBridgeWorker gRPC service."""

    def __init__(self, client: "ServiceBridge") -> None:
        self._client = client

    # ─── Handle ───────────────────────────────────────────────────────────────

    def Handle(self, request, context):  # noqa: N802
        fn = request.fn
        handler = self._client._rpc_handlers.get(fn)
        if handler is None:
            context.abort(grpc.StatusCode.NOT_FOUND, f"no handler for '{fn}'")
            return pb.HandleResponse(success=False, error=f"no handler for '{fn}'")

        handler_opts = self._client._rpc_opts.get(fn, {})
        allowed_callers = handler_opts.get("allowed_callers", [])

        # Enforce mTLS allowedCallers.
        if allowed_callers:
            caller_cn = _extract_peer_cn(context)
            if caller_cn and caller_cn not in allowed_callers:
                context.abort(
                    grpc.StatusCode.PERMISSION_DENIED,
                    f"caller '{caller_cn}' is not in allowedCallers for '{fn}'",
                )
                return pb.HandleResponse(success=False, error="permission_denied")

        # Decode payload (binary Protobuf → dict if schema present, else JSON).
        payload_bytes = request.payload if request.payload else b"{}"
        payload = _decode_payload(fn, payload_bytes, self._client)

        # Extract run_id from gRPC metadata for stream writes.
        run_id = ""
        for key, value in context.invocation_metadata():
            if key == "servicebridge-run-id":
                run_id = value
                break

        from service_bridge.servicebridge import StreamWriter
        stream_writer = StreamWriter(run_id=run_id, client=self._client)

        trace_ctx = {"trace_id": request.trace_id, "span_id": request.span_id}

        try:
            loop = asyncio.new_event_loop()
            try:
                result = loop.run_until_complete(
                    _run_handler(handler, payload, trace_ctx, stream_writer)
                )
            finally:
                loop.close()
        except Exception as exc:
            msg = str(exc)[:2048]
            return pb.HandleResponse(success=False, error=msg)

        # Encode output (dict → binary Protobuf if schema present, else JSON).
        try:
            output = _encode_output(fn, result, self._client)
        except (TypeError, ValueError) as exc:
            return pb.HandleResponse(success=False, error=f"marshal response: {exc}")

        return pb.HandleResponse(output=output, success=True)

    # ─── DeliverMessage ────────────────────────────────────────────────────────

    def DeliverMessage(self, request, context):  # noqa: N802
        # Validate peer CN: only the ServiceBridge control plane may deliver messages.
        caller_cn = _extract_peer_cn(context)
        if caller_cn and caller_cn != "ServiceBridge Server":
            return pb.DeliveryAck(ack=False, error="unauthorized_caller")

        group_name = request.group_name
        entry = self._client._event_handlers.get(group_name)
        if entry is None:
            return pb.DeliveryAck(ack=False, error="consumer_group_handler_missing")

        payload_bytes = request.payload if request.payload else b"{}"
        try:
            payload = json.loads(payload_bytes)
        except json.JSONDecodeError:
            return pb.DeliveryAck(ack=False, reject_reason="invalid_json_payload")

        # Extract run_id for stream writes.
        run_id = ""
        for key, value in context.invocation_metadata():
            if key == "servicebridge-run-id":
                run_id = value
                break

        from service_bridge.servicebridge import EventContext, StreamWriter
        stream_writer = StreamWriter(run_id=run_id, client=self._client)

        should_retry = False
        retry_after_ms = 1000
        reject_reason = ""

        def retry_fn(delay_ms: int = 1000) -> None:
            nonlocal should_retry, retry_after_ms
            should_retry = True
            retry_after_ms = delay_ms

        def reject_fn(reason: str) -> None:
            nonlocal reject_reason
            reject_reason = reason

        ec = EventContext(
            trace_id=request.trace_id,
            span_id=request.parent_span_id,
            topic=request.topic,
            group_name=group_name,
            message_id=request.message_id,
            attempt=request.attempt,
            headers=dict(request.headers),
            stream=stream_writer,
            _retry_fn=retry_fn,
            _reject_fn=reject_fn,
        )

        handler = entry["handler"]
        handler_err = ""
        try:
            loop = asyncio.new_event_loop()
            try:
                loop.run_until_complete(_run_event_handler(handler, payload, ec))
            finally:
                loop.close()
        except Exception as exc:
            handler_err = str(exc)[:2048]
            should_retry = True

        if reject_reason:
            return pb.DeliveryAck(ack=False, reject_reason=reject_reason, error=handler_err)
        if should_retry:
            return pb.DeliveryAck(ack=False, retry_after_ms=retry_after_ms, error=handler_err)
        return pb.DeliveryAck(ack=True)


# ─── Handler dispatch helpers ─────────────────────────────────────────────────

async def _run_handler(handler: Any, payload: Any, trace_ctx: dict, stream_writer: Any) -> Any:
    """Dispatch to an RPC handler, passing RpcContext if the handler accepts it."""
    sig = inspect.signature(handler)
    params = list(sig.parameters.values())

    # Determine if handler expects a context argument (2nd positional param).
    wants_ctx = len(params) >= 2

    if wants_ctx:
        from service_bridge.servicebridge import StreamWriter as _SW
        # Build a simple RpcContext-like namespace.
        ctx = _build_rpc_context(trace_ctx, stream_writer)
        if asyncio.iscoroutinefunction(handler):
            return await handler(payload, ctx)
        return handler(payload, ctx)
    else:
        if asyncio.iscoroutinefunction(handler):
            return await handler(payload)
        return handler(payload)


async def _run_event_handler(handler: Any, payload: Any, ec: Any) -> None:
    if asyncio.iscoroutinefunction(handler):
        await handler(payload, ec)
    else:
        handler(payload, ec)


def _build_rpc_context(trace_ctx: dict, stream_writer: Any) -> "RpcContext":
    return RpcContext(
        trace_id=trace_ctx.get("trace_id", ""),
        span_id=trace_ctx.get("span_id", ""),
        stream=stream_writer,
    )


class RpcContext:
    """Passed as the second argument to extended RPC handlers.

    Provides stream write access and distributed trace identifiers.
    """

    def __init__(self, trace_id: str, span_id: str, stream: Any) -> None:
        self.trace_id = trace_id
        self.span_id = span_id
        self.stream = stream


# ─── Schema encode/decode helpers ─────────────────────────────────────────────

def _decode_payload(fn: str, raw: bytes, client: "ServiceBridge") -> Any:
    """Decode raw bytes to a Python object.  Uses binary Protobuf if schema registered."""
    schema_opts = client._schemas.get(fn)
    if schema_opts and schema_opts.input:
        try:
            md = client._get_compiled_schema(fn, "input", schema_opts.input)
            if md:
                from service_bridge.schema import decode_with_schema
                return decode_with_schema(md, raw)
        except Exception as exc:
            logger.debug("schema decode payload for %r: %s — falling back to JSON", fn, exc)
    try:
        return json.loads(raw)
    except json.JSONDecodeError as exc:
        raise ValueError(f"invalid JSON payload: {exc}") from exc


def _encode_output(fn: str, result: Any, client: "ServiceBridge") -> bytes:
    """Encode a Python object to bytes. Uses binary Protobuf if schema registered."""
    if result is None:
        return b"{}"
    schema_opts = client._schemas.get(fn)
    if schema_opts and schema_opts.output:
        try:
            md = client._get_compiled_schema(fn, "output", schema_opts.output)
            if md:
                from service_bridge.schema import encode_with_schema
                return encode_with_schema(md, result)
        except Exception as exc:
            logger.debug("schema encode output for %r: %s — falling back to JSON", fn, exc)
    return json.dumps(result).encode()


# ─── mTLS peer CN extraction ──────────────────────────────────────────────────

def _extract_peer_cn(context: grpc.ServicerContext) -> str:
    """Extract the Common Name from the gRPC peer's mTLS client certificate.

    Returns "" when mTLS is not in use, the peer has no certificate, or the
    ``cryptography`` package is not installed.
    """
    try:
        ids = context.peer_identities()
        if not ids:
            return ""
        # Try to parse the certificate using the cryptography package.
        try:
            from cryptography import x509
            from cryptography.hazmat.primitives.serialization import Encoding
            cert = x509.load_der_x509_certificate(ids[0])
            attrs = cert.subject.get_attributes_for_oid(x509.NameOID.COMMON_NAME)
            return attrs[0].value if attrs else ""
        except ImportError:
            # cryptography not installed — try raw ASN.1 parsing via ssl module.
            import ssl
            pem = ssl.DER_cert_to_PEM_cert(ids[0]) if hasattr(ssl, "DER_cert_to_PEM_cert") else ""
            if pem:
                # Minimal CN extraction: find "CN=" in the subject line.
                for line in pem.splitlines():
                    if line.startswith("subject=") or "CN=" in line:
                        for part in line.split(","):
                            part = part.strip()
                            if part.startswith("CN="):
                                return part[3:]
            return ""
    except Exception:
        return ""
