"""Identity interceptor for gRPC calls.

Populate identity context (request ID, user ID, workspace ID) for RPC calls
by extracting from metadata and setting context variables.
"""

from __future__ import annotations

from collections.abc import Awaitable, Callable
from typing import TypeVar

import grpc
from grpc import aio

from noteflow.infrastructure.logging import (
    generate_request_id,
    get_logger,
    request_id_var,
    user_id_var,
    workspace_id_var,
)

logger = get_logger(__name__)

# Metadata keys for identity context
METADATA_REQUEST_ID = "x-request-id"
METADATA_USER_ID = "x-user-id"
METADATA_WORKSPACE_ID = "x-workspace-id"

_TRequest = TypeVar("_TRequest")
_TResponse = TypeVar("_TResponse")


def _coerce_metadata_value(value: str | bytes) -> str:
    """Normalize metadata values to string."""
    return value.decode() if isinstance(value, bytes) else value


class IdentityInterceptor(aio.ServerInterceptor):
    """Interceptor that populates identity context for RPC calls.

    Extract user and workspace identifiers from gRPC metadata and
    set them as context variables for use throughout the request.

    Metadata keys:
    - x-request-id: Correlation ID for request tracing
    - x-user-id: User identifier
    - x-workspace-id: Workspace identifier for tenant scoping
    """

    async def intercept_service(
        self,
        continuation: Callable[
            [grpc.HandlerCallDetails],
            Awaitable[grpc.RpcMethodHandler[_TRequest, _TResponse]],
        ],
        handler_call_details: grpc.HandlerCallDetails,
    ) -> grpc.RpcMethodHandler[_TRequest, _TResponse]:
        """Intercept incoming RPC calls to set identity context.

        Args:
            continuation: The next interceptor or handler.
            handler_call_details: Details about the RPC call.

        Returns:
            The RPC handler for this call.
        """
        # Generate or extract request ID
        metadata = dict(handler_call_details.invocation_metadata or [])

        request_id_value = metadata.get(METADATA_REQUEST_ID)
        request_id = (
            _coerce_metadata_value(request_id_value)
            if request_id_value is not None
            else generate_request_id()
        )
        request_id_var.set(request_id)

        # Extract user and workspace IDs from metadata
        if user_id_value := metadata.get(METADATA_USER_ID):
            user_id_var.set(_coerce_metadata_value(user_id_value))

        if workspace_id_value := metadata.get(METADATA_WORKSPACE_ID):
            workspace_id_var.set(_coerce_metadata_value(workspace_id_value))

        logger.debug(
            "Identity context: request=%s user=%s workspace=%s method=%s",
            request_id,
            user_id_var.get(),
            workspace_id_var.get(),
            handler_call_details.method,
        )

        return await continuation(handler_call_details)
