import threading
import queue
import time
import os
import logging
import json
import uuid

# removed requests import as we now validate tokens directly

import grpc
from app.protos 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
from app.config import settings

logger = logging.getLogger(__name__)

# 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.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
        
        # 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 _monitor_mesh(self):
        """Periodically prints status of all nodes in the mesh."""
        while True:
            try:
                time.sleep(10)
                active_nodes = self.registry.list_nodes()
                print("\n" + "="*50)
                print(f"📡 CORTEX MESH DASHBOARD | {len(active_nodes)} Nodes Online")
                print("-" * 50)
                if not active_nodes:
                    print("  No nodes currently connected.")
                for node in active_nodes:
                    stats = node.stats
                    tasks = stats.get("running", [])
                    capability = node.metadata.get("caps", {})
                    print(f"  🟢 {node.node_id:20} | Workers: {stats.get('active_worker_count', 0)} | Running: {len(tasks)} tasks")
                    print(f"      Capabilities: {capability}")
                print("="*50 + "\n", flush=True)
            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
                
                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]
                    
                    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:
                                _node.send_message(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 as e:
                                print(f"[📁⚠️] Failed to broadcast CLEANUP to {_node.node_id}: {e}")
                        
            except Exception as e:
                print(f"[📁⚠️] Mirror Cleanup Thread Error: {e}")
            time.sleep(600)

    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
            node.send_message(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 (Default to PERMISSIVE to match UI expectation for unconfigured nodes)
        mode_str = (sandbox_cfg.get("mode") or "PERMISSIVE").upper()
        grpc_mode = agent_pb2.SandboxPolicy.STRICT if mode_str == "STRICT" else agent_pb2.SandboxPolicy.PERMISSIVE
        
        # 2. Resolve Command Lists (fallback to some safe defaults if enabled but empty)
        allowed = sandbox_cfg.get("allowed_commands", []) or []
        if not allowed and shell_cfg.get("enabled", True):
            allowed = [
                "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"
            ]
            
        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)})")
        
        node.queue.put(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."""
        # 1. Identification Handshake
        node_id = "unknown"
        node = None
        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
            logger.info(f"[*] Node {node_id} Attempting to establish TaskStream...")
        except StopIteration:
            logger.warning("[!] Empty TaskStream received.")
            return
        except Exception as e:
            logger.error(f"[!] TaskStream handshake error: {e}")
            return
            
        node = self.registry.get_node(node_id)
        if not node:
            logger.warning(f"[!] TaskStream rejected: Node {node_id} not registered via SyncConfiguration.")
            return

        logger.info(f"[*] Node {node_id} Online (TaskStream established)")
        
        try:
            self.registry.update_stats(node_id, {"status": "streaming"})
            
            # M6: Offload reconciliation to a background thread to prevent blocking the stream initialization
            # especially when the database (on NFS) is experiencing latency.
            threading.Thread(
                target=self.assistant.reconcile_node, 
                args=(node_id,), 
                daemon=True, 
                name=f"Reconcile-{node_id}"
            ).start()

            def _read_results():
                try:
                    for msg in request_iterator:
                        try:
                            self._handle_client_message(msg, node_id, node)
                        except Exception as inner_e:
                            logger.error(f"[!] Error processing task message from {node_id}: {inner_e}", exc_info=True)
                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:
                    priority_item = node.queue.get(timeout=1.0)
                    msg = priority_item[2] # Unpack (priority, ts, msg)
                    if os.getenv("DEBUG_GRPC"):
                        kind = msg.WhichOneof("payload")
                        logger.info(f"[DEBUG-gRPC] OUTBOUND to {node_id}: {kind}")
                    yield msg
                except queue.Empty:
                    now = time.time()
                    if (now - last_keepalive) > 10.0:
                        last_keepalive = now
                        # Yield a NOP ping (work pool update) to keep proxy connections hot
                        yield agent_pb2.ServerTaskMessage(
                            work_pool_update=agent_pb2.WorkPoolUpdate(available_task_ids=self.pool.list_available())
                        )
                    continue
        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":
                logger.warning(f"[📶] gRPC Stream TERMINATED for {node_id}. Cleaning up.")
                # Fulfill any pending tasks in journal with error immediately
                self.journal.fail_node_tasks(node_id, f"Node {node_id} gRPC stream closed.")
                self.registry.deregister(node_id, record=node)

    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)
                node.send_message(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"):
                print(f"    [📁📥] Received Manifest from {node_id} for {fs.session_id}")
                
                # 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 fs.manifest.files
                    ]
                    self.journal.fulfill(task_id, {"files": files, "path": fs.manifest.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, fs.manifest)
                    if drifts:
                        print(f"    [📁🏃] Drift Detected (Node -> Server): Requesting {len(drifts)} files")
                        # Request node to push these specific files
                        # Priority 1: Drift Reconciliation Request
                        node.send_message(agent_pb2.ServerTaskMessage(
                            file_sync=agent_pb2.FileSyncMessage(
                                session_id=fs.session_id,
                                control=agent_pb2.SyncControl(
                                    action=agent_pb2.SyncControl.REFRESH_MANIFEST, 
                                    path=fs.manifest.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"):
                print(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:
                        print(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 app.protos 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) ---
    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)

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

