"""gRPC patching for AutoDebug capture.

This module captures gRPC unary-unary calls. Streaming RPCs are not
yet supported and will be flagged as unsupported.
"""

from __future__ import annotations

import hashlib
import os
from typing import Any, Callable, Dict, Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from autodebug_probe.recorder import Recorder

# Original functions
_original_unary_unary: Optional[Callable] = None
_original_unary_stream: Optional[Callable] = None
_original_stream_unary: Optional[Callable] = None
_original_stream_stream: Optional[Callable] = None

# Global recorder reference
_recorder: Optional["Recorder"] = None

# Supported RPC types
SUPPORTED_RPC_TYPES = {"unary_unary"}
UNSUPPORTED_RPC_TYPES = {"unary_stream", "stream_unary", "stream_stream"}


def _get_recorder() -> "Recorder":
    """Get the global recorder instance."""
    if _recorder is None:
        from autodebug_probe.recorder import get_recorder
        return get_recorder()
    return _recorder


def _serialize_message(msg: Any) -> Dict[str, Any]:
    """Serialize a protobuf message for recording."""
    result = {
        "type": type(msg).__name__,
    }

    # Try to get serialized bytes
    if hasattr(msg, "SerializeToString"):
        try:
            msg_bytes = msg.SerializeToString()
            result["bytes_hex"] = msg_bytes.hex()
            result["bytes_hash"] = hashlib.sha256(msg_bytes).hexdigest()[:16]
            result["size"] = len(msg_bytes)
        except Exception:
            result["serialization_error"] = True

    # Try to get message as dict (for debugging, opt-in due to ~1-5ms overhead)
    if os.environ.get("AUTODEBUG_GRPC_VERBOSE") and hasattr(msg, "DESCRIPTOR"):
        try:
            from google.protobuf.json_format import MessageToDict
            result["fields"] = MessageToDict(msg)
        except Exception:
            pass

    return result


def _make_wrapped_call(original_call: Callable, method: str):
    """Create a wrapped call function that captures the RPC."""

    def wrapped_call(request, *args, **kwargs):
        recorder = _get_recorder()

        # Serialize request
        request_info = _serialize_message(request)

        # Make the actual call
        response = original_call(request, *args, **kwargs)

        # Serialize response
        response_info = _serialize_message(response)

        # Record the RPC
        recorder.record_grpc({
            "method": method,
            "rpc_type": "unary_unary",
            "request": request_info,
            "response": response_info,
            "method_hash": hashlib.sha256(method.encode()).hexdigest()[:16],
        })

        return response

    return wrapped_call


def _patched_unary_unary(self, method, request_serializer=None, response_deserializer=None, **kwargs):
    """Patched Channel.unary_unary that captures calls."""
    # Get the original callable (pass through any additional kwargs like _registered_method)
    call = _original_unary_unary(self, method, request_serializer, response_deserializer, **kwargs)

    # Wrap the __call__ method
    original_call = call.__call__

    class WrappedCall:
        """Wrapper that intercepts the actual RPC call."""

        def __init__(self, original, method_name):
            self._original = original
            self._method = method_name
            # Copy attributes from original
            for attr in dir(original):
                if not attr.startswith('_') and attr != '__call__':
                    try:
                        setattr(self, attr, getattr(original, attr))
                    except AttributeError:
                        pass

        def __call__(self, request, *args, **kwargs):
            return _make_wrapped_call(self._original, self._method)(request, *args, **kwargs)

        def with_call(self, request, *args, **kwargs):
            """Call with metadata."""
            recorder = _get_recorder()
            request_info = _serialize_message(request)

            response, call_info = self._original.with_call(request, *args, **kwargs)

            response_info = _serialize_message(response)

            recorder.record_grpc({
                "method": self._method,
                "rpc_type": "unary_unary",
                "request": request_info,
                "response": response_info,
                "with_call": True,
            })

            return response, call_info

        def future(self, request, *args, **kwargs):
            """Async future call."""
            recorder = _get_recorder()
            request_info = _serialize_message(request)

            future = self._original.future(request, *args, **kwargs)

            # Record that we made the call (response captured when future resolves)
            recorder.record_event({
                "type": "GRPC_FUTURE",
                "method": self._method,
                "request": request_info,
                "warning": "gRPC futures may not replay deterministically",
            })

            return future

    return WrappedCall(call, method)


def _patched_streaming_rpc(rpc_type: str):
    """Create a patched streaming RPC function that records unsupported warning."""

    def patched(self, method, *args, **kwargs):
        recorder = _get_recorder()
        recorder.record_event({
            "type": "GRPC_UNSUPPORTED",
            "method": method,
            "rpc_type": rpc_type,
            "message": f"gRPC {rpc_type} streaming not yet supported for replay",
        })

        # Call original
        if rpc_type == "unary_stream":
            return _original_unary_stream(self, method, *args, **kwargs)
        elif rpc_type == "stream_unary":
            return _original_stream_unary(self, method, *args, **kwargs)
        elif rpc_type == "stream_stream":
            return _original_stream_stream(self, method, *args, **kwargs)

    return patched


def install() -> None:
    """Install gRPC capture patches."""
    global _original_unary_unary, _original_unary_stream
    global _original_stream_unary, _original_stream_stream

    try:
        import grpc
        from grpc import _channel as grpc_channel
    except ImportError:
        return

    if _original_unary_unary is not None:
        return  # Already patched

    # Patch the concrete Channel class (grpc._channel.Channel), not the abstract base
    # The concrete class overrides the abstract methods
    _original_unary_unary = grpc_channel.Channel.unary_unary
    grpc_channel.Channel.unary_unary = _patched_unary_unary

    # Patch streaming RPCs (unsupported, but record warning)
    _original_unary_stream = grpc_channel.Channel.unary_stream
    _original_stream_unary = grpc_channel.Channel.stream_unary
    _original_stream_stream = grpc_channel.Channel.stream_stream

    grpc_channel.Channel.unary_stream = _patched_streaming_rpc("unary_stream")
    grpc_channel.Channel.stream_unary = _patched_streaming_rpc("stream_unary")
    grpc_channel.Channel.stream_stream = _patched_streaming_rpc("stream_stream")


def uninstall() -> None:
    """Restore original gRPC functions."""
    global _original_unary_unary, _original_unary_stream
    global _original_stream_unary, _original_stream_stream

    try:
        import grpc
        from grpc import _channel as grpc_channel
    except ImportError:
        return

    if _original_unary_unary is not None:
        grpc_channel.Channel.unary_unary = _original_unary_unary
        _original_unary_unary = None

    if _original_unary_stream is not None:
        grpc_channel.Channel.unary_stream = _original_unary_stream
        _original_unary_stream = None

    if _original_stream_unary is not None:
        grpc_channel.Channel.stream_unary = _original_stream_unary
        _original_stream_unary = None

    if _original_stream_stream is not None:
        grpc_channel.Channel.stream_stream = _original_stream_stream
        _original_stream_stream = None


def get_grpc_info() -> Dict[str, Any]:
    """Get gRPC detection info for manifest."""
    try:
        import grpc
        return {
            "detected": True,
            "version": grpc.__version__,
            "supported_rpc_types": list(SUPPORTED_RPC_TYPES),
            "unsupported_rpc_types": list(UNSUPPORTED_RPC_TYPES),
        }
    except ImportError:
        return {"detected": False}
