"""Judge-side gRPC/HTTP bridge runtime implementation."""

from __future__ import annotations

from concurrent import futures
from importlib import import_module
import json
import os
import queue
import sys
import threading
from typing import Any, Callable, Optional

import grpc


def _load_bridge_modules() -> tuple[Any, Any]:
    try:
        return (
            import_module("eval_bridge_pb2"),
            import_module("eval_bridge_pb2_grpc"),
        )
    except ImportError:
        return (
            import_module("agent_genesis.proto.eval_bridge_pb2"),
            import_module("agent_genesis.proto.eval_bridge_pb2_grpc"),
        )


eval_bridge_pb2, eval_bridge_pb2_grpc = _load_bridge_modules()

# HTTP server support (for cloud sandbox compatibility)
try:
    from http.server import HTTPServer, BaseHTTPRequestHandler
    from urllib.parse import urlparse
    _HAS_HTTP_SERVER = True
except ImportError:
    _HAS_HTTP_SERVER = False


class JudgeRuntime:
    def __init__(self) -> None:
        self._inbound_queue: "queue.Queue[Optional[dict[str, Any]]]" = queue.Queue()
        self._outbound_queue: "queue.Queue[dict[str, Any]]" = queue.Queue()

    def send(self, msg: dict[str, Any]) -> None:
        self._outbound_queue.put(dict(msg))

    def recv(self) -> Optional[dict[str, Any]]:
        msg = self._inbound_queue.get()
        if msg is None:
            return None
        if not isinstance(msg, dict):
            return None
        return msg

    def request_next_case_index(self) -> Optional[int]:
        self.send({"type": "case_request"})
        reply = self.recv()
        if not reply:
            return None
        msg_type = str(reply.get("type", ""))
        if msg_type == "case_stop":
            return None
        if msg_type != "case_assign":
            raise RuntimeError(f"unexpected scheduler msg type: {msg_type}")
        return int(reply.get("case_index", -1))

    @staticmethod
    def with_history_events(msg: dict[str, Any], event_or_events: dict | list[dict]) -> dict[str, Any]:
        out = dict(msg)
        out["history_events"] = event_or_events
        return out

    class _BridgeServicer(eval_bridge_pb2_grpc.SandboxBridgeServicer):
        def __init__(self, runtime: "JudgeRuntime") -> None:
            self._runtime = runtime

        def CheckReady(self, request, context):
            _ = (request, context)
            return eval_bridge_pb2.ReadyStatus(
                is_ready=True,
                message="judge bridge ready",
            )

        def SendMessage(self, request, context):
            _ = context
            raw = request.json_message or ""
            if not raw:
                return eval_bridge_pb2.SendAck(ok=False, error="empty message")
            try:
                msg = json.loads(raw)
            except Exception as exc:
                return eval_bridge_pb2.SendAck(ok=False, error=f"invalid json: {exc}")
            if not isinstance(msg, dict):
                return eval_bridge_pb2.SendAck(ok=False, error="message must be json object")
            self._runtime._inbound_queue.put(msg)
            return eval_bridge_pb2.SendAck(ok=True, error="")

        def RecvMessage(self, request, context):
            _ = context
            timeout_ms = int(getattr(request, "timeout_ms", 0) or 0)
            if timeout_ms <= 0:
                msg = self._runtime._outbound_queue.get()
            else:
                try:
                    msg = self._runtime._outbound_queue.get(timeout=timeout_ms / 1000.0)
                except queue.Empty:
                    return eval_bridge_pb2.RecvResponse(
                        has_message=False,
                        json_message="",
                    )
            return eval_bridge_pb2.RecvResponse(
                has_message=True,
                json_message=json.dumps(msg, ensure_ascii=False),
            )


def serve_judge_runtime(main_fn: Callable[[JudgeRuntime], None]) -> None:
    """Start both gRPC and HTTP servers for compatibility.

    gRPC: for backward compatibility with old worker versions
    HTTP: for cloud sandbox compatibility (PPIO, E2B, etc.)
    """
    runtime = JudgeRuntime()
    grpc_port = int(os.getenv("SANDBOX_GRPC_PORT", os.getenv("JUDGE_GRPC_PORT", "50051")))
    http_port = int(os.getenv("SANDBOX_HTTP_PORT", "8080"))

    # Start gRPC server (backward compatibility)
    grpc_server = grpc.server(futures.ThreadPoolExecutor(max_workers=8))
    eval_bridge_pb2_grpc.add_SandboxBridgeServicer_to_server(
        JudgeRuntime._BridgeServicer(runtime),
        grpc_server,
    )
    grpc_server.add_insecure_port(f"[::]:{grpc_port}")
    print(f"[judge_bridge] grpc listening on {grpc_port}", file=sys.stderr, flush=True)
    grpc_server.start()

    # Start HTTP server (cloud compatibility) if available
    http_server = None
    if _HAS_HTTP_SERVER:
        _JudgeHTTPHandler.runtime = runtime
        try:
            http_server = HTTPServer(("0.0.0.0", http_port), _JudgeHTTPHandler)
            print(f"[judge_bridge] http listening on {http_port}", file=sys.stderr, flush=True)
            http_thread = threading.Thread(target=http_server.serve_forever, daemon=True)
            http_thread.start()
        except Exception as e:
            print(f"[judge_bridge] http server failed to start: {e}", file=sys.stderr, flush=True)

    def _run_main() -> None:
        try:
            main_fn(runtime)
        except Exception as exc:
            runtime.send({"type": "error", "error": str(exc)})

    threading.Thread(target=_run_main, daemon=True).start()
    grpc_server.wait_for_termination()
    if http_server:
        http_server.shutdown()


class _JudgeHTTPHandler(BaseHTTPRequestHandler):
    """HTTP handler for JudgeRuntime (cloud-compatible)."""

    runtime: "JudgeRuntime" = None  # Set before starting server

    def log_message(self, format: str, *args: Any) -> None:
        # Suppress default HTTP logging
        pass

    def _send_json(self, data: dict, status: int = 200) -> None:
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.end_headers()
        self.wfile.write(json.dumps(data).encode())

    def _read_json(self) -> dict:
        content_length = int(self.headers.get("Content-Length", 0))
        if content_length == 0:
            return {}
        body = self.rfile.read(content_length)
        return json.loads(body.decode())

    def do_POST(self) -> None:  # noqa: N802
        parsed = urlparse(self.path)
        path = parsed.path.strip("/")

        if path == "check_ready":
            self._send_json({"is_ready": True, "message": "judge bridge ready"})
            return

        if path == "send_message":
            try:
                payload = self._read_json()
                raw = payload.get("json_message", "")
                if not raw:
                    self._send_json({"ok": False, "error": "empty message"})
                    return
                msg = json.loads(raw)
                if not isinstance(msg, dict):
                    self._send_json({"ok": False, "error": "message must be json object"})
                    return
                self.runtime._inbound_queue.put(msg)
                self._send_json({"ok": True, "error": ""})
            except Exception as e:
                self._send_json({"ok": False, "error": str(e)})
            return

        if path == "recv_message":
            try:
                payload = self._read_json()
                timeout_ms = int(payload.get("timeout_ms", 0))

                if timeout_ms <= 0:
                    msg = self.runtime._outbound_queue.get()
                else:
                    try:
                        msg = self.runtime._outbound_queue.get(timeout=timeout_ms / 1000.0)
                    except queue.Empty:
                        self._send_json({"has_message": False, "json_message": ""})
                        return

                self._send_json({
                    "has_message": True,
                    "json_message": json.dumps(msg, ensure_ascii=False)
                })
            except Exception as e:
                self._send_json({"has_message": False, "json_message": "", "error": str(e)})
            return

        self._send_json({"error": "unknown endpoint"}, status=404)


def serve_judge_runtime_http(main_fn: Callable[[JudgeRuntime], None]) -> None:
    """Start HTTP server instead of gRPC (for cloud sandbox compatibility)."""
    if not _HAS_HTTP_SERVER:
        raise RuntimeError("http.server module not available")

    runtime = JudgeRuntime()
    _JudgeHTTPHandler.runtime = runtime

    # Use HTTP port (can be same as gRPC or different)
    port = int(os.getenv("SANDBOX_HTTP_PORT", os.getenv("SANDBOX_GRPC_PORT", "50051")))

    server = HTTPServer(("0.0.0.0", port), _JudgeHTTPHandler)
    print(f"[judge_bridge] http listening on {port}", file=sys.stderr, flush=True)

    # Start HTTP server in background thread
    server_thread = threading.Thread(target=server.serve_forever, daemon=True)
    server_thread.start()

    # Run main function
    try:
        main_fn(runtime)
    except Exception as exc:
        runtime.send({"type": "error", "error": str(exc)})

    server.shutdown()
