snakepit

#!/usr/bin/env python3
"""
Multi-threaded gRPC bridge server for Snakepit.

This server variant uses a ThreadPoolExecutor to handle multiple concurrent requests
within a single Python process. Designed for Python 3.13+ free-threading mode and
CPU-intensive workloads.

Key differences from grpc_server.py:
- ThreadPoolExecutor with configurable thread count
- Thread-safe adapter requirements
- Concurrent request handling
- Thread safety validation (optional)
- Request tracking per thread

Requirements:
- Python 3.13+ (recommended for free-threading)
- Thread-safe adapter implementation
- Thread-safe ML libraries (NumPy, PyTorch, etc.)
"""

import argparse
import asyncio
import inspect
import grpc
import logging
import signal
import sys
import time
import threading
import traceback
from concurrent import futures
from typing import Dict, Optional, Sequence, Tuple
from collections import defaultdict

# Add the package to Python path
sys.path.insert(0, '.')

import snakepit_bridge_pb2 as pb2
import snakepit_bridge_pb2_grpc as pb2_grpc
from snakepit_bridge.session_context import SessionContext
from snakepit_bridge.serialization import TypeSerializer
from snakepit_bridge import telemetry
from google.protobuf.timestamp_pb2 import Timestamp
import json
import pickle

logging.basicConfig(
    format='%(asctime)s - [%(threadName)s] - %(name)s - %(levelname)s - [corr=%(correlation_id)s] %(message)s',
    level=logging.INFO
)

telemetry.setup_tracing()
_log_filter = telemetry.correlation_filter()
logger = logging.getLogger(__name__)
logger.addFilter(_log_filter)
logging.getLogger().addFilter(_log_filter)

logger.info("Loaded grpc_server_threaded.py from %s", __file__)


class ThreadSafetyMonitor:
    """
    Monitor for tracking thread safety issues during execution.

    Detects:
    - Concurrent method access without locking
    - Thread-unsafe library usage
    - Race conditions
    """

    def __init__(self, enabled: bool = False):
        self.enabled = enabled
        self.access_tracker: Dict[str, set] = defaultdict(set)
        self.lock = threading.Lock()
        self.warnings_issued: set = set()

    def record_access(self, method_name: str):
        """Record that current thread accessed a method"""
        if not self.enabled:
            return

        thread_id = threading.get_ident()
        with self.lock:
            self.access_tracker[method_name].add(thread_id)

            # Warn if multiple threads accessing
            if len(self.access_tracker[method_name]) > 1:
                warning_key = f"{method_name}_concurrent"
                if warning_key not in self.warnings_issued:
                    logger.warning(
                        f"⚠️  Concurrent access detected: {method_name} accessed by "
                        f"{len(self.access_tracker[method_name])} different threads"
                    )
                    self.warnings_issued.add(warning_key)

    def get_stats(self) -> dict:
        """Get monitoring statistics"""
        with self.lock:
            return {
                "tracked_methods": len(self.access_tracker),
                "warnings_issued": len(self.warnings_issued),
                "concurrent_accesses": sum(
                    1 for threads in self.access_tracker.values() if len(threads) > 1
                )
            }


class ThreadedBridgeServiceServicer(pb2_grpc.BridgeServiceServicer):
    """
    Multi-threaded gRPC bridge service.

    This servicer can handle multiple concurrent requests by dispatching them
    to a ThreadPoolExecutor. The adapter MUST be thread-safe.

    Thread Safety Requirements:
    - Adapter must implement thread-safe request handling
    - Shared state must use locks or thread-local storage
    - ML libraries must release GIL during computation
    """

    def __init__(
        self,
        adapter_class,
        elixir_address: str,
        max_workers: int,
        enable_safety_checks: bool = False,
        *,
        loop: Optional[asyncio.AbstractEventLoop] = None,
    ):
        logger.info(f"Initializing threaded servicer: max_workers={max_workers}, safety_checks={enable_safety_checks}")

        self.adapter_class = adapter_class
        self.elixir_address = elixir_address
        self.max_workers = max_workers
        self.loop = loop or asyncio.get_event_loop()
        self.server: Optional[grpc.aio.Server] = None

        # Thread safety monitoring
        self.safety_monitor = ThreadSafetyMonitor(enabled=enable_safety_checks)

        # Request tracking
        self.request_count = 0
        self.active_requests = 0
        self.start_time = time.time()
        self.stats_lock = threading.Lock()

        # Validate adapter thread safety
        self._validate_adapter_thread_safety()

        # Create async channel for proxying
        self.elixir_channel = grpc.aio.insecure_channel(elixir_address)
        self.elixir_stub = pb2_grpc.BridgeServiceStub(self.elixir_channel)

        # Create sync channel for SessionContext
        self.sync_elixir_channel = grpc.insecure_channel(elixir_address)
        self.sync_elixir_stub = pb2_grpc.BridgeServiceStub(self.sync_elixir_channel)

        logger.info(f"✅ Threaded servicer initialized. Ready for concurrent requests.")

    async def _proxy_to_elixir(
        self,
        method_name: str,
        request,
        *,
        metadata: Optional[Sequence[Tuple[str, str]]] = None,
    ):
        """Proxy a request to the Elixir backend with correlation metadata."""

        enriched_metadata = telemetry.outgoing_metadata(metadata)
        stub_method = getattr(self.elixir_stub, method_name)
        call = stub_method(request, metadata=enriched_metadata)

        if inspect.isawaitable(call):
            return await call

        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(None, lambda: call)

    def _validate_adapter_thread_safety(self):
        """Validate that the adapter declares thread safety"""
        if not hasattr(self.adapter_class, '__thread_safe__'):
            logger.warning(
                f"⚠️  Adapter {self.adapter_class.__name__} does not declare thread safety. "
                f"Set __thread_safe__ = True if thread-safe, or use process mode instead."
            )
        elif not self.adapter_class.__thread_safe__:
            logger.error(
                f"❌ Adapter {self.adapter_class.__name__} explicitly declares it is NOT thread-safe! "
                f"Use --mode process instead of --mode threaded."
            )
            raise ValueError(f"Adapter {self.adapter_class.__name__} is not thread-safe")

    def _record_request_start(self):
        """Thread-safe request start tracking"""
        with self.stats_lock:
            self.request_count += 1
            self.active_requests += 1
            return self.request_count

    def _record_request_end(self):
        """Thread-safe request end tracking"""
        with self.stats_lock:
            self.active_requests -= 1

    def get_stats(self) -> dict:
        """Get servicer statistics"""
        with self.stats_lock:
            uptime = time.time() - self.start_time
            return {
                "total_requests": self.request_count,
                "active_requests": self.active_requests,
                "max_workers": self.max_workers,
                "uptime_seconds": uptime,
                "requests_per_second": self.request_count / uptime if uptime > 0 else 0,
                "safety_monitor": self.safety_monitor.get_stats()
            }

    async def close(self):
        """Clean up resources"""
        if self.elixir_channel:
            await self.elixir_channel.close()
        if self.sync_elixir_channel:
            self.sync_elixir_channel.close()

    # Health & Session Management

    async def Ping(self, request, context):
        """Health check endpoint"""
        self.safety_monitor.record_access("Ping")
        metadata = context.invocation_metadata()

        with telemetry.span(
            "BridgeService/Ping",
            context_metadata=metadata,
            attributes={"snakepit.thread": threading.current_thread().name},
        ):
            response = pb2.PingResponse()
            response.message = (
                f"Pong from threaded Python [{threading.current_thread().name}]: {request.message}"
            )

            timestamp = Timestamp()
            timestamp.GetCurrentTime()
            response.server_time.CopyFrom(timestamp)

            return response

    async def InitializeSession(self, request, context):
        """Initialize session - proxy to Elixir"""
        self.safety_monitor.record_access("InitializeSession")
        metadata = context.invocation_metadata()
        logger.debug(f"[{threading.current_thread().name}] Proxying InitializeSession: {request.session_id}")

        with telemetry.span(
            "BridgeService/InitializeSession",
            context_metadata=metadata,
            attributes={"snakepit.session_id": request.session_id},
        ):
            return await self._proxy_to_elixir("InitializeSession", request, metadata=metadata)

    async def CleanupSession(self, request, context):
        """Cleanup session - proxy to Elixir"""
        self.safety_monitor.record_access("CleanupSession")
        metadata = context.invocation_metadata()
        logger.debug(f"[{threading.current_thread().name}] Proxying CleanupSession: {request.session_id}")

        with telemetry.span(
            "BridgeService/CleanupSession",
            context_metadata=metadata,
            attributes={"snakepit.session_id": request.session_id},
        ):
            return await self._proxy_to_elixir("CleanupSession", request, metadata=metadata)

    async def GetSession(self, request, context):
        """Get session - proxy to Elixir"""
        metadata = context.invocation_metadata()
        logger.debug(f"[{threading.current_thread().name}] Proxying GetSession: {request.session_id}")

        with telemetry.span(
            "BridgeService/GetSession",
            context_metadata=metadata,
            attributes={"snakepit.session_id": request.session_id},
        ):
            return await self._proxy_to_elixir("GetSession", request, metadata=metadata)

    async def Heartbeat(self, request, context):
        """Heartbeat - proxy to Elixir"""
        metadata = context.invocation_metadata()
        with telemetry.span(
            "BridgeService/Heartbeat",
            context_metadata=metadata,
            attributes={"snakepit.session_id": request.session_id},
        ):
            return await self._proxy_to_elixir("Heartbeat", request, metadata=metadata)

    # Variable Operations - All proxied to Elixir

    async def RegisterVariable(self, request, context):
        metadata = context.invocation_metadata()
        with telemetry.span(
            "BridgeService/RegisterVariable",
            context_metadata=metadata,
            attributes={"snakepit.variable": request.name},
        ):
            return await self._proxy_to_elixir("RegisterVariable", request, metadata=metadata)

    async def GetVariable(self, request, context):
        metadata = context.invocation_metadata()
        with telemetry.span(
            "BridgeService/GetVariable",
            context_metadata=metadata,
            attributes={"snakepit.variable": request.variable_identifier},
        ):
            return await self._proxy_to_elixir("GetVariable", request, metadata=metadata)

    async def SetVariable(self, request, context):
        metadata = context.invocation_metadata()
        with telemetry.span(
            "BridgeService/SetVariable",
            context_metadata=metadata,
            attributes={"snakepit.variable": request.variable_identifier},
        ):
            return await self._proxy_to_elixir("SetVariable", request, metadata=metadata)

    async def GetVariables(self, request, context):
        metadata = context.invocation_metadata()
        with telemetry.span(
            "BridgeService/GetVariables",
            context_metadata=metadata,
            attributes={"snakepit.variable_count": len(request.variable_identifiers)},
        ):
            return await self._proxy_to_elixir("GetVariables", request, metadata=metadata)

    async def SetVariables(self, request, context):
        metadata = context.invocation_metadata()
        with telemetry.span(
            "BridgeService/SetVariables",
            context_metadata=metadata,
            attributes={"snakepit.variable_count": len(request.updates)},
        ):
            return await self._proxy_to_elixir("SetVariables", request, metadata=metadata)

    async def ListVariables(self, request, context):
        metadata = context.invocation_metadata()
        with telemetry.span(
            "BridgeService/ListVariables",
            context_metadata=metadata,
            attributes={"snakepit.pattern": request.pattern},
        ):
            return await self._proxy_to_elixir("ListVariables", request, metadata=metadata)

    async def DeleteVariable(self, request, context):
        metadata = context.invocation_metadata()
        with telemetry.span(
            "BridgeService/DeleteVariable",
            context_metadata=metadata,
            attributes={"snakepit.variable": request.variable_identifier},
        ):
            return await self._proxy_to_elixir("DeleteVariable", request, metadata=metadata)

    # Tool Execution - Thread-safe concurrent execution

    async def ExecuteTool(self, request, context):
        """Execute tool with thread-safe handling"""
        request_id = self._record_request_start()
        thread_name = threading.current_thread().name
        self.safety_monitor.record_access("ExecuteTool")
        metadata = context.invocation_metadata()
        start_time = time.time()

        try:
            with telemetry.span(
                "BridgeService/ExecuteTool",
                context_metadata=metadata,
                attributes={
                    "snakepit.session_id": request.session_id,
                    "snakepit.tool": request.tool_name,
                    "snakepit.thread": thread_name,
                    "snakepit.request_id": request_id,
                },
            ):
                logger.info(
                    f"[{thread_name}] ExecuteTool #{request_id}: {request.tool_name} "
                    f"(session: {request.session_id})"
                )

                # Ensure session exists
                init_request = pb2.InitializeSessionRequest(session_id=request.session_id)
                try:
                    self.sync_elixir_stub.InitializeSession(
                        init_request,
                        metadata=telemetry.outgoing_metadata(metadata),
                    )
                except grpc.RpcError as e:
                    if e.code() != grpc.StatusCode.ALREADY_EXISTS:
                        logger.debug(f"InitializeSession: {e}")

                # Create ephemeral context and adapter
                session_context = SessionContext(self.sync_elixir_stub, request.session_id)
                adapter = self.adapter_class()
                adapter.set_session_context(session_context)

                # Register tools if needed
                if hasattr(adapter, 'register_with_session'):
                    adapter.register_with_session(request.session_id, self.sync_elixir_stub)

                # Initialize adapter
                if hasattr(adapter, 'initialize'):
                    if asyncio.iscoroutinefunction(adapter.initialize):
                        await adapter.initialize()
                    else:
                        adapter.initialize()

                # Decode parameters
                arguments = {
                    key: TypeSerializer.decode_any(any_msg)
                    for key, any_msg in request.parameters.items()
                }
                for key, binary_val in request.binary_parameters.items():
                    arguments[key] = pickle.loads(binary_val)

                # Execute tool
                if not hasattr(adapter, 'execute_tool'):
                    raise NotImplementedError("Adapter does not support execute_tool")

                if inspect.iscoroutinefunction(adapter.execute_tool):
                    result_data = await adapter.execute_tool(
                        tool_name=request.tool_name,
                        arguments=arguments,
                        context=session_context
                    )
                else:
                    result_data = adapter.execute_tool(
                        tool_name=request.tool_name,
                        arguments=arguments,
                        context=session_context
                    )

                # Encode response
                response = pb2.ExecuteToolResponse(success=True)

                # Type inference
                if isinstance(result_data, dict):
                    result_type = "map"
                elif isinstance(result_data, str):
                    result_type = "string"
                elif isinstance(result_data, (int, float)):
                    result_type = "float"
                elif isinstance(result_data, bool):
                    result_type = "boolean"
                elif isinstance(result_data, list):
                    result_type = "list"
                else:
                    result_type = "string"

                any_msg, binary_data = TypeSerializer.encode_any(result_data, result_type)
                response.result.CopyFrom(any_msg)
                if binary_data:
                    response.binary_result = binary_data

                response.execution_time_ms = int((time.time() - start_time) * 1000)

                logger.info(
                    f"[{thread_name}] ExecuteTool #{request_id} completed in "
                    f"{response.execution_time_ms}ms"
                )

                return response

        except Exception as e:
            logger.error(f"[{thread_name}] ExecuteTool #{request_id} failed: {e}", exc_info=True)
            response = pb2.ExecuteToolResponse()
            response.success = False
            response.error_message = str(e)
            return response

        finally:
            self._record_request_end()

    async def ExecuteStreamingTool(self, request, context):
        """Execute streaming tool with thread-safe handling"""
        request_id = self._record_request_start()
        thread_name = threading.current_thread().name
        self.safety_monitor.record_access("ExecuteStreamingTool")
        metadata = context.invocation_metadata()

        try:
            with telemetry.span(
                "BridgeService/ExecuteStreamingTool",
                context_metadata=metadata,
                attributes={
                    "snakepit.session_id": request.session_id,
                    "snakepit.tool": request.tool_name,
                    "snakepit.thread": thread_name,
                    "snakepit.request_id": request_id,
                },
            ):
                logger.info(
                    f"[{thread_name}] ExecuteStreamingTool #{request_id}: {request.tool_name}"
                )

                # Similar to ExecuteTool but yields chunks
                init_request = pb2.InitializeSessionRequest(session_id=request.session_id)
                try:
                    self.sync_elixir_stub.InitializeSession(
                        init_request,
                        metadata=telemetry.outgoing_metadata(metadata),
                    )
                except grpc.RpcError as e:
                    if e.code() != grpc.StatusCode.ALREADY_EXISTS:
                        logger.debug(f"InitializeSession: {e}")

                session_context = SessionContext(self.sync_elixir_stub, request.session_id)
                adapter = self.adapter_class()
                adapter.set_session_context(session_context)

                if hasattr(adapter, 'register_with_session'):
                    adapter.register_with_session(request.session_id, self.sync_elixir_stub)

                if hasattr(adapter, 'initialize'):
                    if asyncio.iscoroutinefunction(adapter.initialize):
                        await adapter.initialize()
                    else:
                        adapter.initialize()

                arguments = {
                    key: TypeSerializer.decode_any(any_msg)
                    for key, any_msg in request.parameters.items()
                }
                for key, binary_val in request.binary_parameters.items():
                    arguments[key] = pickle.loads(binary_val)

                if not hasattr(adapter, 'execute_tool'):
                    await context.abort(
                        grpc.StatusCode.UNIMPLEMENTED,
                        "Adapter does not support tool execution"
                    )
                    return

                if inspect.iscoroutinefunction(adapter.execute_tool):
                    stream_iterator = await adapter.execute_tool(
                        tool_name=request.tool_name,
                        arguments=arguments,
                        context=session_context
                    )
                else:
                    stream_iterator = adapter.execute_tool(
                        tool_name=request.tool_name,
                        arguments=arguments,
                        context=session_context
                    )

                loop = self.loop
                queue: asyncio.Queue = asyncio.Queue()
                sentinel = object()
                chunk_id_counter = 0
                marked_final = False

                def encode_payload(payload):
                    if payload is None:
                        return b""
                    if isinstance(payload, (bytes, bytearray, memoryview)):
                        return bytes(payload)
                    if isinstance(payload, str):
                        return payload.encode("utf-8")
                    return json.dumps(payload).encode("utf-8")

                def build_chunk(payload, *, is_final=False, metadata=None):
                    nonlocal chunk_id_counter, marked_final
                    chunk_id_counter += 1
                    chunk = pb2.ToolChunk(
                        chunk_id=f"{request.tool_name}-{chunk_id_counter}",
                        data=encode_payload(payload),
                        is_final=is_final,
                    )

                    if metadata and isinstance(metadata, dict):
                        chunk.metadata.update({str(k): str(v) for k, v in metadata.items()})

                    if is_final:
                        marked_final = True

                    return chunk

                async def enqueue_chunk(payload, *, is_final=False, metadata=None):
                    await queue.put(build_chunk(payload, is_final=is_final, metadata=metadata))

                def unpack_chunk(chunk_data):
                    payload = chunk_data
                    metadata = None
                    is_final = False

                    if hasattr(chunk_data, "data"):
                        payload = getattr(chunk_data, "data")

                    if hasattr(chunk_data, "is_final"):
                        is_final = bool(getattr(chunk_data, "is_final"))

                    if hasattr(chunk_data, "metadata"):
                        meta_attr = getattr(chunk_data, "metadata")
                        if isinstance(meta_attr, dict):
                            metadata = meta_attr

                    return payload, is_final, metadata

                async def drain_async(iterator):
                    async for chunk_data in iterator:
                        payload, is_final, metadata = unpack_chunk(chunk_data)
                        await enqueue_chunk(payload, is_final=is_final, metadata=metadata)

                def drain_sync(iterator):
                    for chunk_data in iterator:
                        payload, is_final, metadata = unpack_chunk(chunk_data)
                        fut = asyncio.run_coroutine_threadsafe(
                            enqueue_chunk(payload, is_final=is_final, metadata=metadata),
                            loop,
                        )
                        fut.result()

                async def produce_chunks():
                    error_raised = False
                    try:
                        if hasattr(stream_iterator, "__aiter__"):
                            await drain_async(stream_iterator)
                        elif hasattr(stream_iterator, "__iter__"):
                            await asyncio.to_thread(drain_sync, stream_iterator)
                        else:
                            payload, is_final, metadata = unpack_chunk(stream_iterator)
                            await enqueue_chunk(payload, is_final=is_final, metadata=metadata)
                    except Exception as iteration_error:
                        error_raised = True
                        await queue.put(iteration_error)
                    finally:
                        if not error_raised and not marked_final:
                            await enqueue_chunk(None, is_final=True)
                        await queue.put(sentinel)

                asyncio.create_task(produce_chunks())

                while True:
                    item = await queue.get()
                    if item is sentinel:
                        break
                    if isinstance(item, Exception):
                        logger.error(
                            "[%s] ExecuteStreamingTool #%s raised %s",
                            thread_name,
                            request_id,
                            item,
                            exc_info=True,
                        )
                        await context.abort(grpc.StatusCode.INTERNAL, str(item))
                        return
                    yield item

        except Exception as e:
            logger.error(f"[{thread_name}] ExecuteStreamingTool #{request_id} failed: {e}", exc_info=True)
            await context.abort(grpc.StatusCode.INTERNAL, str(e))

        finally:
            self._record_request_end()

    # Placeholder methods

    async def WatchVariables(self, request, context):
        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
        context.set_details('WatchVariables not implemented')
        return
        yield

    async def AddDependency(self, request, context):
        metadata = context.invocation_metadata()
        logger.debug("Proxying AddDependency (threaded)")
        with telemetry.span(
            "BridgeService/AddDependency",
            context_metadata=metadata,
        ):
            return await self._proxy_to_elixir("AddDependency", request, metadata=metadata)

    async def StartOptimization(self, request, context):
        metadata = context.invocation_metadata()
        logger.debug("Proxying StartOptimization (threaded)")
        with telemetry.span(
            "BridgeService/StartOptimization",
            context_metadata=metadata,
        ):
            return await self._proxy_to_elixir("StartOptimization", request, metadata=metadata)

    async def StopOptimization(self, request, context):
        metadata = context.invocation_metadata()
        logger.debug("Proxying StopOptimization (threaded)")
        with telemetry.span(
            "BridgeService/StopOptimization",
            context_metadata=metadata,
        ):
            return await self._proxy_to_elixir("StopOptimization", request, metadata=metadata)

    async def GetVariableHistory(self, request, context):
        metadata = context.invocation_metadata()
        logger.debug("Proxying GetVariableHistory (threaded)")
        with telemetry.span(
            "BridgeService/GetVariableHistory",
            context_metadata=metadata,
        ):
            return await self._proxy_to_elixir("GetVariableHistory", request, metadata=metadata)

    async def RollbackVariable(self, request, context):
        metadata = context.invocation_metadata()
        logger.debug("Proxying RollbackVariable (threaded)")
        with telemetry.span(
            "BridgeService/RollbackVariable",
            context_metadata=metadata,
        ):
            return await self._proxy_to_elixir("RollbackVariable", request, metadata=metadata)

    def set_server(self, server):
        """Set server reference for graceful shutdown"""
        self.server = server


async def wait_for_elixir_server(elixir_address: str, max_retries: int = 30, initial_delay: float = 0.05):
    """Wait for Elixir gRPC server to become available"""
    delay = initial_delay
    for attempt in range(1, max_retries + 1):
        try:
            channel = grpc.aio.insecure_channel(elixir_address)
            stub = pb2_grpc.BridgeServiceStub(channel)
            request = pb2.PingRequest(message="connection_test")
            await asyncio.wait_for(stub.Ping(request), timeout=1.0)
            await channel.close()
            logger.info(f"✅ Connected to Elixir server at {elixir_address} after {attempt} attempt(s)")
            return True
        except (grpc.aio.AioRpcError, asyncio.TimeoutError, Exception) as e:
            if attempt < max_retries:
                logger.debug(f"Elixir server not ready (attempt {attempt}/{max_retries}), retrying in {delay:.2f}s...")
                await asyncio.sleep(delay)
                delay = min(delay * 2, 2.0)
            else:
                logger.error(f"❌ Failed to connect to Elixir server after {max_retries} attempts")
                return False
    return False


async def serve_threaded(
    port: int,
    adapter_module: str,
    elixir_address: str,
    max_workers: int,
    enable_safety_checks: bool,
    shutdown_event: asyncio.Event
):
    """Start multi-threaded gRPC server"""
    logger.info(f"🚀 Starting threaded gRPC server: port={port}, max_workers={max_workers}")

    # Wait for Elixir server
    if not await wait_for_elixir_server(elixir_address):
        logger.error("Cannot start: Elixir server unavailable")
        sys.exit(1)

    # Import adapter
    module_parts = adapter_module.split('.')
    module_name = '.'.join(module_parts[:-1])
    class_name = module_parts[-1]

    try:
        module = __import__(module_name, fromlist=[class_name])
        adapter_class = getattr(module, class_name)
    except (ImportError, AttributeError) as e:
        logger.error(f"Failed to import adapter {adapter_module}: {e}")
        sys.exit(1)

    # Create server with ThreadPoolExecutor
    server = grpc.aio.server(
        futures.ThreadPoolExecutor(max_workers=max_workers),
        options=[
            ('grpc.max_send_message_length', 100 * 1024 * 1024),
            ('grpc.max_receive_message_length', 100 * 1024 * 1024),
            ('grpc.max_concurrent_streams', max_workers),
            ('grpc.so_reuseport', 0),
        ]
    )

    loop = asyncio.get_running_loop()

    servicer = ThreadedBridgeServiceServicer(
        adapter_class,
        elixir_address,
        max_workers,
        enable_safety_checks,
        loop=loop,
    )
    servicer.set_server(server)

    pb2_grpc.add_BridgeServiceServicer_to_server(servicer, server)

    # Bind to port
    try:
        actual_port = server.add_insecure_port(f'[::]:{port}')
        if actual_port == 0 and port != 0:
            logger.error(f"Failed to bind to port {port}")
            sys.exit(1)
    except Exception as e:
        logger.error(f"Exception binding to port {port}: {e}")
        sys.exit(1)

    await server.start()

    # Signal ready
    logger.info(f"GRPC_READY:{actual_port}")
    logger.info(f"✅ Threaded gRPC server ready on port {actual_port} with {max_workers} worker threads")

    # Wait for shutdown
    server_task = asyncio.create_task(server.wait_for_termination())

    try:
        await shutdown_event.wait()
        logger.info("Shutdown signal received")
    finally:
        server_task.cancel()
        await servicer.close()
        await server.stop(0.5)
        logger.info("Server stopped gracefully")


def main():
    parser = argparse.ArgumentParser(description='Snakepit Multi-Threaded gRPC Server')
    parser.add_argument('--port', type=int, default=0, help='Port to listen on')
    parser.add_argument('--adapter', type=str, required=True, help='Adapter class path')
    parser.add_argument('--elixir-address', type=str, required=True, help='Elixir server address')
    parser.add_argument('--max-workers', type=int, default=10, help='Thread pool size')
    parser.add_argument('--thread-safety-check', action='store_true', help='Enable thread safety checks')
    parser.add_argument('--snakepit-run-id', type=str, default='', help='Snakepit run ID')

    args = parser.parse_args()

    # Signal handlers
    shutdown_event = None

    def handle_signal(signum, frame):
        if shutdown_event and not shutdown_event.is_set():
            asyncio.get_running_loop().call_soon_threadsafe(shutdown_event.set)

    signal.signal(signal.SIGTERM, handle_signal)
    signal.signal(signal.SIGINT, handle_signal)

    # Run server
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    shutdown_event = asyncio.Event()

    try:
        loop.run_until_complete(serve_threaded(
            args.port,
            args.adapter,
            args.elixir_address,
            args.max_workers,
            args.thread_safety_check,
            shutdown_event
        ))
    except BaseException as e:
        logger.error(f"Unhandled exception: {type(e).__name__}: {e}")
        traceback.print_exc()
        sys.exit(1)
    finally:
        loop.close()


if __name__ == '__main__':
    main()