from __future__ import annotations

import json
import logging
import time
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__)

# HTTP transport support (for cloud sandbox compatibility)
try:
    import requests
    _HAS_REQUESTS = True
except ImportError:
    _HAS_REQUESTS = False


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):
    _SERVICE_CANDIDATES: tuple[str, ...] = (
        "agent_genesis.evaluation.SandboxBridge",
        "evaluation.SandboxBridge",
        "SandboxBridge",
    )

    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)
        self._active_service = self._SERVICE_CANDIDATES[0]
        self._bind_rpc_methods(self._active_service)

    def _bind_rpc_methods(self, service_name: str) -> None:
        self._active_service = service_name
        if not hasattr(self._channel, "unary_unary"):
            self._check_ready_rpc = self._stub.CheckReady
            self._send_message_rpc = self._stub.SendMessage
            self._recv_message_rpc = self._stub.RecvMessage
            return
        base = f"/{service_name}"
        self._check_ready_rpc = self._channel.unary_unary(
            f"{base}/CheckReady",
            request_serializer=eval_bridge_pb2.Empty.SerializeToString,
            response_deserializer=eval_bridge_pb2.ReadyStatus.FromString,
        )
        self._send_message_rpc = self._channel.unary_unary(
            f"{base}/SendMessage",
            request_serializer=eval_bridge_pb2.ProtocolEnvelope.SerializeToString,
            response_deserializer=eval_bridge_pb2.SendAck.FromString,
        )
        self._recv_message_rpc = self._channel.unary_unary(
            f"{base}/RecvMessage",
            request_serializer=eval_bridge_pb2.RecvRequest.SerializeToString,
            response_deserializer=eval_bridge_pb2.RecvResponse.FromString,
        )

    def _try_switch_service_name(self, probe_timeout_seconds: float) -> bool:
        if not hasattr(self._channel, "unary_unary"):
            return False
        probe_timeout = max(0.2, min(2.0, float(probe_timeout_seconds)))
        not_found_codes = {
            getattr(grpc.StatusCode, "UNIMPLEMENTED", None),
            getattr(grpc.StatusCode, "NOT_FOUND", None),
            getattr(grpc.StatusCode, "UNAVAILABLE", None),
        }
        for service_name in self._SERVICE_CANDIDATES:
            if service_name == self._active_service:
                continue
            probe_rpc = self._channel.unary_unary(
                f"/{service_name}/CheckReady",
                request_serializer=eval_bridge_pb2.Empty.SerializeToString,
                response_deserializer=eval_bridge_pb2.ReadyStatus.FromString,
            )
            try:
                resp = probe_rpc(eval_bridge_pb2.Empty(), timeout=probe_timeout)
            except grpc.RpcError as err:
                if err.code() in not_found_codes:
                    continue
                logger.debug(
                    "gRPC bridge service probe failed (%s, %s): %s",
                    self.target,
                    service_name,
                    err,
                )
                continue
            self._bind_rpc_methods(service_name)
            logger.info(
                "gRPC bridge service resolved (%s): %s",
                self.target,
                service_name,
            )
            return bool(resp.is_ready)
        return False

    def wait_for_ready(self, timeout: int) -> bool:
        ready_timeout = max(1, int(timeout))
        deadline = time.monotonic() + ready_timeout
        retry_interval = 0.5
        last_error: Optional[str] = None
        not_ready_codes = {
            getattr(grpc.StatusCode, "UNIMPLEMENTED", None),
            getattr(grpc.StatusCode, "NOT_FOUND", None),
            getattr(grpc.StatusCode, "UNAVAILABLE", None),
        }
        while True:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                if last_error:
                    logger.warning(
                        "gRPC wait_for_ready timeout (%s, service=%s): %s",
                        self.target,
                        self._active_service,
                        last_error,
                    )
                return False
            try:
                grpc.channel_ready_future(self._channel).result(timeout=min(1.0, remaining))
                resp = self._check_ready_rpc(
                    eval_bridge_pb2.Empty(),
                    timeout=min(5, int(remaining) + 1),
                )
                if resp.is_ready:
                    return True
                last_error = f"bridge returned is_ready={resp.is_ready}: {resp.message}"
            except grpc.FutureTimeoutError:
                last_error = "channel not ready before probe timeout"
            except grpc.RpcError as err:
                code = err.code()
                if code in not_ready_codes:
                    code_name = getattr(code, "name", str(code))
                    details_fn = getattr(err, "details", None)
                    details = details_fn() if callable(details_fn) else str(err)
                    last_error = f"{code_name}: {details or err}"
                    if self._try_switch_service_name(probe_timeout_seconds=remaining):
                        return True
                else:
                    logger.warning(
                        "gRPC wait_for_ready rpc failed (%s): %s", self.target, err
                    )
                    return False
            sleep_time = min(retry_interval, max(0.05, deadline - time.monotonic()))
            time.sleep(sleep_time)

    def send_message(self, msg: dict[str, Any], timeout: int) -> None:
        payload = json.dumps(msg, ensure_ascii=False)
        try:
            ack = self._send_message_rpc(
                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._recv_message_rpc(
                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()


class HttpSandboxTransport(SandboxTransport):
    """HTTP REST transport for sandbox bridge (cloud-compatible).

    This transport replaces gRPC with simple HTTP POST requests,
    making it compatible with cloud sandbox providers that only
    support HTTP/HTTPS ingress (like PPIO, E2B, etc.).

    API Endpoints:
        POST /check_ready  -> {"is_ready": bool, "message": str}
        POST /send_message -> {"json_message": str} -> {"ok": bool, "error": str}
        POST /recv_message -> {"timeout_ms": int} -> {"has_message": bool, "json_message": str}
    """

    def __init__(self, host_or_target: str) -> None:
        if not _HAS_REQUESTS:
            raise RuntimeError("requests library required for HTTP transport; install with: pip install requests")

        raw = str(host_or_target or "").strip()
        if not raw:
            raise ValueError("empty HTTP host/target")

        # Handle protocol prefix
        if raw.startswith("https://"):
            self.base_url = raw.rstrip("/")
            self.verify_ssl = True
        elif raw.startswith("http://"):
            self.base_url = raw.rstrip("/")
            self.verify_ssl = False
        else:
            # Default to HTTPS for cloud sandboxes
            self.base_url = f"https://{raw}"
            self.verify_ssl = True

        # Session for connection reuse
        self._session = requests.Session()
        self._session.headers.update({
            "Content-Type": "application/json",
            "Accept": "application/json",
        })

        logger.debug("HTTP transport initialized: %s", self.base_url)

    def _url(self, endpoint: str) -> str:
        return f"{self.base_url}/{endpoint}"

    def _post(self, endpoint: str, payload: dict, timeout: float) -> dict:
        """POST JSON and return parsed response."""
        try:
            resp = self._session.post(
                self._url(endpoint),
                json=payload,
                timeout=timeout,
                verify=self.verify_ssl,
            )
            resp.raise_for_status()
            return resp.json()
        except requests.exceptions.RequestException as e:
            raise SandboxTransportConnectionError(
                f"HTTP POST failed ({self.base_url}/{endpoint}): {e}"
            ) from e

    def wait_for_ready(self, timeout: int) -> bool:
        """Poll /check_ready until ready or timeout (with exponential backoff for 502s)."""
        ready_timeout = max(1, int(timeout))
        deadline = time.monotonic() + ready_timeout
        retry_interval = 1.0  # Start with 1 second
        max_interval = 5.0    # Cap at 5 seconds
        last_error: Optional[str] = None

        # Initial delay for cloud gateway warmup (PPIO/E2B need time to configure ingress)
        logger.debug("HTTP transport: waiting 5s for cloud gateway warmup...")
        time.sleep(5)

        while True:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                if last_error:
                    logger.warning("HTTP wait_for_ready timeout (%s): %s", self.base_url, last_error)
                return False

            try:
                resp = self._post("check_ready", {}, timeout=min(5, remaining))
                if resp.get("is_ready"):
                    return True
                last_error = f"not ready: {resp.get('message', 'unknown')}"
            except requests.exceptions.HTTPError as e:
                # Handle 502 Bad Gateway (PPIO gateway not ready yet) and other HTTP errors
                status = e.response.status_code if hasattr(e, 'response') and e.response else 0
                last_error = f"HTTP {status}: {e}"
                # Exponential backoff for gateway errors (502, 503, 504)
                if status in (502, 503, 504):
                    retry_interval = min(retry_interval * 1.5, max_interval)
            except (requests.exceptions.ConnectionError, requests.exceptions.Timeout) as e:
                last_error = f"connection error: {e}"
                retry_interval = min(retry_interval * 1.2, max_interval)
            except Exception as e:
                last_error = str(e)
                retry_interval = min(retry_interval * 1.2, max_interval)

            sleep_time = min(retry_interval, max(0.1, deadline - time.monotonic()))
            time.sleep(sleep_time)

    def send_message(self, msg: dict[str, Any], timeout: int) -> None:
        """Send message via HTTP POST /send_message."""
        payload = {"json_message": json.dumps(msg, ensure_ascii=False)}
        try:
            resp = self._post("send_message", payload, timeout=max(1, min(int(timeout), 120)))
            if not resp.get("ok"):
                raise SandboxTransportError(
                    f"bridge rejected message ({self.base_url}): {resp.get('error', 'unknown')}"
                )
        except requests.exceptions.Timeout as e:
            raise SandboxTransportConnectionError(
                f"HTTP send timeout ({self.base_url}): {e}"
            ) from e

    def recv_message(self, timeout: int) -> Optional[str]:
        """Receive message via HTTP POST /recv_message (long-polling)."""
        wait_seconds = max(1, int(timeout))
        payload = {"timeout_ms": wait_seconds * 1000}
        # HTTP timeout = message timeout + 5s buffer
        http_timeout = wait_seconds + 5

        try:
            resp = self._post("recv_message", payload, timeout=http_timeout)
            if not resp.get("has_message"):
                return None
            return resp.get("json_message") or None
        except requests.exceptions.Timeout:
            # Timeout means no message available
            return None
        except Exception as e:
            raise SandboxTransportConnectionError(
                f"HTTP recv failed ({self.base_url}): {e}"
            ) from e

    def close(self) -> None:
        """Close HTTP session."""
        self._session.close()
