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

class MeshTerminalControlTool(BaseToolPlugin):
    @property
    def name(self) -> str:
        return "mesh_terminal_control"

    @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]]:
        cmd = args.get("command", "")
        timeout = int(args.get("timeout", 30))
        no_abort = args.get("no_abort", False)
        
        node_id = context.get("node_id")
        node_ids = context.get("node_ids")
        resolved_sid = context.get("session_id")
        assistant = context.get("assistant")

        if node_id in ["hub", "server", "local"] or (node_ids and any(nid in ["hub", "server", "local"] for nid in node_ids)):
            def _hub_command(**kwargs):
                return self._execute_hub_command(context.get("services"), kwargs.get("cmd"), kwargs.get("timeout"), kwargs.get("resolved_sid"))
            return _hub_command, {"cmd": cmd, "timeout": timeout, "resolved_sid": resolved_sid}
        elif node_ids and isinstance(node_ids, list):
            return assistant.dispatch_swarm, {
                "node_ids": node_ids, "cmd": cmd, "timeout": timeout, "session_id": resolved_sid, "no_abort": no_abort
            }
        elif node_id:
            return assistant.dispatch_single, {
                "node_id": node_id, "cmd": cmd, "timeout": timeout, "session_id": resolved_sid, "no_abort": no_abort
            }
        
        return None, {"success": False, "error": "node_id or node_ids is required"}

    def _execute_hub_command(self, services, cmd: str, timeout: int = 30, resolved_sid: str = None) -> dict:
        import subprocess
        cwd = os.getcwd()
        if resolved_sid and getattr(services, "orchestrator", None):
            try:
                cwd = services.orchestrator.mirror.get_workspace_path(resolved_sid)
            except: pass
            
        try:
            proc = subprocess.run(
                cmd, shell=True, capture_output=True, text=True, timeout=timeout, cwd=cwd
            )
            return {
                "status": "SUCCESS" if proc.returncode == 0 else "FAILED",
                "stdout": proc.stdout,
                "stderr": proc.stderr,
                "exit_code": proc.returncode,
                "node_id": "hub"
            }
        except subprocess.TimeoutExpired as e:
            return {"status": "TIMEOUT", "stdout": e.stdout or "", "stderr": e.stderr or "", "error": "Command timed out on Hub"}
        except Exception as e:
            return {"status": "ERROR", "error": str(e)}
