from typing import Dict, Any, Tuple, Callable, Optional
from app.core.tools.base import BaseToolPlugin

class MeshWaitTasksTool(BaseToolPlugin):
    @property
    def name(self) -> str:
        return "mesh_wait_tasks"
        
    @property
    def retries(self) -> int:
        return 0

    def prepare_task(self, args: Dict[str, Any], context: Dict[str, Any]) -> Tuple[Optional[Callable], Dict[str, Any]]:
        timeout = int(args.get("timeout", 30))
        no_abort = args.get("no_abort", False)
        task_map = args.get("task_map", {})
        
        assistant = context.get("assistant")
        if not assistant:
             return None, {"success": False, "error": "Orchestrator not available"}
        
        if not task_map:
            return None, {"success": False, "error": "task_map is required"}
        
        if len(task_map) == 1:
            nid, tid = next(iter(task_map.items()))
            return assistant.wait_for_task, {"node_id": nid, "task_id": tid, "timeout": timeout, "no_abort": no_abort}
        else:
            return assistant.wait_for_swarm, {"task_map": task_map, "timeout": timeout, "no_abort": no_abort}
