import threading
import queue
import time
import os
import logging
import json
import uuid
from weakref import WeakValueDictionary

# removed requests import as we now validate tokens directly

import grpc
from mesh_core import agent_pb2, agent_pb2_grpc
from app.core.grpc.core.journal import TaskJournal
from app.core.grpc.core.pool import GlobalWorkPool
from app.core.grpc.core.mirror import GhostMirrorManager
from app.core.grpc.services.assistant import TaskAssistant
from app.core.grpc.utils.crypto import sign_payload, sign_bytes
from app.config import settings

from mesh_core.engines import MeshServerCore
from mesh_core.transport import GrpcServerTransport

logger = logging.getLogger(__name__)

# M4: Token validation is now handled directly via NodeRegistryService

# M4: Token validation is now handled directly via NodeRegistryService

class AgentOrchestrator(agent_pb2_grpc.AgentOrchestratorServicer):
    """Integrated gRPC Servicer for Agent Orchestration within AI Hub."""
    def __init__(self, registry):
        self.registry = registry  # Injected NodeRegistryService
        self.mesh_core = MeshServerCore()
        self.registry.mesh_core = self.mesh_core
        self.journal = TaskJournal()
        self.pool = GlobalWorkPool()
        self.mirror = GhostMirrorManager(storage_root=os.path.join(settings.DATA_DIR, "mirrors"))
        self.io_locks = {} # key -> threading.Lock
        self.io_locks_lock = threading.Lock()
        self.assistant = TaskAssistant(self.registry, self.journal, self.pool, self.mirror)
        self.pool.on_new_work = self._broadcast_work
        self.manifest_accumulators = {} 
        self.manifest_accumulators_lock = threading.Lock()

        # Bind Mesh Events
        self.mesh_core.on_message_received = self._on_mesh_message
        self.mesh_core.on_node_online = self._on_mesh_online
        self.mesh_core.on_node_offline = self._on_mesh_offline
        
        # 4. Mesh Observation (Aggregated Health Dashboard)
        threading.Thread(target=self._monitor_mesh, daemon=True, name="MeshMonitor").start()

        # 5. Mirror Cleanup Loop (Purge archived sessions)
        threading.Thread(target=self._mirror_cleanup_loop, daemon=True, name="MirrorCleanup").start()

    def _on_mesh_online(self, node):
        logger.info(f"[Mesh] Node Online: {node.node_id}")

    def _on_mesh_offline(self, node):
        logger.warning(f"[Mesh] Node Offline: {node.node_id}")

    def _on_mesh_message(self, node_id, msg):
        node = self.registry.get_node(node_id)
        if node:
            self._handle_client_message(msg, node_id, node)

    def _monitor_mesh(self):
        """Periodically records status of all nodes in the mesh for diagnostics."""
        log_path = os.path.join(settings.DATA_DIR, "mesh_status.log") # Use DATA_DIR for visibility
        while True:
            try:
                time.sleep(30) # Assignment #5: Min 30s sleep to prevent CPU spikes
                active_nodes = self.registry.list_nodes()
                
                status_lines = []
                status_lines.append(f"📡 CORTEX MESH DASHBOARD | {len(active_nodes)} Nodes Online | {time.ctime()}")
                status_lines.append("-" * 50)
                if not active_nodes:
                    status_lines.append("  No nodes currently connected.")
                for node in active_nodes:
                    stats = node.stats
                    tasks = stats.get("running", [])
                    capability = node.metadata.get("caps", {})
                    status_lines.append(f"  🟢 {node.node_id:20} | Workers: {stats.get('active_worker_count', 0)} | Running: {len(tasks)} tasks")
                status_lines.append("="*50 + "\n")
                
                # Dedicated status log ensures stdout stays clean for test summaries
                with open(log_path, "w") as f:
                    f.write("\n".join(status_lines))
                    
            except Exception as e:
                logger.error(f"[MeshMonitor] Error: {e}")

    def _mirror_cleanup_loop(self):
        """Periodically purges ghost mirrors for archived sessions."""
        time.sleep(10) # Initial quick cleanup shortly after boot
        while True:
            try:
                from app.db.session import get_db_session
                from app.db.models import Session
                
                # M6: Use a local list to release DB lock ASAP
                active_ids = []
                with get_db_session() as db:
                    # Fetch all workspace IDs for non-archived sessions
                    active_sessions = db.query(Session).filter(
                        Session.is_archived == False,
                        Session.sync_workspace_id.isnot(None)
                    ).all()
                    active_ids = [s.sync_workspace_id for s in active_sessions]
                
                # Performance: Purge and broadcast outside the DB session
                if self.mirror:
                    self.mirror.purge_orphaned(active_ids)
                    
                    # M6: Distributed Lock Scavenging
                    for sid in active_ids:
                        self.mirror.purge_stale_locks(sid, ttl=60) # 1-minute TTL

                    # M6: Broadcast CLEANUP to all connected nodes to do local cleanup
                    live_nodes = self.registry.list_nodes()
                    for _node in live_nodes:
                        try:
                            self.mesh_core.dispatch(_node.node_id, agent_pb2.ServerTaskMessage(
                                file_sync=agent_pb2.FileSyncMessage(
                                    session_id="system",
                                    control=agent_pb2.SyncControl(
                                        action=agent_pb2.SyncControl.CLEANUP,
                                        request_paths=active_ids
                                    )
                                )
                            ), priority=0)
                        except Exception:
                            pass
                        
            except Exception as e:
                logger.error(f"[📁⚠️] Mirror Cleanup Thread Error: {e}")
            time.sleep(600) # Assignment #5: Run every 10 minutes

    def _broadcast_work(self, _):
        """Pushes work notifications to all active nodes."""
        live_nodes = self.registry.list_nodes()
        for node in live_nodes:
            logger.debug(f"[📢] Broadcasting availability to {node.node_id}")
            # Priority 1: Interactive/Work Dispatch
            self.mesh_core.dispatch(node.node_id, agent_pb2.ServerTaskMessage(
                work_pool_update=agent_pb2.WorkPoolUpdate(available_task_ids=self.pool.list_available())
            ), priority=1)
        
    def _build_sandbox_policy(self, skill_cfg):
        """Builds a SandboxPolicy protobuf from raw skill configuration."""
        shell_cfg = skill_cfg.get("shell", {}) if isinstance(skill_cfg, dict) else {}
        if shell_cfg is None: shell_cfg = {}
        
        # Get sandbox config
        sandbox_cfg = shell_cfg.get("sandbox", {}) if isinstance(shell_cfg, dict) else {}
        if sandbox_cfg is None: sandbox_cfg = {}
        
        # 1. Resolve Mode (M6 FIX: Hub now defaults to PERMISSIVE/PASSIVE to match UI expectation)
        raw_mode = (sandbox_cfg.get("mode") or "PERMISSIVE").upper()
        # Aliasing: UI calls it PASSIVE, Backend calls it PERMISSIVE
        if raw_mode in ["PERMISSIVE", "PASSIVE"]:
            grpc_mode = agent_pb2.SandboxPolicy.PERMISSIVE
        else:
            grpc_mode = agent_pb2.SandboxPolicy.STRICT
        
        # 2. Resolve Command Lists (M6 FIX: Expanded with common Windows/PowerShell tools)
        allowed = sandbox_cfg.get("allowed_commands", []) or []
        if not allowed and shell_cfg.get("enabled", True):
            allowed = [
                # Linux/macOS Defaults
                "ls", "cat", "echo", "pwd", "uname", "curl", "python3", "git",
                "lscpu", "free", "df", "uptime", "nproc", "grep", "awk", "sed",
                "hostname", "id", "whoami", "ip", "ping", "ps", "top", "which",
                "sh", "bash", "zsh", "set", "export", "env", "find", "mkdir",
                "touch", "rm", "mv", "cp", "tail", "head", "less", "more",
                # Windows/PowerShell Defaults (M6)
                "ipconfig", "systeminfo", "ver", "dir", "cmd", "powershell", "pwsh",
                "Get-CimInstance", "Select-Object", "Measure-Object", "Get-Process",
                "Get-Date", "Get-NetIPAddress", "cls", "type", "where", "findstr",
                "tasklist", "netstat", "whoami", "hostname"
            ]
            
        denied = sandbox_cfg.get("denied_commands", []) or []
        sensitive = sandbox_cfg.get("sensitive_commands", []) or []
        jail = shell_cfg.get("cwd_jail", "")

        return agent_pb2.SandboxPolicy(
            mode=grpc_mode,
            allowed_commands=allowed,
            denied_commands=denied,
            sensitive_commands=sensitive,
            working_dir_jail=jail,
            skill_config_json=json.dumps(skill_cfg)
        )

    def push_policy(self, node_id, skill_config):
        """Pushes a live policy update to a connected node."""
        node = self.registry.get_node(node_id)
        if not node:
            logger.warning(f"[🔒] Cannot push policy to {node_id}: Node Offline")
            return
        
        policy = self._build_sandbox_policy(skill_config)
        logger.info(f"[🔒] Pushing LIVE policy update to {node_id} (Mode: {policy.mode}, Allowed: {len(policy.allowed_commands)})")
        
        self.mesh_core.dispatch(node_id, agent_pb2.ServerTaskMessage(
            policy_update=policy
        ))

    def SyncConfiguration(self, request, context):
        """M4 Authenticated Handshake: Validate invite_token via Hub DB, then send policy."""
        node_id = request.node_id
        invite_token = request.auth_token
        logger.info(f"[🔑] SyncConfiguration REQUEST from {node_id} (token prefix: {invite_token[:4]}...)")

        # --- M4: Token validation via Hub API (M6: switched to direct call to avoid deadlock) ---
        user_id = "default"
        skill_cfg = { }
        
        try:
            payload = self.registry.validate_invite_token(node_id, invite_token)
            if not payload.get("valid"):
                reason = payload.get("reason", "Token rejected") or "Token rejected"
                logger.warning(f"[🔒] SyncConfiguration REJECTED {node_id}: {reason}")
                return agent_pb2.RegistrationResponse(success=False, error_message=reason)
            
            user_id = payload.get("user_id", "default")
            skill_cfg = payload.get("skill_config", {})
            logger.info(f"[🔑] Token validated for {node_id} (owner: {user_id})")
        except Exception as e:
            logger.error(f"[🔒] SyncConfiguration FAILED internal validation for {node_id}: {e}")
            return agent_pb2.RegistrationResponse(success=False, error_message=f"Internal Hub Error: {e}")
        
        policy = self._build_sandbox_policy(skill_cfg)
        logger.info(f"[🔑] Handshake successful for {node_id} (owner: {user_id})")

        # Register the node in the centralized AI Hub registry
        self.registry.register(request.node_id, user_id, {
            "desc": request.node_description,
            "caps": dict(request.capabilities),
        })

        return agent_pb2.RegistrationResponse(
            success=True,
            policy=policy
        )

    def TaskStream(self, request_iterator, context):
        """Persistent bi-directional stream for dispatching work and collecting results."""
        node_id = "unknown"
        try:
            first_msg = next(request_iterator)
            if not first_msg.HasField("announce"):
                logger.warning("[!] Initial TaskStream message MUST be 'announce'")
                return
            node_id = first_msg.announce.node_id
        except Exception as e:
            logger.error(f"[!] TaskStream handshake error: {e}")
            return
            
        node = self.registry.get_node(node_id)
        if not node:
            if not self.registry.try_recovery_register(node_id):
                logger.warning(f"[!] TaskStream rejected: Node {node_id} not registered.")
                return
            node = self.registry.get_node(node_id)

        # M6: Eject 'Zombie' connections
        if node.stream:
            logger.warning(f"[*] Node {node_id} has a STALE TaskStream. Ejecting Zombie...")
            try: node.stream.abort(grpc.StatusCode.ABORTED, "New connection established.")
            except: pass
            time.sleep(0.1)

        # 1. Initialize Transport
        transport = GrpcServerTransport(context, signer=sign_bytes)
        node.stream = context # Still need for assistant.reconcile legacy calls
        
        # 2. Register in Mesh Core
        self.mesh_core.register_node(node_id, node.user_id, node.metadata, transport)
        
        try:
            self.registry.update_stats(node_id, {"status": "streaming"})
            
            # Start Background Reader
            def _read_results():
                try:
                    for msg in request_iterator:
                        self.mesh_core.handle_inbound(node_id, msg)
                except Exception as e:
                    logger.warning(f"Results listener closed for {node_id}: {e}")

            threading.Thread(target=_read_results, daemon=True, name=f"Results-{node_id}").start()

            # 3. Work Dispatcher (Main Stream)
            last_keepalive = 0
            while context.is_active():
                try:
                    # In this refactored version, the 'MeshNodeRecord' in MeshServerCore
                    # holds the transport, which has the send_queue.
                    # But we also have the legacy node.queue (from registry).
                    # For now, we drain both or migrate one.
                    # Let's drain transport.send_queue and node.queue.
                    
                    msg = None
                    try:
                        priority_item = transport.send_queue.get(timeout=0.05)
                        msg = priority_item[2]
                    except queue.Empty:
                        # Fallback to legacy queue for backwards compatibility
                        try:
                            priority_item = node.queue.get(timeout=0.05)
                            msg = priority_item[2]
                            # Sign it if it wasn't signed (Legacy path)
                            msg.signature = ""
                            msg_bytes = msg.SerializeToString(deterministic=True)
                            msg.signature = sign_bytes(msg_bytes)
                        except queue.Empty:
                            pass

                    if msg:
                        yield msg
                    else:
                        now = time.time()
                        if (now - last_keepalive) > 10.0:
                            last_keepalive = now
                            msg = agent_pb2.ServerTaskMessage(
                                work_pool_update=agent_pb2.WorkPoolUpdate(available_task_ids=self.pool.list_available())
                            )
                            msg_bytes = msg.SerializeToString(deterministic=True)
                            msg.signature = sign_bytes(msg_bytes)
                            yield msg
                except Exception as e:
                    logger.error(f"[!] Dispatch Error for {node_id}: {e}")
                    break
        except StopIteration:
            logger.info(f"[📶] gRPC Stream StopIteration for {node_id}")
        except Exception as e:
            logger.error(f"[!] TaskStream Error for {node_id}: {e}")
        finally:
            if node_id != "unknown" and node:
                # 0. Only nullify the record's stream if THIS context is the one owning it.
                # This prevents a new successful stream from being cleared by a late-arriving cleanup of a zombie.
                if node.stream is context:
                    node.stream = None
                
                logger.warning(f"[📶] gRPC Stream TERMINATED for {node_id}. Cleaning up.")
                
                # M6: Clean up orphaned I/O locks for this node's sessions to prevent memory leak
                with self.io_locks_lock:
                    to_purge = [k for k in self.io_locks.keys() if k.startswith(f"session:")] # More precise check needed if we have session ids
                    # Currently io_locks are keyed as "{session_id}:{path}"
                    # We can't easily map back to node_id here without more state, 
                    # but we can rely on the 10-minute periodic worker for an absolute safety net.
                    pass

                # Hand off failure handling to the Assistant's Recovery Policy (Ownership Shift)
                self.assistant.recover_node_tasks(node_id)
                self.registry.deregister(node_id, record=node)
                self.mesh_core.deregister_node(node_id)

    def _handle_client_message(self, msg, node_id, node):
        kind = msg.WhichOneof('payload')
        if os.getenv("DEBUG_GRPC"):
            # Exclude large blobs like terminal_out or file_data from full log unless needed
            if kind == 'skill_event' and msg.skill_event.WhichOneof('data') == 'terminal_out':
                logger.info(f"[DEBUG-gRPC] INBOUND from {node_id}: {kind} (terminal_out, size={len(msg.skill_event.terminal_out)})")
            elif kind == 'file_sync' and msg.file_sync.HasField('file_data'):
                logger.info(f"[DEBUG-gRPC] INBOUND from {node_id}: {kind} (file_data, path={msg.file_sync.file_data.path}, size={len(msg.file_sync.file_data.chunk)})")
            else:
                logger.info(f"[DEBUG-gRPC] INBOUND from {node_id}: {kind} | {msg}")
        
        if kind == 'task_claim':
            task_id = msg.task_claim.task_id
            success, payload = self.pool.claim(task_id, node_id)
            
            node.send_message(agent_pb2.ServerTaskMessage(
                claim_status=agent_pb2.TaskClaimResponse(
                    task_id=task_id, 
                    granted=success, 
                    reason="Task successfully claimed" if success else "Task already claimed by another node"
                )
            ), priority=1)
            # M6: Notify UI that a node is claiming a global task
            self.registry.emit(node_id, "task_claim", {"task_id": task_id, "granted": success})

            if success:
                sig = sign_payload(payload)
                self.mesh_core.dispatch(node_id, agent_pb2.ServerTaskMessage(
                    task_request=agent_pb2.TaskRequest(
                        task_id=task_id, 
                        payload_json=payload, 
                        signature=sig
                    )
                ), priority=1)
        
        elif kind == 'task_response':
            tr = msg.task_response
            res_obj = {"stdout": tr.stdout, "stderr": tr.stderr, "status": tr.status}
            self.journal.fulfill(tr.task_id, res_obj)
            
            # M6: Emit to EventBus for UI streaming
            event_type = "task_complete" if tr.status == agent_pb2.TaskResponse.SUCCESS else "task_error"
            self.registry.emit(node_id, event_type, res_obj, task_id=tr.task_id)
        
        elif kind == 'skill_event':
            se = msg.skill_event
            data_kind = se.WhichOneof('data')
            event_data = {"session_id": se.session_id, "task_id": se.task_id, "skill": "shell"}
            
            if data_kind == 'terminal_out':
                event_data["data"] = se.terminal_out
                event_data["type"] = "output"
                # NEW: Capture stream buffer in journal for tool callback/timeout
                if se.task_id:
                    self.journal.append_stream(se.task_id, se.terminal_out)
            elif data_kind == 'prompt':
                event_data["data"] = se.prompt
                event_data["type"] = "prompt"
                
                # EDGE INTELLIGENCE: Signal SubAgent a prompt was detected
                if se.task_id:
                    self.journal.signal_prompt(se.task_id)
            
            self.registry.emit(node_id, "skill_event", event_data)
        

        elif kind == 'file_sync':
            fs = msg.file_sync
            task_id = fs.task_id
            
            if fs.HasField("file_data"):
                # Keyed Lock per file to prevent write/hash/swap race
                lock_key = f"{fs.session_id}:{fs.file_data.path}"
                with self.io_locks_lock:
                    if lock_key not in self.io_locks:
                        self.io_locks[lock_key] = threading.Lock()
                    lock = self.io_locks[lock_key]
                
                with lock:
                    self.mirror.write_file_chunk(fs.session_id, fs.file_data)
                
                # M6: Clean up the lock entry after finalization to prevent memory leak
                if fs.file_data.is_final:
                    with self.io_locks_lock:
                        if lock_key in self.io_locks:
                            del self.io_locks[lock_key]
                
                self.assistant.broadcast_file_chunk(fs.session_id, node_id, fs.file_data)
                # M6: Emit sync progress (rarely to avoid flood, but good for large pushes)
                if fs.file_data.chunk_index % 10 == 0:
                    self.registry.emit(node_id, "sync_progress", {"path": fs.file_data.path, "chunk": fs.file_data.chunk_index})
            
            elif fs.HasField("manifest"):
                logger.debug(f"[📁📥] Received Manifest from {node_id} for {fs.session_id} (chunk: {fs.manifest.chunk_index}, final: {fs.manifest.is_final})")
                
                key = (node_id, fs.session_id)
                with self.manifest_accumulators_lock:
                    if key not in self.manifest_accumulators:
                        self.manifest_accumulators[key] = []
                    self.manifest_accumulators[key].extend(fs.manifest.files)
                    
                    is_final = fs.manifest.is_final
                    root_path = fs.manifest.root_path
                    
                    if is_final:
                        accumulated_files = self.manifest_accumulators.pop(key)
                
                if is_final:
                    # Create a complete manifest object for downstream processing
                    full_manifest = agent_pb2.DirectoryManifest(
                        root_path=root_path,
                        files=accumulated_files
                    )
                    
                    # M6: Handle interactive 'ls' result correlation
                    if task_id and task_id.startswith("fs-ls-"):
                        files = [
                            {"path": f.path, "name": os.path.basename(f.path) or f.path, "is_dir": f.is_dir, "size": f.size} 
                            for f in accumulated_files
                        ]
                        self.journal.fulfill(task_id, {"files": files, "path": root_path})
                    
                    # M6: Only reconcile for real user sessions, not for the modular explorer.
                    if fs.session_id != "__fs_explorer__":
                        drifts = None
                        # Do not reconcile on shallow manifests triggered by interactive FS tools
                        if task_id and any(task_id.startswith(p) for p in ("fs-ls-", "fs-write-", "fs-rm-")):
                            pass
                        else:
                            drifts = self.mirror.reconcile(fs.session_id, full_manifest)
                        if drifts:
                            logger.info(f"[📁🏃] Drift Detected (Node -> Server): Requesting {len(drifts)} files")
                            # Request node to push these specific files
                            # Priority 1: Drift Reconciliation Request
                            self.mesh_core.dispatch(node_id, agent_pb2.ServerTaskMessage(
                                file_sync=agent_pb2.FileSyncMessage(
                                     session_id=fs.session_id,
                                     control=agent_pb2.SyncControl(
                                         action=agent_pb2.SyncControl.REFRESH_MANIFEST, 
                                         path=root_path,
                                         request_paths=drifts
                                     )
                                )
                            ), priority=1)
                        else:
                            self.registry.emit(node_id, "sync_status", {"message": "Synchronized (Node -> Server)", "code": 0})

            elif fs.HasField("status"):
                logger.debug(f"[📁] Sync Status from {node_id}: {fs.status.message}")
                
                # M6: Handle interactive write/rm/ls/cat result correlation from node-side error status
                if task_id and task_id.startswith("fs-"):
                    success = fs.status.code == 0
                    if not success:
                        self.journal.fulfill(task_id, {"error": fs.status.message})
                    elif task_id.startswith("fs-write-") or task_id.startswith("fs-rm-"):
                        self.journal.fulfill(task_id, {"success": True, "message": fs.status.message})
                
                self.registry.emit(node_id, "sync_status", {"message": fs.status.message, "code": fs.status.code})
                if fs.status.code == agent_pb2.SyncStatus.RECONCILE_REQUIRED:
                    # Phase 3: Notify active orchestrators of drift
                    self.registry.emit(node_id, "mesh_observation", {
                        "type": "drift_warning",
                        "session_id": fs.session_id,
                        "paths": list(fs.status.reconcile_paths),
                        "message": f"Critical Drift: Node {node_id} has hash mismatches in {fs.session_id}."
                    })
                    for path in fs.status.reconcile_paths:
                        self.assistant.push_file(node_id, fs.session_id, path)
            
            elif fs.HasField("control"):
                ctrl = fs.control
                if ctrl.action == agent_pb2.SyncControl.DELETE:
                    path_to_del = ctrl.path
                    # Guard: never process deletes for internal temp/lock files.
                    # Nodes may fire spurious DELETE events for these when os.replace() is used.
                    if path_to_del.endswith(".cortex_tmp") or path_to_del.endswith(".cortex_lock"):
                        print(f"    [📁🚫] Ignored temp/lock DELETE from {node_id}: {path_to_del}")
                    else:
                        logger.debug(f"[📁🗑️] Node requested DELETE on mirror: {path_to_del}")
                        self.mirror.delete_file(fs.session_id, path_to_del)
                        # Broadcast delete to all other session nodes for mesh consistency
                        self.assistant.broadcast_delete(fs.session_id, node_id, path_to_del)

    def ReportHealth(self, request_iterator, context):
        """Collect Health Metrics and Feed Policy Updates."""
        for hb in request_iterator:
            self.registry.update_stats(hb.node_id, {
                "active_worker_count": hb.active_worker_count, 
                "cpu_usage_percent": hb.cpu_usage_percent,
                "memory_usage_percent": hb.memory_usage_percent,
                "cpu_count": hb.cpu_count,
                "memory_used_gb": hb.memory_used_gb,
                "memory_total_gb": hb.memory_total_gb,
                "running": list(hb.running_task_ids),
                # M6 Rich Metrics
                "cpu_usage_per_core": list(hb.cpu_usage_per_core),
                "cpu_freq_mhz": hb.cpu_freq_mhz,
                "memory_available_gb": hb.memory_available_gb,
                "load_avg": list(hb.load_avg)
            })
            yield agent_pb2.HealthCheckResponse(server_time_ms=int(time.time()*1000))

class LoggingInterceptor(grpc.ServerInterceptor):
    def intercept_service(self, continuation, handler_call_details):
        method = handler_call_details.method
        logger.info(f"[gRPC] Incoming RPC Call: {method}")
        return continuation(handler_call_details)

def serve_grpc(registry, port=50051):
    """Starts the gRPC Orchestrator Server for Agent Nodes."""
    import grpc
    from concurrent import futures
    from mesh_core import agent_pb2_grpc
    
    # M6: Increase message size limits to 128MB (prevents manifest sync drops)
    options = [
        ('grpc.max_receive_message_length', 128 * 1024 * 1024),
        ('grpc.max_send_message_length', 128 * 1024 * 1024),
        ('grpc.keepalive_time_ms', 10000),      # Send keepalive ping every 10s
        ('grpc.keepalive_timeout_ms', 10000),    # Wait 10s for pong
        ('grpc.keepalive_permit_without_calls', True),
    ]
    server = grpc.server(
        futures.ThreadPoolExecutor(max_workers=10), 
        options=options,
        interceptors=(LoggingInterceptor(),)
    )
    orchestrator = AgentOrchestrator(registry)
    agent_pb2_grpc.add_AgentOrchestratorServicer_to_server(orchestrator, server)
    
    addr = f"[::]:{port}"
    server.add_insecure_port(addr)
    
    # --- Enable Reflection (M6 Debugging) ---
    try:
        from grpc_reflection.v1alpha import reflection
        SERVICE_NAMES = (
            agent_pb2.DESCRIPTOR.services_by_name['AgentOrchestrator'].full_name,
            reflection.SERVICE_NAME,
        )
        reflection.enable_server_reflection(SERVICE_NAMES, server)
    except ImportError:
        logger.warning("[gRPC] grpcio-reflection not found. Disabling server reflection.")
    except Exception as re:
        logger.error(f"[gRPC] Reflection setup failed: {re}")

    logger.info(f"🚀 CORTEX gRPC Orchestrator starting on {addr}")
    server.start()
    return server, orchestrator

