"""gRPC replay for AutoDebug.

This module provides replay of captured gRPC unary-unary calls during test replay.
"""

from __future__ import annotations

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

if TYPE_CHECKING:
    from autodebug_replay.transcripts import TranscriptIndex

# Original functions
_original_unary_unary: Optional[Callable] = None

# Global transcript index
_index: Optional["TranscriptIndex"] = None

# Response class registry for deserialization
_response_classes: Dict[str, type] = {}


def register_response_class(name: str, cls: type) -> None:
    """Register a response class for deserialization.

    Args:
        name: The class name as recorded in transcripts
        cls: The protobuf message class
    """
    _response_classes[name] = cls


def _deserialize_response(response_info: Dict[str, Any]) -> Any:
    """Deserialize a recorded response back to a protobuf message."""
    response_type = response_info.get("type")
    bytes_hex = response_info.get("bytes_hex")

    if not bytes_hex:
        raise ValueError(f"No bytes recorded for response type {response_type}")

    # Try to find the response class
    if response_type in _response_classes:
        cls = _response_classes[response_type]
        msg = cls()
        msg.ParseFromString(bytes.fromhex(bytes_hex))
        return msg

    # If no class registered, return a mock response
    return _MockResponse(response_info)


class _MockResponse:
    """Mock response when actual protobuf class is not available."""

    def __init__(self, response_info: Dict[str, Any]):
        self._info = response_info
        self._bytes = bytes.fromhex(response_info.get("bytes_hex", ""))

        # Set fields from recorded data
        fields = response_info.get("fields", {})
        for key, value in fields.items():
            setattr(self, key, value)

    def SerializeToString(self) -> bytes:
        return self._bytes

    def __repr__(self):
        return f"MockResponse({self._info.get('type')})"


class _ReplayCall:
    """Replay callable that returns recorded response."""

    def __init__(self, method: str, index: "TranscriptIndex"):
        self._method = method
        self._index = index

    def __call__(self, request, *args, **kwargs):
        from autodebug_replay.transcripts import ReplayDivergence

        # Get next gRPC record from transcript
        record = self._index.expect("GRPC")

        # Verify method matches
        if record["method"] != self._method:
            raise ReplayDivergence("GRPC_METHOD_MISMATCH", {
                "expected": record["method"],
                "actual": self._method,
            })

        # Verify RPC type
        if record.get("rpc_type") != "unary_unary":
            raise ReplayDivergence("GRPC_RPC_TYPE_MISMATCH", {
                "expected": record.get("rpc_type"),
                "actual": "unary_unary",
            })

        # Return deserialized response
        return _deserialize_response(record["response"])

    def with_call(self, request, *args, **kwargs):
        """Call with metadata - returns (response, call_info)."""
        response = self(request, *args, **kwargs)
        # Return mock call info
        return response, _MockCallInfo()

    def future(self, request, *args, **kwargs):
        """Async future call."""
        from autodebug_replay.transcripts import ReplayDivergence
        raise ReplayDivergence("GRPC_FUTURE_NOT_SUPPORTED", {
            "method": self._method,
            "message": "gRPC futures cannot be replayed deterministically",
        })


class _MockCallInfo:
    """Mock gRPC call info for replay."""

    def trailing_metadata(self):
        return []

    def code(self):
        try:
            import grpc
            return grpc.StatusCode.OK
        except ImportError:
            return 0

    def details(self):
        return ""


def _patched_unary_unary(self, method, request_serializer=None, response_deserializer=None, **kwargs):
    """Patched Channel.unary_unary that replays from transcript."""
    return _ReplayCall(method, _index)


def install(index: "TranscriptIndex") -> None:
    """Install gRPC replay patches.

    Args:
        index: TranscriptIndex containing recorded gRPC calls
    """
    global _original_unary_unary, _index

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

    if _original_unary_unary is not None:
        return  # Already patched

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


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

    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

    _index = None
