# SLAM / Nav2 subprocess management, map upload, locations
import asyncio
import json
import math
import os
import shlex
import signal
import subprocess

from .config import log, _subprocesses, _ros2_cache, _ros2_lock, _ros2_node, LOCATIONS_FILE
from .utils import _safe_name, _safe_dir


# ── Subprocess management ──

def _launch_ros2(name: str, cmd: str) -> tuple[bool, str]:
    """Launch a ROS2 subprocess. Returns (success, message)."""
    proc = _subprocesses.get(name)
    if proc and proc.poll() is None:
        return False, f"{name} is already running (pid={proc.pid})"
    try:
        full_cmd = f"source /etc/env.sh && {cmd}"
        p = subprocess.Popen(["bash", "-c", full_cmd], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, preexec_fn=os.setsid)
        _subprocesses[name] = p
        log.info(f"{name} launched: pid={p.pid} cmd={cmd}")
        return True, f"{name} started (pid={p.pid})"
    except Exception as e:
        log.error(f"{name} launch failed: {e}")
        return False, str(e)


def _kill_ros2(name: str) -> tuple[bool, str]:
    """Kill a ROS2 subprocess. Returns (success, message)."""
    proc = _subprocesses.get(name)
    if not proc or proc.poll() is not None:
        _subprocesses[name] = None
        return True, f"{name} is not running"
    try:
        os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
        proc.wait(timeout=5)
    except Exception:
        try:
            os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
        except Exception:
            pass
    _subprocesses[name] = None
    log.info(f"{name} stopped")
    return True, f"{name} stopped"


# ── SLAM handler ──

async def _handle_slam(command: str, params: dict) -> dict:
    """Handle SLAM_START, SLAM_STOP, SLAM_SAVE."""
    loop = asyncio.get_event_loop()

    if command == "SLAM_START":
        launch_file = _safe_name(params.get("launch_file", "online_async_launch"), default="online_async_launch")
        if not launch_file.endswith(".py"):
            launch_file += ".py"
        ok, msg = await loop.run_in_executor(None, _launch_ros2, "slam",
            f"ros2 launch slam_toolbox {shlex.quote(launch_file)}")
        return {"status": "ok" if ok else "error", "message": msg}

    elif command == "SLAM_STOP":
        ok, msg = await loop.run_in_executor(None, _kill_ros2, "slam")
        return {"status": "ok" if ok else "error", "message": msg}

    elif command == "SLAM_SAVE":
        map_name = _safe_name(params.get("map_name", "map"))
        save_dir = _safe_dir(params.get("save_dir", "~/maps"))
        auto_upload = params.get("auto_upload", False)
        os.makedirs(save_dir, exist_ok=True)
        try:
            safe_path = shlex.quote(f"{save_dir}/{map_name}")
            cmd = f"source /etc/env.sh && ros2 run nav2_map_server map_saver_cli -f {safe_path} --ros-args -p save_map_timeout:=10.0"
            result = subprocess.run(["bash", "-c", cmd], capture_output=True, text=True, timeout=15)
            if result.returncode != 0:
                return {"status": "error", "message": result.stderr[:200] or "map save failed"}
            save_result = {"status": "ok", "message": f"Map saved to {save_dir}/{map_name}",
                           "path": f"{save_dir}/{map_name}"}
            if auto_upload:
                upload_result = await _upload_map_files(save_dir, map_name)
                save_result["upload"] = upload_result
            return save_result
        except Exception as e:
            return {"status": "error", "message": str(e)}

    return {"status": "error", "message": f"Unknown SLAM command: {command}"}


# ── Map upload ──

async def _upload_map_files(save_dir: str, map_name: str, metadata: dict = None) -> dict:
    """PGM + YAML map files chunked upload to Gateway."""
    import struct
    from .ws_client import _send_service_request

    CHUNK_SIZE = 256 * 1024

    pgm_path = os.path.join(save_dir, f"{map_name}.pgm")
    yaml_path = os.path.join(save_dir, f"{map_name}.yaml")

    if not os.path.exists(pgm_path):
        return {"success": False, "error": f"PGM file not found: {pgm_path}"}

    parts = []
    if os.path.exists(yaml_path):
        parts.append(("yaml", open(yaml_path, "rb").read()))
    parts.append(("pgm", open(pgm_path, "rb").read()))

    header = json.dumps({"files": [{"name": p[0], "size": len(p[1])} for p in parts]}).encode()
    payload = struct.pack(">I", len(header)) + header
    for _, data in parts:
        payload += data

    chunks_hex = []
    for i in range(0, len(payload), CHUNK_SIZE):
        chunks_hex.append(payload[i:i + CHUNK_SIZE].hex())

    map_metadata = metadata or {}
    if os.path.exists(yaml_path):
        try:
            yaml_content = open(yaml_path).read()
            for line in yaml_content.split("\n"):
                if line.startswith("resolution:"):
                    map_metadata["resolution"] = float(line.split(":")[1].strip())
                elif line.startswith("origin:"):
                    map_metadata["origin"] = line.split(":")[1].strip()
        except Exception:
            pass
    map_metadata["format"] = "pgm+yaml"

    svc_data = {
        "service": "map_upload",
        "name": map_name,
        "map_type": "local",
        "chunks": chunks_hex,
        "metadata": map_metadata,
    }

    log.info(f"Uploading map '{map_name}': {len(payload)} bytes, {len(chunks_hex)} chunks")
    result = await _send_service_request(svc_data, timeout=60.0)
    return result


# ── Nav2 handler ──

async def _handle_nav2(command: str, params: dict) -> dict:
    """Handle NAV2_GOAL, NAV2_CANCEL, NAV2_STATUS."""
    loop = asyncio.get_event_loop()

    if command == "NAV2_GOAL":
        x = float(params.get("x", 0))
        y = float(params.get("y", 0))
        theta = float(params.get("theta", 0))
        try:
            cmd = f'source /etc/env.sh && python3 -c "import rclpy; from geometry_msgs.msg import PoseStamped; from nav2_simple_commander.robot_navigator import BasicNavigator; import math; rclpy.init(); nav=BasicNavigator(); goal=PoseStamped(); goal.header.frame_id=\\"map\\"; goal.pose.position.x={x}; goal.pose.position.y={y}; goal.pose.orientation.z=math.sin({theta}/2); goal.pose.orientation.w=math.cos({theta}/2); nav.goToPose(goal); print(\\"goal_sent\\"); rclpy.shutdown()"'
            result = await loop.run_in_executor(None, lambda: subprocess.run(["bash", "-c", cmd], capture_output=True, text=True, timeout=10))
            if "goal_sent" in (result.stdout or ""):
                return {"status": "ok", "message": f"Nav2 goal sent: x={x}, y={y}, theta={theta}"}
            return {"status": "error", "message": result.stderr[:200] or "nav2 goal failed"}
        except Exception as e:
            return {"status": "error", "message": str(e)}

    elif command == "NAV2_CANCEL":
        try:
            cmd = 'source /etc/env.sh && ros2 action send_goal /navigate_to_pose nav2_msgs/action/NavigateToPose "{}" --cancel'
            await loop.run_in_executor(None, lambda: subprocess.run(["bash", "-c", cmd], capture_output=True, timeout=5))
            return {"status": "ok", "message": "Navigation cancelled"}
        except Exception as e:
            return {"status": "error", "message": str(e)}

    elif command == "NAV2_STATUS":
        try:
            cmd = "source /etc/env.sh && ros2 node list 2>/dev/null | grep -c nav2"
            result = await loop.run_in_executor(None, lambda: subprocess.run(["bash", "-c", cmd], capture_output=True, text=True, timeout=5))
            count = int(result.stdout.strip()) if result.stdout.strip().isdigit() else 0
            slam_active = _subprocesses.get("slam") is not None and _subprocesses["slam"].poll() is None
            return {"status": "ok", "nav2_nodes": count, "slam_active": slam_active}
        except Exception as e:
            return {"status": "error", "message": str(e)}

    return {"status": "error", "message": f"Unknown nav2 command: {command}"}


# ── Locations ──

def _load_locations() -> dict:
    try:
        with open(LOCATIONS_FILE) as f:
            return json.load(f)
    except Exception:
        return {}


def _save_locations(locs: dict):
    with open(LOCATIONS_FILE, "w") as f:
        json.dump(locs, f, indent=2)
