import time
from typing import Dict, Any, List, Optional, Union
import json
import os
import hashlib
import zlib
import logging
import shutil
import threading
import asyncio
from functools import partial
from app.core.grpc.utils.crypto import sign_payload
from mesh_core import agent_pb2
from app.db.session import get_db_session
from app.db.models import Session

logger = logging.getLogger(__name__)

class TaskAssistant:
    """The 'Brain' of the Orchestrator: High-Level AI API for Dispatching Tasks."""
    def __init__(self, registry, journal, pool, mirror=None):
        self.registry = registry
        self.journal = journal
        self.pool = pool
        self.mirror = mirror
        self.memberships = {} # session_id -> list(node_id)
        self.membership_lock = threading.Lock()
        self.retry_limit = 2
        self.idempotent_task_types = {"shell", "file_sync", "ls", "cat"}

    def recover_node_tasks(self, node_id: str):
        """
        The Orchestrator's Resurrection Logic:
        Identifies interrupted tasks from a dead node and decides whether to retry them.
        """
        logger.info(f"[📋🏥] Orchestrator: Recovering tasks for failed node {node_id}")
        
        # 1. Identify all tasks currently in-flight on this node
        interrupted_tasks = []
        for shard in self.journal.shards:
            with shard["lock"]:
                for tid, t in shard["tasks"].items():
                    if t.get("node_id") == node_id and t.get("result") is None:
                        # Capture task details before failing it in the journal
                        interrupted_tasks.append({
                            "task_id": tid,
                            "node_id": node_id,
                            "type": "shell", # Default, should ideally be stored in journal
                            "payload": t.get("stream_buffer"), # Partial context
                            "created_at": t.get("created_at")
                        })

        if not interrupted_tasks:
            logger.info(f"[📋🏥] No active tasks to recover for {node_id}.")
            return

        # 2. Mark them as failed in the Journal so UI/API callers don't hang
        # We use a special status "RETRYING" if we plan to resurrection
        self.journal.fail_node_tasks(node_id, f"Node {node_id} disconnected. Orchestrator attempting recovery.")

        # 3. Decision Phase: Resurrection
        for task in interrupted_tasks:
            tid = task["task_id"]
            
            # Policy: Only retry tasks that are idempotent or haven't timed out
            # In this simple implementation, we check the task_type if available
            is_safe_to_retry = task["type"] in self.idempotent_task_types
            
            if is_safe_to_retry:
                logger.info(f"[📋🔄] Task {tid} is idempotent. Re-injecting into Global Work Pool for a different node.")
                # Note: This is a simplified re-injection. Real implementation would 
                # need the full TaskRequest payload stored in the journal.
                # For now, we emit a recovery signal.
                self.registry.emit(node_id, "task_retry", {"task_id": tid, "original_node": node_id})
                
                # If we had the full command, we would call self.pool.add(task_request)
            else:
                logger.warning(f"[📋❌] Task {tid} is NOT idempotent. Cannot safely retry automatically.")

    def push_workspace(self, node_id, session_id):
        """Initial unidirectional push from server ghost mirror to a node."""
        if not self.mirror: return
        
        # 1. Ensure Server Mirror exists immediately
        manifest = self.mirror.generate_manifest(session_id)

        # 2. Track relationship for recovery/reconciliation
        with self.membership_lock:
            if session_id not in self.memberships:
                self.memberships[session_id] = []
            if node_id not in self.memberships[session_id]:
                self.memberships[session_id].append(node_id)

        # 3. If node is online, push actual data
        node = self.registry.get_node(node_id)
        if not node:
            logger.info(f"[📁📤] Workspace {session_id} prepared on server for offline node {node_id}")
            return
            
        print(f"[📁📤] Initiating Workspace Push for Session {session_id} to {node_id}")
        
        # Send Manifest to Node. The node will compare this with its local state
        # and send back RECONCILE_REQUIRED for any files it is missing.
        # This prevents the "Double Push" race where the server blasts data
        # while the node is still trying to decide what it needs.
        node.send_message(agent_pb2.ServerTaskMessage(
            file_sync=agent_pb2.FileSyncMessage(
                session_id=session_id,
                manifest=manifest
            )
        ), priority=1)
        
        # NOTE: Proactive parallel push removed. Manifest-driven reactive sync is cleaner.

    def push_file(self, node_id, session_id, rel_path):
        """Pushes a specific file to a node (used for drift recovery)."""
        node = self.registry.get_node(node_id)
        if not node: return
        
        workspace = self.mirror.get_workspace_path(session_id)
        abs_path = os.path.join(workspace, rel_path)
        
        if not os.path.exists(abs_path):
            print(f"    [📁❓] Requested file {rel_path} not found in mirror")
            return

        # Line-rate Optimization: 4MB chunks + No Software Throttling
        hasher = hashlib.sha256()
        file_size = os.path.getsize(abs_path)
        
        try:
            with open(abs_path, "rb") as f:
                total_chunks = (file_size + 4 * 1024 * 1024 - 1) // (4 * 1024 * 1024) if file_size > 0 else 1
                # Use DataChunker for standardized segmenting
                from mesh_core.utils import DataChunker
                
                # Construct payload
                # Note: We still iterate because we need to track index/final for business logic
                for index, chunk in enumerate(DataChunker.chunk_file(f)):
                    hasher.update(chunk)
                    offset = f.tell() - len(chunk)
                    is_final = f.tell() >= file_size
                    
                    node.send_message(agent_pb2.ServerTaskMessage(
                        file_sync=agent_pb2.FileSyncMessage(
                            session_id=session_id,
                            file_data=agent_pb2.FilePayload(
                                path=rel_path,
                                chunk=zlib.compress(chunk),
                                chunk_index=index,
                                is_final=is_final,
                                hash=hasher.hexdigest() if is_final else "",
                                offset=offset,
                                compressed=True,
                                total_chunks=total_chunks if "total_chunks" in agent_pb2.FilePayload.DESCRIPTOR.fields_by_name else 0,
                                total_size=file_size if "total_size" in agent_pb2.FilePayload.DESCRIPTOR.fields_by_name else 0
                            )
                        )
                    ), priority=2)
                    if is_final: break
        except Exception as e:
            logger.error(f"[📁📤] Line-rate push error for {rel_path}: {e}")

    def clear_workspace(self, node_id, session_id):
        """Sends a SyncControl command to purge the local sync directory on a node, and removes from active mesh."""
        print(f"    [📁🧹] Instructing node {node_id} to purge workspace for session {session_id}")
        if session_id in self.memberships and node_id in self.memberships[session_id]:
            self.memberships[session_id].remove(node_id)
        
        node = self.registry.get_node(node_id)
        if not node: return
        node.send_message(agent_pb2.ServerTaskMessage(
            file_sync=agent_pb2.FileSyncMessage(
                session_id=session_id,
                control=agent_pb2.SyncControl(action=agent_pb2.SyncControl.PURGE, path=".")
            )
        ), priority=1)

    def reconcile_node(self, node_id):
        """Forces a re-sync check for all sessions this node belongs to and purges dead sessions."""
        print(f"    [📁🔄] Triggering Resync Check for {node_id}...")
        
        active_sessions = []
        try:
            with get_db_session() as db:
                sessions = db.query(Session).filter(
                    Session.is_archived == False,
                    Session.sync_workspace_id.isnot(None)
                ).all()
                
                with self.membership_lock:
                    for s in sessions:
                        attached = s.attached_node_ids or []
                        if node_id in attached:
                            active_sessions.append(s.sync_workspace_id)
                            if s.sync_workspace_id not in self.memberships:
                                self.memberships[s.sync_workspace_id] = []
                            if node_id not in self.memberships[s.sync_workspace_id]:
                                self.memberships[s.sync_workspace_id].append(node_id)
                
                # Aggressive memory cleanup: Purge orphaned session memberships
                current_active_workspace_ids = {s.sync_workspace_id for s in sessions}
                with self.membership_lock:
                    to_purge = [sid for sid in self.memberships.keys() if sid not in current_active_workspace_ids]
                    for sid in to_purge:
                        del self.memberships[sid]
        except Exception as e:
            print(f"    [📁⚠️] Failed to fetch active sessions for node reconciliation: {e}")
            # Fallback to in-memory if DB fails
            with self.membership_lock:
                for sid, nodes in self.memberships.items():
                    if node_id in nodes:
                        active_sessions.append(sid)
                
        # Send proactive cleanup payload with the active sessions whitelist
        node = self.registry.get_node(node_id)
        if node:
            node.send_message(agent_pb2.ServerTaskMessage(
                file_sync=agent_pb2.FileSyncMessage(
                    session_id="global",
                    control=agent_pb2.SyncControl(
                        action=agent_pb2.SyncControl.CLEANUP,
                        request_paths=active_sessions
                    )
                )
            ), priority=0)
            
        for sid in active_sessions:
            # 1. Re-push manifest to trigger node-side drift check
            self.push_workspace(node_id, sid)
            
            # 2. PROACTIVE: Ensure the node starts watching this session
            # This is critical for catching 'dd' or 'cp' writes on the node
            self.control_sync(node_id, sid, action="START")
            
            # Add a small delay to prevent saturating the gRPC stream for multiple sessions
            time.sleep(0.5)

    def broadcast_file_chunk(self, session_id: str, sender_node_id: str, file_payload):
        """Broadcasts a file chunk received from one node to all other nodes in the mesh."""
        with self.membership_lock:
            session_members = self.memberships.get(session_id, [])
            destinations = [n for n in session_members if n != sender_node_id]
        
        if destinations:
            print(f"    [📁📢] Broadcasting {file_payload.path} from {sender_node_id} to: {', '.join(destinations)}")
        
        def _send_to_node(nid):
            node = self.registry.get_node(nid)
            if node:
                # Forward the exact same FileSyncMessage (Priority 2 for Sync Data)
                node.send_message(agent_pb2.ServerTaskMessage(
                    file_sync=agent_pb2.FileSyncMessage(
                        session_id=session_id,
                        file_data=file_payload
                    )
                ), priority=2)

        # M6 Fix: Do not use BoundedThreadPoolExecutor for chunks because it scrambles order!
        # node.queue limits provide the backpressure. We iterate and put sequentially.
        for nid in destinations:
            _send_to_node(nid)
    def broadcast_delete(self, session_id: str, sender_node_id: str, rel_path: str):
        """Broadcasts a file deletion from one node to all other nodes in the session mesh."""
        with self.membership_lock:
            session_members = self.memberships.get(session_id, [])
            destinations = [n for n in session_members if n != sender_node_id]

        if destinations:
            print(f"    [📁🗑️📢] Broadcasting DELETE {rel_path} from {sender_node_id} to: {', '.join(destinations)}")

        tid = f"fs-rm-bcast-{int(time.time()*1000)}"
        for nid in destinations:
            node = self.registry.get_node(nid)
            if node:
                node.send_message(agent_pb2.ServerTaskMessage(
                    file_sync=agent_pb2.FileSyncMessage(
                        session_id=session_id,
                        task_id=tid,
                        control=agent_pb2.SyncControl(
                            action=agent_pb2.SyncControl.DELETE,
                            path=rel_path
                        )
                    )
                ), priority=2)

    def lock_workspace(self, node_id, session_id):
        """Disables user-side synchronization from a node during AI refactors."""
        self.control_sync(node_id, session_id, action="LOCK")

    def unlock_workspace(self, node_id, session_id):
        """Re-enables user-side synchronization from a node."""
        self.control_sync(node_id, session_id, action="UNLOCK")

    def request_manifest(self, node_id, session_id, path="."):
        """Requests a full directory manifest from a node for drift checking."""
        node = self.registry.get_node(node_id)
        if not node: return
        node.send_message(agent_pb2.ServerTaskMessage(
            file_sync=agent_pb2.FileSyncMessage(
                session_id=session_id,
                control=agent_pb2.SyncControl(action=agent_pb2.SyncControl.REFRESH_MANIFEST, path=path)
            )
        ), priority=1)

    def control_sync(self, node_id, session_id, action="START", path="."):
        """Sends a SyncControl command to a node (e.g. START_WATCHING, LOCK)."""
        node = self.registry.get_node(node_id)
        if not node: return

        action_map = {
            "START": agent_pb2.SyncControl.START_WATCHING,
            "STOP": agent_pb2.SyncControl.STOP_WATCHING,
            "LOCK": agent_pb2.SyncControl.LOCK,
            "UNLOCK": agent_pb2.SyncControl.UNLOCK,
            "RESYNC": agent_pb2.SyncControl.RESYNC
        }
        proto_action = action_map.get(action, agent_pb2.SyncControl.START_WATCHING)

        # Track for recovery & broadcast
        if session_id not in self.memberships:
            self.memberships[session_id] = []
        if node_id not in self.memberships[session_id]:
            self.memberships[session_id].append(node_id)
        
        node.send_message(agent_pb2.ServerTaskMessage(
            file_sync=agent_pb2.FileSyncMessage(
                session_id=session_id,
                control=agent_pb2.SyncControl(action=proto_action, path=path)
            )
        ), priority=1)

    # ==================================================================
    #  Modular FS Explorer / Mesh Navigation
    # ==================================================================

    def ls(self, node_id: str, path: str = ".", timeout=10, session_id="__fs_explorer__", force_remote: bool = False):
        """Requests a directory listing from a node (waits for response)."""
        # Phase 1: Local Mirror Fast-Path
        if session_id != "__fs_explorer__" and self.mirror and not force_remote:
            workspace = self.mirror.get_workspace_path(session_id)
            abs_path = os.path.normpath(os.path.join(workspace, path.lstrip("/")))
            if os.path.exists(abs_path) and os.path.isdir(abs_path):
                files = []
                try:
                    for entry in os.scandir(abs_path):
                        if entry.name in [".cortex_sync"] and path in [".", "", "/"]:
                            continue
                        
                        # Native Orphan Syslink Cleanup
                        if entry.is_symlink() and not os.path.exists(entry.path):
                            try:
                                os.unlink(entry.path)
                                logger.info(f"[📁🧹] Cleaned up broken ghost symlink during refresh: {entry.name}")
                            except: pass
                            continue
                            
                        rel = os.path.relpath(entry.path, workspace)
                        try:
                            f_size = entry.stat().st_size if entry.is_file() else 0
                        except: f_size = 0

                        files.append({
                            "path": rel,
                            "name": entry.name,
                            "is_dir": entry.is_dir() if not entry.is_symlink() else os.path.isdir(entry.path),
                            "size": f_size,
                            "is_synced": True
                        })
                    return {"files": files, "path": path}
                except Exception as e:
                    logger.error(f"[📁📂] Local ls error for {session_id}/{path}: {e}")
            elif node_id in ["hub", "server", "local"]:
                return {"files": [], "path": path}

        node = self.registry.get_node(node_id)
        if not node:
            if node_id in ["hub", "server", "local"]:
                return {"files": [], "path": path}
            return {"error": "Offline"}

        tid = f"fs-ls-{int(time.time()*1000)}"
        event = self.journal.register(tid, node_id)
        
        node.send_message(agent_pb2.ServerTaskMessage(
            file_sync=agent_pb2.FileSyncMessage(
                session_id=session_id,
                task_id=tid,
                control=agent_pb2.SyncControl(action=agent_pb2.SyncControl.LIST, path=path)
            )
        ), priority=1)

        if event.wait(timeout):
            res = self.journal.get_result(tid)
            self.journal.pop(tid)
            
            # Proactive Mirroring: start fetching content so dots turn green
            if res and "files" in res and session_id != "__fs_explorer__":
                self._proactive_explorer_sync(node_id, res["files"], session_id)

            return res
        self.journal.pop(tid)
        return {"error": "Timeout"}

    async def als(self, node_id: str, path: str = ".", timeout=10, session_id="__fs_explorer__", force_remote: bool = False):
        """Async wrapper for ls."""
        return await asyncio.to_thread(self.ls, node_id, path, timeout, session_id, force_remote)

    def _proactive_explorer_sync(self, node_id, files, session_id):
        """Starts background tasks to mirror files to Hub so dots turn green."""
        for f in files:
            if f.get("is_dir"): continue
            if not f.get("is_synced") and f.get("size", 0) < 1024 * 512: # Skip large files
                # M6: Use shared registry executor instead of spawning loose threads
                if self.registry.executor:
                    self.registry.executor.submit(self.cat, node_id, f["path"], 15, session_id)

    def cat(self, node_id: str, path: str, timeout=15, session_id="__fs_explorer__", force_remote: bool = False):
        """Requests file content from a node (waits for result)."""
        # Phase 1: Local Mirror Fast-Path
        if session_id != "__fs_explorer__" and self.mirror and not force_remote:
            workspace = self.mirror.get_workspace_path(session_id)
            abs_path = os.path.normpath(os.path.join(workspace, path.lstrip("/")))
            if os.path.exists(abs_path) and os.path.isfile(abs_path):
                try:
                    # Try reading as text
                    with open(abs_path, 'r', encoding='utf-8', errors='ignore') as f:
                        content = f.read()
                    return {"content": content, "path": path}
                except Exception as e:
                    logger.error(f"[📁📄] Local cat error for {session_id}/{path}: {e}")
            else:
                # File is not in the mirror — it was either never written or already deleted.
                # Don't round-trip to the node: return instantly so verify-after-delete
                # doesn't stall for the full 15s journal timeout.
                return {"error": "File not found"}

        node = self.registry.get_node(node_id)
        if not node:
            if node_id in ["hub", "server", "local"]:
                return {"error": "File not found"}
            return {"error": "Offline"}

        # For 'cat', we might get multiple chunks, but TaskJournal fulfill 
        # usually happens on the final chunk. We'll handle chunking in server.
        tid = f"fs-cat-{int(time.time()*1000)}"
        event = self.journal.register(tid, node_id)
        
        node.send_message(agent_pb2.ServerTaskMessage(
            file_sync=agent_pb2.FileSyncMessage(
                session_id=session_id,
                task_id=tid,
                control=agent_pb2.SyncControl(action=agent_pb2.SyncControl.READ, path=path)
            )
        ), priority=1)

        if event.wait(timeout):
            res = self.journal.get_result(tid)
            self.journal.pop(tid)
            # res usually contains {content, path}. grpc_server already writes it to mirror.
            return res
        self.journal.pop(tid)
        return {"error": "Timeout"}

    async def acat(self, node_id: str, path: str, timeout=15, session_id="__fs_explorer__", force_remote: bool = False):
        """Async wrapper for cat."""
        return await asyncio.to_thread(self.cat, node_id, path, timeout, session_id, force_remote)

    def write(self, node_id: str, path: str, content: Union[bytes, str] = b"", is_dir: bool = False, timeout=10, session_id="__fs_explorer__"):
        """Creates or updates a file/directory on a node (waits for status)."""
        if isinstance(content, str):
            content = content.encode("utf-8")
        node = self.registry.get_node(node_id)
        
        # --- MESH SYNC MODE (Session-Aware) ---
        # If we have a session, we prioritize writing to the local mirror even if the target node is offline.
        # The node will reconcile its state from the Hub mirror upon reconnection.
        if self.mirror and session_id != "__fs_explorer__":
            workspace_mirror = self.mirror.get_workspace_path(session_id)
            dest = os.path.normpath(os.path.join(workspace_mirror, path.lstrip("/")))
            try:
                if is_dir:
                    os.makedirs(dest, exist_ok=True)
                else:
                    os.makedirs(os.path.dirname(dest), exist_ok=True)
                    with open(dest, "wb") as f:
                        f.write(content)
                
                # Multi-node broadcast for sessions
                targets = []
                if session_id != "__fs_explorer__":
                    targets = self.memberships.get(session_id, [node_id])
                else:
                    targets = [node_id]

                print(f"[📁✏️] AI Write: {path} (Session: {session_id}) -> Dispatching to {len(targets)} nodes")
                
                for target_nid in targets:
                    target_node = self.registry.get_node(target_nid)
                    if not target_node: continue
                    
                    tid = f"fs-write-{int(time.time()*1000)}"
                    target_node.send_message(agent_pb2.ServerTaskMessage(
                        file_sync=agent_pb2.FileSyncMessage(
                            session_id=session_id,
                            task_id=tid,
                            control=agent_pb2.SyncControl(
                                action=agent_pb2.SyncControl.WRITE, 
                                path=path, 
                                content=content,
                                is_dir=is_dir
                            )
                        )
                    ), priority=2)
                
                return {"success": True, "message": f"Synchronized to local mirror and dispatched to {len(targets)} nodes"}
            except Exception as e:
                logger.error(f"[📁✏️] Local mirror write error: {e}")
                return {"error": str(e)}

        if not node:
            if node_id in ["hub", "server", "local"]:
                return {"success": True, "message": "Operation completed on Hub local node"}
            return {"error": f"Node {node_id} Offline"}

        # Legacy/Explorer path: await node confirmation
        tid = f"fs-write-{int(time.time()*1000)}"
        event = self.journal.register(tid, node_id)
        
        node.send_message(agent_pb2.ServerTaskMessage(
            file_sync=agent_pb2.FileSyncMessage(
                session_id=session_id,
                task_id=tid,
                control=agent_pb2.SyncControl(
                    action=agent_pb2.SyncControl.WRITE, 
                    path=path, 
                    content=content,
                    is_dir=is_dir
                )
            )
        ), priority=2)

        if event.wait(timeout):
            res = self.journal.get_result(tid)
            self.journal.pop(tid)
            return res
        self.journal.pop(tid)
        return {"error": "Timeout"}

    async def awrite(self, node_id: str, path: str, content: Union[bytes, str] = b"", is_dir: bool = False, timeout=10, session_id="__fs_explorer__"):
        """Async wrapper for write."""
        return await asyncio.to_thread(self.write, node_id, path, content, is_dir, timeout, session_id)

    def inspect_drift(self, node_id: str, path: str, session_id: str):
        """Returns a unified diff between Hub local mirror and Node's actual file."""
        if not self.mirror: return {"error": "Mirror not available"}
        
        # 1. Get Local Content
        workspace = self.mirror.get_workspace_path(session_id)
        local_abs = os.path.normpath(os.path.join(workspace, path.lstrip("/")))
        local_content = ""
        if os.path.exists(local_abs) and os.path.isfile(local_abs):
            try:
                with open(local_abs, 'r', encoding='utf-8', errors='ignore') as f:
                    local_content = f.read()
            except: pass

        # 2. Get Remote Content (Force Bypass Fast-Path)
        print(f"    [📁🔍] Inspecting Drift: Fetching remote content for {path} on {node_id}")
        remote_res = self.cat(node_id, path, session_id=session_id, force_remote=True)
        if "error" in remote_res:
             return {"error": f"Failed to fetch remote content: {remote_res['error']}"}
        
        remote_content = remote_res.get("content", "")

        # 3. Create Diff
        import difflib
        diff = difflib.unified_diff(
            local_content.splitlines(keepends=True),
            remote_content.splitlines(keepends=True),
            fromfile=f"hub://{session_id}/{path}",
            tofile=f"node://{node_id}/{path}"
        )
        
        diff_text = "".join(diff)
        return {
            "path": path,
            "has_drift": local_content != remote_content,
            "diff": diff_text,
            "local_size": len(local_content),
            "remote_size": len(remote_content)
        }

    def rm(self, node_id: str, path: str, timeout=10, session_id="__fs_explorer__"):
        """Requests deletion of a file/directory on a node (waits for status)."""
        node = self.registry.get_node(node_id)
        
        # --- MESH SYNC MODE (Session-Aware) ---
        if self.mirror and session_id != "__fs_explorer__":
            workspace_mirror = self.mirror.get_workspace_path(session_id)
            dest = os.path.normpath(os.path.join(workspace_mirror, path.lstrip("/")))
            try:
                if os.path.exists(dest):
                    if os.path.isdir(dest): import shutil; shutil.rmtree(dest)
                    else: os.remove(dest)
                
                # Broadcast delete to other nodes
                targets = [nid for nid in self.registry._nodes if nid != node_id]
                for target_id in targets:
                    target_node = self.registry.get_node(target_id)
                    if target_node:
                        target_node.send_message(agent_pb2.ServerTaskMessage(
                            file_sync=agent_pb2.FileSyncMessage(
                                session_id=session_id,
                                task_id=f"sync-rm-{int(time.time()*1000)}",
                                control=agent_pb2.SyncControl(action=agent_pb2.SyncControl.DELETE, path=path)
                            )
                        ), priority=1)
            except Exception as e:
                logger.error(f"[📁🗑️] Sync delete error for {session_id}/{path}: {e}")

        if not node:
            return {"error": "Offline"}

        tid = f"fs-rm-{int(time.time()*1000)}"
        event = self.journal.register(tid, node_id)
        
        node.send_message(agent_pb2.ServerTaskMessage(
            file_sync=agent_pb2.FileSyncMessage(
                session_id=session_id,
                task_id=tid,
                control=agent_pb2.SyncControl(action=agent_pb2.SyncControl.DELETE, path=path)
            )
        ), priority=1)

        if event.wait(timeout):
            res = self.journal.get_result(tid)
            self.journal.pop(tid)
            return res
        self.journal.pop(tid)
        return {"error": "Timeout"}

    async def arm(self, node_id: str, path: str, timeout=10, session_id="__fs_explorer__"):
        """Async wrapper for rm."""
        return await asyncio.to_thread(self.rm, node_id, path, timeout, session_id)

    def move(self, session_id: str, old_path: str, new_path: str):
        """Orchestrates an atomic move/rename across the mesh."""
        if not self.mirror: return {"error": "Mirror not available"}
        
        try:
            # 1. Hub-side move
            self.mirror.move_file(session_id, old_path, new_path)
            
            # 2. Broadcast Deletion of old path to all mesh nodes
            self.broadcast_delete(session_id, "hub", old_path)
            
            # 3. Trigger reconciliation for new path across mesh
            # Send latest manifest to nodes so they Pull the new path
            with self.membership_lock:
                members = self.memberships.get(session_id, [])
            for nid in members:
                self.push_workspace(nid, session_id)
                
            return {"success": True, "message": f"Moved {old_path} -> {new_path}"}
        except Exception as e:
            logger.error(f"[📁🚚] Mesh move error: {e}")
            return {"error": str(e)}

    def copy(self, session_id: str, old_path: str, new_path: str):
        """Orchestrates a copy operation across the mesh."""
        if not self.mirror: return {"error": "Mirror not available"}
        
        try:
            # 1. Hub-side copy
            self.mirror.copy_file(session_id, old_path, new_path)
            
            # 2. Re-push workspace to trigger node-side pulls for the new copy
            with self.membership_lock:
                members = self.memberships.get(session_id, [])
            for nid in members:
                self.push_workspace(nid, session_id)
                
            return {"success": True, "message": f"Copied {old_path} -> {new_path}"}
        except Exception as e:
            logger.error(f"[📁👯] Mesh copy error: {e}")
            return {"error": str(e)}

    def stat(self, session_id: str, path: str):
        """Retrieves file metadata from the Hub mirror (zero mesh latency)."""
        if not self.mirror: return {"error": "Mirror not available"}
        return self.mirror.stat_file(session_id, path)

    def dispatch_swarm(self, node_ids, cmd, timeout=120, session_id=None, no_abort=False):
        """Dispatches a command to multiple nodes in parallel and waits for all results."""
        from concurrent.futures import ThreadPoolExecutor, as_completed
        
        results = {}
        with ThreadPoolExecutor(max_workers=max(1, len(node_ids))) as executor:
            future_to_node = {
                executor.submit(self.dispatch_single, nid, cmd, timeout, session_id, no_abort): nid 
                for nid in node_ids
            }
            # Use as_completed to avoid blocking on a slow node when others are finished
            for future in as_completed(future_to_node):
                node_id = future_to_node[future]
                try:
                    results[node_id] = future.result()
                except Exception as exc:
                    results[node_id] = {"error": str(exc)}
        
        return results

    def dispatch_single(self, node_id, cmd, timeout=120, session_id=None, no_abort=False):
        """Dispatches a shell command to a specific node."""
        import uuid
        node = self.registry.get_node(node_id)
        if not node: return {"error": f"Node {node_id} Offline"}

        # Use UUID to prevent timestamp collisions in high-speed swarm dispatch
        tid = f"task-{uuid.uuid4().hex[:12]}"
        event = self.journal.register(tid, node_id)
        
        # 12-Factor Signing Logic
        sig = sign_payload(cmd)
        req = agent_pb2.ServerTaskMessage(task_request=agent_pb2.TaskRequest(
            task_id=tid, task_type="shell", payload_json=cmd, signature=sig, session_id=session_id,
            timeout_ms=timeout * 1000))
        
        # Diagnostic logging for swarm performance tracking
        node_stats = node.stats
        active_cnt = node_stats.get('active_worker_count', 'unknown')
        logger.info(f"[📤] Dispatching shell {tid} to {node_id} (Timeout: {timeout}s | Active Workers: {active_cnt})")
        
        self.registry.emit(node_id, "task_assigned", {"command": cmd, "session_id": session_id}, task_id=tid)
        node.send_message(req, priority=1)
        self.registry.emit(node_id, "task_start", {"command": cmd}, task_id=tid)
        
        # Immediate peek if timeout is 0
        if timeout == 0:
            return {"status": "RUNNING", "stdout": "", "task_id": tid}

        if event.wait(timeout):
            res = self.journal.get_result(tid)
            # pop only if fully done
            if res.get("status") != "RUNNING":
                self.journal.pop(tid)
            return res
        
        # M6: Timeout recovery.
        if no_abort:
            logger.info(f"[⏳] Shell task {tid} TIMEOUT (no_abort=True). Leaving alive on {node_id}.")
            res = self.journal.get_result(tid) or {}
            res["task_id"] = tid
            res["status"] = "TIMEOUT_PENDING"
            return res

        logger.warning(f"[⚠️] Shell task {tid} TIMEOUT after {timeout}s on {node_id}. Sending ABORT.")
        try:
            node.send_message(agent_pb2.ServerTaskMessage(task_cancel=agent_pb2.TaskCancelRequest(task_id=tid)), priority=0)
        except: pass
        
        # Return partial result captured in buffer before popping
        res = self.journal.get_result(tid)
        self.journal.pop(tid)
        return res if res else {"error": "Timeout", "stdout": "", "stderr": "", "status": "TIMEOUT", "task_id": tid}

    async def adispatch_single(self, node_id, cmd, timeout=120, session_id=None, no_abort=False):
        """Async wrapper for dispatch_single."""
        return await asyncio.to_thread(self.dispatch_single, node_id, cmd, timeout, session_id, no_abort)

    async def adispatch_swarm(self, node_ids: List[str], command: str, timeout=30, session_id="__fs_explorer__"):
        """Async wrapper for dispatch_swarm."""
        return await asyncio.to_thread(self.dispatch_swarm, node_ids, command, timeout, session_id)

    async def await_for_swarm(self, task_map: Dict[str, str], timeout=30, no_abort=False):
        """Async wrapper for wait_for_swarm."""
        return await asyncio.to_thread(self.wait_for_swarm, task_map, timeout, no_abort)

    def wait_for_swarm(self, task_map, timeout=30, no_abort=False):
        """Waits for multiple tasks (map of node_id -> task_id) in parallel."""
        from concurrent.futures import ThreadPoolExecutor
        
        results = {}
        with ThreadPoolExecutor(max_workers=max(1, len(task_map))) as executor:
            # item = (node_id, task_id)
            future_to_node = {
                executor.submit(self.wait_for_task, nid, tid, timeout, no_abort): nid 
                for nid, tid in task_map.items()
            }
            for fut in future_to_node:
                nid = future_to_node[fut]
                try: results[nid] = fut.result()
                except Exception as e: results[nid] = {"error": str(e)}
        return results

    def wait_for_task(self, node_id, task_id, timeout=30, no_abort=False):
        """Waits for an existing task in the journal."""
        # Check journal first
        event = self.journal.get_task_event(task_id)
        if not event:
            return {"error": f"Task {task_id} not found in journal (finished or expired)", "status": "NOT_FOUND"}
        
        # Immediate peek if timeout is 0 or event is already set
        if timeout == 0 or event.is_set():
            res = self.journal.get_result(task_id)
            if res.get("status") != "RUNNING":
                self.journal.pop(task_id)
            return res

        logger.info(f"[⏳] Re-waiting for task {task_id} on {node_id} for {timeout}s")
        if event.wait(timeout):
            res = self.journal.get_result(task_id)
            if res.get("status") != "RUNNING":
                self.journal.pop(task_id)
            return res
        
        if no_abort:
            res = self.journal.get_result(task_id) or {}
            res["task_id"] = task_id
            res["status"] = "TIMEOUT_PENDING"
            return res

        logger.warning(f"[⚠️] Wait for task {task_id} TIMEOUT again. Sending ABORT.")
        node = self.registry.get_node(node_id)
        if node:
            try: node.send_message(agent_pb2.ServerTaskMessage(task_cancel=agent_pb2.TaskCancelRequest(task_id=task_id)), priority=0)
            except: pass
        
        res = self.journal.get_result(task_id)
        self.journal.pop(task_id)
        return res if res else {"error": "Timeout", "status": "TIMEOUT", "task_id": task_id}
