"""
gRPC Adapter — Expose a gRPC service as Brynq tools.

Accepts a gRPC endpoint (e.g. ``grpc://localhost:50051``) and provides
tools to call unary RPCs on the service.  Service and method metadata are
discovered via gRPC server reflection when available, or can be supplied
manually via a service descriptor.

Uses ``grpcio`` and ``grpcio-reflection`` (lazy imported) for connection
and reflection.  All stdlib except those two which are imported lazily.
"""

from __future__ import annotations

import json
from typing import Any, Dict, List

from runtime.adapters.base import BaseAdapter

_VERSION = "0.1.0"


class GRPCAdapter(BaseAdapter):
    """Adapter that exposes a gRPC service as callable tools.

    Parameters
    ----------
    endpoint:
        A gRPC endpoint URL (e.g. ``grpc://localhost:50051``).
        The ``grpc://`` scheme prefix is stripped before connecting.
    """

    def __init__(self, endpoint: str) -> None:
        # Strip scheme prefix if present
        clean = endpoint
        for prefix in ("grpc://", "grpcs://"):
            if clean.lower().startswith(prefix):
                clean = clean[len(prefix):]
                break
        self._endpoint = clean
        self._secure = endpoint.lower().startswith("grpcs://")
        self._channel: Any = None
        self._tools: List[Dict[str, Any]] | None = None

    # ── BaseAdapter interface ────────────────────────────────────────────

    def introspect(self) -> Dict[str, Any]:
        """Return discovered gRPC service methods as tools."""
        self._ensure_tools()
        return {
            "tools": list(self._tools),  # type: ignore[arg-type]
            "version": _VERSION,
            "endpoint": self._endpoint,
        }

    def execute(self, tool_name: str, args: Dict[str, Any]) -> Dict[str, Any]:
        """Execute a gRPC method by tool name.

        Sends a unary RPC call to the service with the given arguments
        serialised as JSON.

        Parameters
        ----------
        tool_name:
            Name of the tool to invoke (must be one returned by ``introspect``).
        args:
            Keyword arguments forwarded as the request message fields.
        """
        self._ensure_tools()

        tool = self._find_tool(tool_name)
        if tool is None:
            return {"status": "error", "result": f"Unknown tool: {tool_name}"}

        try:
            result = self._call_method(tool, args)
            return {"status": "ok", "result": result}
        except Exception as exc:  # noqa: BLE001
            return {"status": "error", "result": str(exc)}

    # ── Tool discovery ───────────────────────────────────────────────────

    def _ensure_tools(self) -> None:
        """Build the tool list lazily (once) via server reflection."""
        if self._tools is not None:
            return

        try:
            self._tools = self._reflect_services()
        except Exception:  # noqa: BLE001
            # Reflection not available — expose a generic call tool
            self._tools = [
                {
                    "name": "call",
                    "description": (
                        f"Call a gRPC method on {self._endpoint}. "
                        "Provide 'service', 'method', and 'payload' in args."
                    ),
                    "parameters": [
                        {"name": "service", "description": "Fully qualified service name", "required": True},
                        {"name": "method", "description": "Method name to invoke", "required": True},
                        {"name": "payload", "description": "JSON request payload", "required": False},
                    ],
                    "service": None,
                    "method": None,
                },
            ]

    def _reflect_services(self) -> List[Dict[str, Any]]:
        """Use gRPC server reflection to discover services and methods."""
        grpc = _import_grpc()
        reflection = _import_grpc_reflection()

        channel = self._get_channel(grpc)
        stub = reflection.ServerReflectionStub(channel)

        # List all services
        request = reflection.ServerReflectionRequest(list_services="")
        responses = stub.ServerReflectionInfo(iter([request]))

        tools: List[Dict[str, Any]] = []
        for response in responses:
            if not response.list_services_response:
                continue
            for svc in response.list_services_response.service:
                service_name = svc.name
                # Skip the reflection service itself
                if "ServerReflection" in service_name:
                    continue
                methods = self._get_service_methods(stub, reflection, service_name)
                for method_name, method_desc in methods:
                    tool_id = f"{service_name}/{method_name}"
                    tools.append({
                        "name": tool_id,
                        "description": f"gRPC {service_name}.{method_name}",
                        "parameters": self._extract_method_params(method_desc),
                        "service": service_name,
                        "method": method_name,
                    })

        if not tools:
            raise ValueError("No services found via reflection")

        return tools

    @staticmethod
    def _get_service_methods(
        stub: Any, reflection: Any, service_name: str
    ) -> List[tuple]:
        """Get method descriptors for a service via reflection."""
        request = reflection.ServerReflectionRequest(
            file_containing_symbol=service_name
        )
        responses = stub.ServerReflectionInfo(iter([request]))

        methods: List[tuple] = []
        for response in responses:
            if not response.file_descriptor_response:
                continue
            from google.protobuf import descriptor_pb2  # noqa: E402

            for fd_bytes in response.file_descriptor_response.file_descriptor_proto:
                fd_proto = descriptor_pb2.FileDescriptorProto()
                fd_proto.ParseFromString(fd_bytes)
                for svc in fd_proto.service:
                    if svc.name in service_name:
                        for method in svc.method:
                            methods.append((method.name, method))
        return methods

    @staticmethod
    def _extract_method_params(method_desc: Any) -> List[Dict[str, Any]]:
        """Extract parameter info from a method descriptor."""
        return [
            {
                "name": "payload",
                "description": (
                    f"JSON request payload for {getattr(method_desc, 'name', 'method')}"
                ),
                "required": False,
            },
        ]

    # ── Execution ────────────────────────────────────────────────────────

    def _call_method(self, tool: Dict[str, Any], args: Dict[str, Any]) -> Any:
        """Make a unary gRPC call."""
        grpc = _import_grpc()
        channel = self._get_channel(grpc)

        service = tool.get("service") or args.get("service")
        method = tool.get("method") or args.get("method")

        if not service or not method:
            raise ValueError("Both 'service' and 'method' are required")

        payload = args.get("payload", {})
        if isinstance(payload, str):
            payload = json.loads(payload)

        # Build the full method path
        method_path = f"/{service}/{method}"

        # Serialise payload as JSON bytes for generic invocation
        request_bytes = json.dumps(payload).encode("utf-8")

        # Use the generic unary-unary channel callable
        call = channel.unary_unary(
            method_path,
            request_serializer=_identity,
            response_deserializer=_identity,
        )

        response_bytes = call(request_bytes)

        # Try to parse as JSON, fall back to raw string
        try:
            return json.loads(response_bytes)
        except (json.JSONDecodeError, TypeError):
            if isinstance(response_bytes, bytes):
                return response_bytes.decode("utf-8", errors="replace")
            return str(response_bytes)

    # ── Channel management ───────────────────────────────────────────────

    def _get_channel(self, grpc: Any) -> Any:
        """Get or create the gRPC channel (cached)."""
        if self._channel is not None:
            return self._channel

        if self._secure:
            credentials = grpc.ssl_channel_credentials()
            self._channel = grpc.secure_channel(self._endpoint, credentials)
        else:
            self._channel = grpc.insecure_channel(self._endpoint)

        return self._channel

    def _find_tool(self, tool_name: str) -> Dict[str, Any] | None:
        """Look up a tool by name."""
        assert self._tools is not None  # noqa: S101
        for tool in self._tools:
            if tool["name"] == tool_name:
                return tool
        return None


# ── helpers ──────────────────────────────────────────────────────────────


def _identity(data: bytes) -> bytes:
    """Pass-through serialiser / deserialiser for generic gRPC calls."""
    return data


def _import_grpc():  # type: ignore[no-untyped-def]
    """Lazily import grpcio."""
    try:
        import grpc  # noqa: E402
        return grpc
    except ImportError:
        raise ImportError(
            "grpcio is required for gRPC connections. "
            "Install it with: pip install grpcio"
        ) from None


def _import_grpc_reflection():  # type: ignore[no-untyped-def]
    """Lazily import grpc_reflection."""
    try:
        from grpc_reflection.v1alpha import reflection_pb2, reflection_pb2_grpc  # noqa: E402

        class _ReflectionHelper:
            """Thin wrapper exposing the reflection stub and request types."""
            ServerReflectionStub = reflection_pb2_grpc.ServerReflectionStub
            ServerReflectionRequest = reflection_pb2.ServerReflectionRequest

        return _ReflectionHelper()
    except ImportError:
        raise ImportError(
            "grpcio-reflection is required for gRPC server reflection. "
            "Install it with: pip install grpcio-reflection"
        ) from None
