from __future__ import annotations

import json
import logging
from abc import ABC, abstractmethod
from typing import Any, Optional

import grpc

from .proto import eval_bridge_pb2, eval_bridge_pb2_grpc

logger: logging.Logger = logging.getLogger(__name__)


class SandboxTransportError(RuntimeError):
    """Transport-level failure between worker and sandbox bridge."""


class SandboxTransportConnectionError(SandboxTransportError):
    """Connection failure to sandbox bridge."""


class SandboxTransport(ABC):
    @abstractmethod
    def wait_for_ready(self, timeout: int) -> bool:
        raise NotImplementedError

    @abstractmethod
    def send_message(self, msg: dict[str, Any], timeout: int) -> None:
        raise NotImplementedError

    @abstractmethod
    def recv_message(self, timeout: int) -> Optional[str]:
        raise NotImplementedError

    @abstractmethod
    def close(self) -> None:
        raise NotImplementedError


def resolve_grpc_target(host_or_target: str) -> tuple[str, bool]:
    raw = str(host_or_target or "").strip()
    if not raw:
        raise ValueError("empty grpc host/target")

    secure = True
    if raw.startswith("https://"):
        raw = raw[len("https://") :]
        secure = True
    elif raw.startswith("http://"):
        raw = raw[len("http://") :]
        secure = False

    if "/" in raw:
        raw = raw.split("/", 1)[0]

    local_prefixes = ("127.0.0.1", "localhost")
    is_local = raw.startswith(local_prefixes)

    if ":" not in raw:
        raw = f"{raw}:{50051 if is_local else 443}"

    if is_local and secure:
        secure = False

    return raw, secure


class GrpcSandboxTransport(SandboxTransport):
    def __init__(
        self,
        host_or_target: str,
        *,
        options: Optional[list[tuple[str, int]]] = None,
    ) -> None:
        target, secure = resolve_grpc_target(host_or_target)
        self.target = target
        self.secure = secure

        channel_options = options or [
            ("grpc.max_receive_message_length", 20 * 1024 * 1024),
            ("grpc.max_send_message_length", 20 * 1024 * 1024),
            ("grpc.keepalive_time_ms", 30_000),
            ("grpc.keepalive_timeout_ms", 10_000),
            ("grpc.keepalive_permit_without_calls", 1),
        ]

        if secure:
            creds = grpc.ssl_channel_credentials()
            self._channel = grpc.secure_channel(target, creds, options=channel_options)
        else:
            self._channel = grpc.insecure_channel(target, options=channel_options)
        self._stub = eval_bridge_pb2_grpc.SandboxBridgeStub(self._channel)

    def wait_for_ready(self, timeout: int) -> bool:
        ready_timeout = max(1, int(timeout))
        try:
            grpc.channel_ready_future(self._channel).result(timeout=ready_timeout)
            resp = self._stub.CheckReady(
                eval_bridge_pb2.Empty(),
                timeout=min(5, ready_timeout),
            )
            return bool(resp.is_ready)
        except grpc.FutureTimeoutError:
            return False
        except grpc.RpcError as err:
            logger.warning("gRPC wait_for_ready rpc failed (%s): %s", self.target, err)
            return False

    def send_message(self, msg: dict[str, Any], timeout: int) -> None:
        payload = json.dumps(msg, ensure_ascii=False)
        try:
            ack = self._stub.SendMessage(
                eval_bridge_pb2.ProtocolEnvelope(json_message=payload),
                timeout=max(1, min(int(timeout), 120)),
            )
        except grpc.RpcError as err:
            raise SandboxTransportConnectionError(
                f"gRPC send failed ({self.target}): {err}"
            ) from err
        if not ack.ok:
            raise SandboxTransportError(
                f"bridge rejected message ({self.target}): {ack.error}"
            )

    def recv_message(self, timeout: int) -> Optional[str]:
        wait_seconds = max(1, int(timeout))
        try:
            resp = self._stub.RecvMessage(
                eval_bridge_pb2.RecvRequest(timeout_ms=wait_seconds * 1000),
                timeout=wait_seconds + 2,
            )
        except grpc.RpcError as err:
            if err.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
                return None
            raise SandboxTransportConnectionError(
                f"gRPC recv failed ({self.target}): {err}"
            ) from err

        if not resp.has_message:
            return None
        return resp.json_message or None

    def close(self) -> None:
        self._channel.close()
