actions.py

# actions.py
# Usage: python actions.py --action update_code [--branch feature-branch]

"""
actions.py - Drone Action Executor

This script allows for executing various actions on a drone, such as takeoff, land, hold, test, reboot, and update code.
It utilizes MAVSDK for drone communication and provides visual feedback via LEDs.

Usage Examples:
    python actions.py --action takeoff --altitude 20
    python actions.py --action update_code --branch feature-branch
    python actions.py --action update_code  # Update code without specifying a branch

"""

import argparse
import asyncio
import csv
import glob
import logging
import logging.handlers
import os
import socket
import subprocess
import sys
import time

import psutil
from mavsdk import System
from src.led_controller import LEDController
from src.params import Params

# =======================
# Configuration Constants
# =======================

GRPC_PORT = Params.DEFAULT_GRPC_PORT
UDP_PORT = Params.mavsdk_port
HW_ID = None  # Will be set by read_hw_id()

# =======================
# Configure Logging
# =======================

# Create logs directory if it doesn't exist
logs_directory = os.path.join("logs", "action_logs")
os.makedirs(logs_directory, exist_ok=True)

# Configure the logger
logger = logging.getLogger("action_logger")
logger.setLevel(logging.DEBUG)

# Create formatter
formatter = logging.Formatter(
    fmt='%(asctime)s - %(levelname)s - %(name)s: %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)

# Create console handler with a higher log level
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setLevel(logging.INFO)
console_handler.setFormatter(formatter)

# Create rotating file handler
log_file = os.path.join(logs_directory, "actions.log")
file_handler = logging.handlers.RotatingFileHandler(
    log_file, maxBytes=10*1024*1024, backupCount=5
)
file_handler.setLevel(logging.DEBUG)
file_handler.setFormatter(formatter)

# Add handlers to the logger
logger.addHandler(console_handler)
logger.addHandler(file_handler)

# =======================
# Helper Functions
# =======================

def check_mavsdk_server_running(port):
    """
    Checks if the MAVSDK server is running on the specified gRPC port.

    Args:
        port (int): The gRPC port to check.

    Returns:
        tuple: (is_running (bool), pid (int or None))
    """
    for proc in psutil.process_iter(['pid', 'name']):
        try:
            for conn in proc.connections(kind='inet'):
                if conn.laddr.port == port:
                    return True, proc.info['pid']
        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
            pass
    return False, None

def wait_for_port(port, host='localhost', timeout=10.0):
    """
    Wait until a port starts accepting TCP connections.

    Args:
        port (int): The port to check.
        host (str): The hostname to check.
        timeout (float): The maximum time to wait in seconds.

    Returns:
        bool: True if the port is open, False if the timeout was reached.
    """
    start_time = time.time()
    while True:
        try:
            with socket.create_connection((host, port), timeout=1):
                return True
        except (ConnectionRefusedError, socket.timeout, OSError) as e:
            logger.debug(f"Waiting for port {port} on {host}: {e}")
            if time.time() - start_time >= timeout:
                return False
            time.sleep(0.1)

def start_mavsdk_server(grpc_port, udp_port):
    """
    Starts the MAVSDK server on the specified gRPC and UDP ports.
    If a server is already running on the gRPC port, it terminates it before starting a new one.

    Args:
        grpc_port (int): The gRPC port for MAVSDK server.
        udp_port (int): The UDP port for MAVSDK server.

    Returns:
        subprocess.Popen: The subprocess running the MAVSDK server.
    """
    is_running, pid = check_mavsdk_server_running(grpc_port)
    if is_running:
        logger.info(f"MAVSDK server already running on port {grpc_port}. Terminating...")
        try:
            psutil.Process(pid).terminate()
            psutil.Process(pid).wait(timeout=5)
            logger.info(f"Terminated existing MAVSDK server with PID: {pid}")
        except psutil.NoSuchProcess:
            logger.warning(f"No process found with PID: {pid} to terminate.")
        except psutil.TimeoutExpired:
            logger.warning(f"Process with PID: {pid} did not terminate gracefully. Killing it.")
            psutil.Process(pid).kill()
            psutil.Process(pid).wait()
            logger.info(f"Killed MAVSDK server with PID: {pid}")

    logger.info(f"Starting MAVSDK server on gRPC port: {grpc_port}, UDP port: {udp_port}")

    # Define absolute path to mavsdk_server
    mavsdk_server_path = os.path.join(os.getcwd(), "mavsdk_server")

    if not os.path.isfile(mavsdk_server_path):
        logger.error(f"mavsdk_server executable not found at '{mavsdk_server_path}'.")
        sys.exit(1)

    try:
        mavsdk_server = subprocess.Popen(
            [mavsdk_server_path, "-p", str(grpc_port), f"udp://:{udp_port}"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE
        )

        # Log MAVSDK server output for debugging
        asyncio.create_task(log_mavsdk_output(mavsdk_server))

        # Wait until the server is listening on the gRPC port
        if not wait_for_port(grpc_port, timeout=10):
            logger.error(f"MAVSDK server did not start listening on port {grpc_port} within timeout.")
            mavsdk_server.terminate()
            sys.exit(1)

        logger.info("MAVSDK server is now listening on gRPC port.")
        return mavsdk_server
    except FileNotFoundError:
        logger.error("mavsdk_server executable not found. Ensure it's in the current directory.")
        sys.exit(1)
    except Exception:
        logger.exception("Failed to start MAVSDK server")
        sys.exit(1)

async def log_mavsdk_output(mavsdk_server):
    """
    Asynchronously logs the stdout and stderr of the MAVSDK server.

    Args:
        mavsdk_server (subprocess.Popen): The subprocess running the MAVSDK server.
    """
    try:
        while True:
            line = await asyncio.get_event_loop().run_in_executor(None, mavsdk_server.stdout.readline)
            if not line:
                break
            logger.debug(f"MAVSDK Server: {line.decode().strip()}")
    except Exception:
        logger.exception("Error while reading MAVSDK server stdout")

    try:
        while True:
            line = await asyncio.get_event_loop().run_in_executor(None, mavsdk_server.stderr.readline)
            if not line:
                break
            logger.error(f"MAVSDK Server Error: {line.decode().strip()}")
    except Exception:
        logger.exception("Error while reading MAVSDK server stderr")

def read_hw_id():
    """
    Reads the hardware ID from files with the '.hwID' extension.

    Returns:
        int or None: The hardware ID if found, else None.
    """
    hwid_files = glob.glob('*.hwID')
    if hwid_files:
        filename = hwid_files[0]
        hw_id_str = os.path.splitext(filename)[0]  # Get filename without extension
        logger.info(f"Hardware ID file detected: {filename}")
        try:
            hw_id = int(hw_id_str)
            logger.info(f"Hardware ID {hw_id} detected.")
            return hw_id
        except ValueError:
            logger.error(f"Invalid hardware ID format in file: {filename}. Expected an integer.")
            return None
    else:
        logger.warning("Hardware ID file not found.")
        return None

def read_config(filename=Params.config_csv_name):
    """
    Reads the drone configuration from a CSV file based on the hardware ID.

    Returns:
        dict or None: The drone configuration if found, else None.
    """
    global HW_ID
    logger.info("Reading drone configuration...")
    try:
        with open(filename, newline='') as csvfile:
            reader = csv.DictReader(csvfile)
            for row in reader:
                hw_id = int(row.get('hw_id', -1))
                if hw_id == HW_ID:
                    drone_config = {
                        'hw_id': hw_id,
                        'pos_id': int(row.get('pos_id', -1)),
                        'x': float(row.get('x', 0.0)),
                        'y': float(row.get('y', 0.0)),
                        'ip': row.get('ip', ''),
                        'udp_port': int(row.get('mavlink_port', UDP_PORT)),
                        'grpc_port': int(row.get('debug_port', GRPC_PORT)),
                        'gcs_ip': row.get('gcs_ip', ''),
                    }
                    logger.info(f"Drone configuration found: {drone_config}")
                    return drone_config
        logger.warning(f"No matching hardware ID {HW_ID} found in the configuration file.")
    except FileNotFoundError:
        logger.error(f"Config file '{filename}' not found.")
    except Exception:
        logger.exception(f"Error reading config file '{filename}'")
    return None

def stop_mavsdk_server(mavsdk_server):
    """
    Terminates the MAVSDK server subprocess.

    Args:
        mavsdk_server (subprocess.Popen): The subprocess running the MAVSDK server.
    """
    try:
        if mavsdk_server.poll() is None:  # Process is still running
            logger.info("Stopping MAVSDK server...")
            mavsdk_server.terminate()
            try:
                mavsdk_server.wait(timeout=5)
                logger.info("MAVSDK server terminated gracefully.")
            except subprocess.TimeoutExpired:
                logger.warning("MAVSDK server did not terminate gracefully. Killing it.")
                mavsdk_server.kill()
                mavsdk_server.wait()
                logger.info("MAVSDK server killed forcefully.")
        else:
            logger.debug("MAVSDK server has already terminated.")
    except Exception:
        logger.exception("Error stopping MAVSDK server")

async def set_parameters(drone, parameters):
    """
    Sets one or more parameters on the drone.

    Args:
        drone (System): The MAVSDK drone system.
        parameters (dict): Dictionary of parameter names and their integer values.
    """
    for param_name, param_value in parameters.items():
        try:
            logger.info(f"Setting parameter '{param_name}' to {param_value}")
            await drone.param.set_param_int(param_name, param_value)
            logger.info(f"Parameter '{param_name}' set to {param_value} successfully.")
        except Exception:
            logger.exception(f"Failed to set parameter '{param_name}'")

async def perform_action(action, altitude=None, parameters=None, branch=None):
    """
    Orchestrates the entire action execution process, including starting and stopping the MAVSDK server.

    Args:
        action (str): The action to perform.
        altitude (float, optional): The altitude for takeoff.
        parameters (dict, optional): Parameters to set on the drone.
        branch (str, optional): The branch name for code update.
    """
    logger.info("Starting to perform action...")

    if action == "update_code":
        await update_code(branch)
        return  # No need to proceed further

    # Retrieve Hardware ID
    global HW_ID
    HW_ID = read_hw_id()
    if HW_ID is None:
        logger.error("Hardware ID could not be determined. Exiting...")
        return

    # Read Drone Configuration
    drone_config = read_config()
    if not drone_config:
        logger.error("Drone configuration not found. Exiting...")
        return

    # Define MAVSDK Ports since it's on each system it's constant
    grpc_port = GRPC_PORT
    udp_port = UDP_PORT

    logger.info(f"gRPC Port: {grpc_port}, UDP Port: {udp_port}")

    # Start the MAVSDK server
    mavsdk_server = start_mavsdk_server(grpc_port, udp_port)

    # Initialize the MAVSDK drone system
    drone = System(mavsdk_server_address="localhost", port=grpc_port)
    logger.info("Attempting to connect to drone...")

    try:
        await drone.connect(system_address=f"udp://:{udp_port}")
    except Exception:
        logger.exception("Failed to connect to MAVSDK server")
        stop_mavsdk_server(mavsdk_server)
        return

    # Check connection state with a timeout
    logger.info("Checking connection state...")
    connected = False
    start_time = time.time()
    timeout = 10  # seconds
    try:
        async for state in drone.core.connection_state():
            if state.is_connected:
                logger.info(f"Drone connected on UDP Port: {udp_port} and gRPC Port: {grpc_port}")
                connected = True
                break
            if time.time() - start_time > timeout:
                logger.error("Timed out waiting for connection to drone.")
                break
    except Exception:
        logger.exception("Error while checking connection state")
        stop_mavsdk_server(mavsdk_server)
        return

    if not connected:
        logger.error("Could not establish a connection with the drone.")
        stop_mavsdk_server(mavsdk_server)
        return

    try:
        # Set parameters if provided
        if parameters:
            await set_parameters(drone, parameters)

        # Perform action
        if action == "takeoff":
            await takeoff(drone, altitude)
        elif action == "land":
            await land(drone)
        elif action == "hold":
            await hold(drone)
        elif action == "test":
            await test(drone)
        elif action == "reboot_fc":
            await reboot(drone, fc_flag=True, sys_flag=False)
        elif action == "reboot_sys":
            await reboot(drone, fc_flag=False, sys_flag=True)
        else:
            logger.error(f"Invalid action specified: '{action}'")
    except Exception:
        logger.exception(f"Error performing action '{action}'")
    finally:
        stop_mavsdk_server(mavsdk_server)
        logger.info("Action completed.")

# =======================
# Action Implementations
# =======================

async def takeoff(drone, altitude):
    """
    Executes the takeoff action.

    Args:
        drone (System): The MAVSDK drone system.
        altitude (float): The altitude to take off to.
    """
    led_controller = LEDController.get_instance()

    # Indicate takeoff initiation with yellow color
    led_controller.set_color(255, 255, 0)  # Yellow
    await asyncio.sleep(0.5)

    try:
        await drone.action.set_takeoff_altitude(float(altitude))
        await drone.action.arm()

        # Indicate arming with white color
        led_controller.set_color(255, 255, 255)  # White
        await asyncio.sleep(0.5)

        await drone.action.takeoff()

        # Indicate successful takeoff with green blinks
        for _ in range(3):
            led_controller.set_color(0, 255, 0)  # Green
            await asyncio.sleep(0.2)
            led_controller.turn_off()
            await asyncio.sleep(0.2)

        logger.info("Takeoff successful.")
    except Exception:
        logger.exception("Takeoff failed")
        # Indicate takeoff failure with red blinks
        for _ in range(3):
            led_controller.set_color(255, 0, 0)  # Red
            await asyncio.sleep(0.2)
            led_controller.turn_off()
            await asyncio.sleep(0.2)
    finally:
        # Turn off LEDs after feedback
        led_controller.turn_off()

async def land(drone):
    """
    Executes the land action.

    Args:
        drone (System): The MAVSDK drone system.
    """
    led_controller = LEDController.get_instance()

    # Indicate landing initiation with yellow color
    led_controller.set_color(255, 255, 0)  # Yellow
    await asyncio.sleep(0.5)

    try:
        await drone.action.hold()  # Switch to Hold mode
        await asyncio.sleep(1)  # Wait for a short period

        # Indicate landing in progress with blue slow pulse
        for _ in range(3):
            led_controller.set_color(0, 0, 255)  # Blue
            await asyncio.sleep(0.5)
            led_controller.turn_off()
            await asyncio.sleep(0.5)

        await drone.action.land()  # Execute land command

        # Indicate successful landing with green blinks
        for _ in range(3):
            led_controller.set_color(0, 255, 0)  # Green
            await asyncio.sleep(0.2)
            led_controller.turn_off()
            await asyncio.sleep(0.2)

        logger.info("Landing successful.")
    except Exception:
        logger.exception("Landing failed")
        # Indicate landing failure with red blinks
        for _ in range(3):
            led_controller.set_color(255, 0, 0)  # Red
            await asyncio.sleep(0.2)
            led_controller.turn_off()
            await asyncio.sleep(0.2)
    finally:
        # Turn off LEDs after feedback
        led_controller.turn_off()

async def hold(drone):
    """
    Executes the hold position action.

    Args:
        drone (System): The MAVSDK drone system.
    """
    led_controller = LEDController.get_instance()

    # Indicate hold command received with blue color
    led_controller.set_color(0, 0, 255)  # Blue
    await asyncio.sleep(0.5)

    try:
        await drone.action.hold()  # Switch to Hold mode

        # Indicate successful hold with a solid blue color
        led_controller.set_color(0, 0, 255)  # Solid Blue
        await asyncio.sleep(1)

        logger.info("Hold position successful.")
    except Exception:
        logger.exception("Hold failed")
        # Indicate hold failure with red blinks
        for _ in range(3):
            led_controller.set_color(255, 0, 0)  # Red
            await asyncio.sleep(0.2)
            led_controller.turn_off()
            await asyncio.sleep(0.2)
    finally:
        # Turn off LEDs after feedback
        led_controller.turn_off()

async def test(drone):
    """
    Executes the test action, which includes arming and disarming the drone with LED feedback.

    Args:
        drone (System): The MAVSDK drone system.
    """
    led_controller = LEDController.get_instance()

    try:
        # Step 1: Set LEDs to red before attempting to arm
        led_controller.set_color(255, 0, 0)  # Red
        await asyncio.sleep(1)  # Wait to show red color

        # Step 2: Arm the drone
        await drone.action.arm()

        # Indicate arming with white color
        led_controller.set_color(255, 255, 255)  # White
        await asyncio.sleep(1)  # Wait to show white color

        # Step 3: Change LED colors during the 3-second wait
        led_controller.set_color(0, 0, 255)  # Blue
        await asyncio.sleep(1)  # Wait to show blue color

        led_controller.set_color(0, 255, 0)  # Green
        await asyncio.sleep(1)  # Wait to show green color

        # Step 4: Disarm the drone
        await drone.action.disarm()
        logger.info("Test action successful.")

        # Step 5: Turn off LEDs after disarming
        led_controller.turn_off()

    except Exception:
        logger.exception("Test failed")
        # Indicate test failure with red blinks
        for _ in range(3):
            led_controller.set_color(255, 0, 0)  # Red
            await asyncio.sleep(0.2)
            led_controller.turn_off()
            await asyncio.sleep(0.2)
    finally:
        # Ensure LEDs are turned off
        led_controller.turn_off()

async def reboot(drone, fc_flag, sys_flag, force_reboot=True):
    """
    Reboots the flight controller, system, or both, with optional forced reboot.

    Args:
        drone (System): MAVSDK drone system.
        fc_flag (bool): Whether to reboot the flight controller.
        sys_flag (bool): Whether to reboot the Raspberry Pi system.
        force_reboot (bool): Whether to force a system reboot if the initial reboot fails.
    """
    led_controller = LEDController.get_instance()

    # Indicate reboot initiation with yellow color
    led_controller.set_color(255, 255, 0)  # Yellow
    await asyncio.sleep(0.5)  # Brief feedback

    if fc_flag:
        try:
            # Attempt to reboot the flight controller (drone)
            await drone.action.reboot()

            # Indicate success with green blinks
            for _ in range(3):
                led_controller.set_color(0, 255, 0)  # Green
                await asyncio.sleep(0.2)
                led_controller.turn_off()
                await asyncio.sleep(0.2)

            logger.info("Flight controller reboot successful.")
        except Exception:
            logger.exception("Flight controller reboot failed")

            # Indicate failure with red blinks
            for _ in range(3):
                led_controller.set_color(255, 0, 0)  # Red
                await asyncio.sleep(0.2)
                led_controller.turn_off()
                await asyncio.sleep(0.2)

            if force_reboot:
                logger.info("Force reboot enabled after flight controller reboot failure.")

    if sys_flag:
        try:
            # Attempt system (Raspberry Pi) reboot
            logger.info("Initiating system reboot...")
            led_controller.turn_off()

            # Use the reboot_system() function to reboot
            await reboot_system()
        except Exception:
            logger.exception("System reboot failed")

            if force_reboot:
                logger.info("Force reboot enabled, attempting forced system reboot.")
                try:
                    await reboot_system()
                except Exception:
                    logger.exception("Forced system reboot failed")

    # Final LED cleanup
    led_controller.turn_off()

async def reboot_system():
    """
    Reboots the Raspberry Pi using the D-Bus system interface.
    """
    try:
        # Use D-Bus to initiate the reboot without sudo
        process = await asyncio.create_subprocess_exec(
            'dbus-send', '--system', '--print-reply', '--dest=org.freedesktop.login1',
            '/org/freedesktop/login1', 'org.freedesktop.login1.Manager.Reboot', 'boolean:true',
            stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
        )
        stdout, stderr = await process.communicate()

        if process.returncode != 0:
            logger.error(f"System reboot via D-Bus failed: {stderr.decode().strip()}")
        else:
            logger.info("System reboot command executed successfully via D-Bus.")
    except Exception:
        logger.exception("System reboot failed")

async def update_code(branch=None):
    """
    Executes the update_https script to pull code from the requested repo.

    Args:
        branch (str, optional): The branch name to pull from. If None, default behavior is used.
    """
    led_controller = LEDController.get_instance()

    # Indicate update initiation with yellow color
    led_controller.set_color(255, 255, 0)  # Yellow
    await asyncio.sleep(0.5)

    try:
        # Build the command
        script_path = os.path.join('tools', 'update_https')
        command = [script_path]

        # Append branch name if provided
        if branch:
            command.append(branch)

        # Append '--sitl' if in simulation mode
        if Params.sim_mode:
            command.append('--sitl')

        logger.info(f"Executing update script: {' '.join(command)}")

        # Execute the bash script
        process = await asyncio.create_subprocess_exec(
            *command,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        stdout, stderr = await process.communicate()

        if process.returncode != 0:
            logger.error(f"Update script failed: {stderr.decode().strip()}")
            # Indicate failure with red blinks
            for _ in range(3):
                led_controller.set_color(255, 0, 0)  # Red
                await asyncio.sleep(0.2)
                led_controller.turn_off()
                await asyncio.sleep(0.2)
        else:
            logger.info(f"Update script executed successfully: {stdout.decode().strip()}")
            # Indicate success with green blinks
            for _ in range(3):
                led_controller.set_color(0, 255, 0)  # Green
                await asyncio.sleep(0.2)
                led_controller.turn_off()
                await asyncio.sleep(0.2)
    except Exception:
        logger.exception("Update code action failed")
        # Indicate failure with red blinks
        for _ in range(3):
            led_controller.set_color(255, 0, 0)  # Red
            await asyncio.sleep(0.2)
            led_controller.turn_off()
            await asyncio.sleep(0.2)
    finally:
        # Turn off LEDs after feedback
        led_controller.turn_off()

# =======================
# Entry Point
# =======================

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Perform actions with drones.")
    parser.add_argument('--action', type=str, required=True,
                        help='Action to perform: takeoff, land, hold, test, reboot_fc, reboot_sys, update_code')
    parser.add_argument('--altitude', type=float, default=10,
                        help='Altitude for takeoff')
    parser.add_argument('--param', action='append', nargs=2, metavar=('param_name', 'param_value'),
                        help='Set parameters in the form param_name param_value')
    parser.add_argument('--branch', type=str, help='Branch name for code update')

    args = parser.parse_args()

    # Convert parameter arguments to dictionary
    parameters = {param[0]: int(param[1]) for param in args.param} if args.param else None

    try:
        asyncio.run(perform_action(args.action, args.altitude, parameters, args.branch))
    except Exception:
        logger.exception("An unexpected error occurred")
    finally:
        logger.info("Operation completed.")