#!/usr/bin/python3
import os
import json
import argparse
import subprocess

# COLORS!
RESET = "\033[0m"
GRN = "\033[1;32m"
LIT_GRN = "\033[92m"
RED = "\033[1;31m"
WARN = "\033[33m"
CLEAR ="\033c"
BOLD="\033[1m"
RESET_BOLD="\033[2m"

SERVICE_NAME="voxl-joystick-server"
CONFIG_FILE = "/etc/modalai/voxl-joystick-server.conf"

config_header = """/**
 * voxl-joystick-server Configuration File
 */
"""

default_config = """
{
  "yaw": "R",
  "roll": "X",
  "pitch": "Y",
  "throttle": "Z",
  "arm_button": 0,
  "kill_button": -1,
  "acro_fm_button": -1,
  "manual_fm_button": 1,
  "altctl_fm_button": 2,
  "offboard_fm_button": -1,
  "posctl_fm_button": 3,
  "mission_fm_button": -1,
  "mission_rtl_button": -1,
  "center_zero_throttle": false,
  "update_rate_hz": 200,
  "gcs_mavlink_system_id": 255,
  "mavlink_target_system_id": 1,
  "transmitter_ip":  "10.223.0.101",
  "receiver_configs":       
    [
      {
          "ip":     "10.223.0.100",
          "port":   14550
      }
    ]
}
"""

# Parse args
def parse_args() -> dict:
    """
    Parse and validate args
    """
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument('-h', '--help',        help='Show this help message and exit', required=False, action='store_true')

    # Enable/Disable
    parser.add_argument('-d', '--disable', help='Disable the service', required=False, action='store_true')

    # Factory Enable
    parser.add_argument('-f', '--factory_enable', help='Factory enable the service', required=False, action='store_true')
    parser.add_argument('-I', '--factory_enable_fiber', help='Setup for fiber operation', required=False, action='store_true')

    # Wizard
    parser.add_argument('-w', '--wizard', help='Start interactive wizard', required=False, action='store_true')

    return vars(parser.parse_args())

# Print help message
def print_usage() -> None:
    """
    Print usage/help message
    """

    print("Description: Configure voxl-joystick-server settings.")
    print("")
    print("Usage: voxl-configure-joystick-server -h -d -f -I -w")
    print("")
    print("Optional Arguments:")
    print(" -h, --help                  Show this help message and exit")
    print(" -d, --disable               Disable the service")
    print(" -f, --factory_enable        Factory Enable the service")
    print(" -I, --factory_enable_fiber  Setup for fiber operation")
    print(" -w, --wizard                Start interactive wizard")
    print("")

# Check if a service is running
def service_running(service_name: str) -> bool:
    """
    Checks if a systemd service is running or not

    @service_name: name of service to check
    @return: True/False depending on service state
    """
    cmd = ['systemctl', 'is-active', service_name]
    result = subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    if result.returncode == 0:
        return True
    else:
        return False
    
# Check if a service is enabled
def service_enabled(service_name: str) -> bool:
    """
    Checks if a systemd service is enabled or not

    @service_name: name of service to check
    @return: True/False depending on service state
    """
    cmd = ['systemctl', 'is-enabled', service_name]
    result = subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    if result.returncode == 0:
        return True
    else:
        return False

# Restart a service
def restart_service(service_name: str) -> bool:
    """
    Restarts a systemd service

    @service_name: name of service to restart
    @return: True/False depending on if successful
    """
    print(f"Restarting {service_name}, this may take a few seconds...")
    result = subprocess.run(f"systemctl restart {service_name}".split(), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    if result.returncode != 0:
        print(f"Failed to restart {service_name}.")
        return False
    
    return True

# Enable and start a service
def enable_and_start_service(service_name: str) -> bool:
    """
    Enables and starts a systemd service

    @service_name: name of service to enable and starts
    @return: True/False depending on if successful
    """
    enable_cmd = ['systemctl', 'enable', service_name]
    start_cmd = ['systemctl', 'start', service_name]

    # Enable service
    result = subprocess.run(enable_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    if result.returncode != 0:
        print(f"Failed to enable {service_name}.")
        return False
    
    # Start service
    result = subprocess.run(start_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    if result.returncode != 0:
        print(f"Failed to start {service_name}.")
        return False
    
    return True

# Disable and stop a service
def disable_and_stop_service(service_name: str) -> bool:
    """
    Disables and stops a systemd service

    @service_name: name of service to disable and stop
    @return: True/False depending on if successful
    """
    disable_cmd = ['systemctl', 'disable', service_name]
    stop_cmd = ['systemctl', 'stop', service_name]

    # Disable service
    result = subprocess.run(disable_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    if result.returncode != 0:
        print(f"Failed to disable {service_name}.")
        return False
    
    # Stop service
    result = subprocess.run(stop_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    if result.returncode != 0:
        print(f"Failed to stop {service_name}.")
        return False
    
    return True

# Enable a service
def enable_service(service_name: str) -> bool:
    """
    Enables a systemd service

    @service_name: name of service to enable
    @return: True/False depending on if successful
    """
    enable_cmd = ['systemctl', 'enable', service_name]

    # Enable service
    result = subprocess.run(enable_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    if result.returncode != 0:
        print(f"Failed to enable {service_name}.")
        return False
    
    return True

# Disable a service
def disable_service(service_name: str) -> bool:
    """
    Disables a systemd service

    @service_name: name of service to disable
    @return: True/False depending on if successful
    """
    disable_cmd = ['systemctl', 'disable', service_name]

    # Disable service
    result = subprocess.run(disable_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    if result.returncode != 0:
        print(f"Failed to disable {service_name}.")
        return False
    
    return True

# Helper to load JSON file with optional comments
def load_json_with_comments(file_path: str) -> dict:
    """
    Load JSON from a file, filtering out lines containing '*' (comment lines).

    @file_path: Path to the JSON file
    @return: Parsed JSON as a dictionary
    """
    with open(file_path, 'r') as f:
        filtered_lines = [line for line in f if '*' not in line]
        return json.loads(''.join(filtered_lines))

# Helper to save JSON file with header comments
def save_json_with_header(file_path: str, config: dict) -> None:
    """
    Save JSON to a file with header comments preserved.

    @file_path: Path to the JSON file
    @config: Dictionary to save as JSON
    """
    with open(file_path, 'w') as f:
        f.write(config_header)
        json.dump(config, f, indent=4)
        f.write('\n')

# Read a file
def read_file(file: str) -> None:
    """
    Read the contents of a file given a file name.
    """
    if not os.path.exists(file):
        print(f"{BOLD}{RED}[ERROR]{RESET_BOLD} File {file} does not exist!")
        return 
    
    with open(file, 'r') as config:
        print(config.read().strip())

# Create a configuration file with default settings  
def write_default_config() -> None:
    """
    Write the default config file.
    """
    config = json.loads(default_config)
    print("Printing config")
    for key, value in config.items():
        print(f"\t{key}: {value}")

    # Ensure the config file exists
    if not os.path.exists(CONFIG_FILE):
        print(f"Config file '{CONFIG_FILE}' not found. Creating with default values.")
        with open(CONFIG_FILE, "w") as file:
            json.dump(default_config, file, indent=4)
        return 
    
    with open(CONFIG_FILE, 'w') as config_file:
        config_file.truncate(0)
        config_file.write(default_config)

# Create a configuration file with default settings
def write_default_config() -> None:
    """
    Write the default config file.
    """
    config = json.loads(default_config)
    print("Printing config")
    for key, value in config.items():
        print(f"\t{key}: {value}")
    
    # Ensure the config file exists
    if not os.path.exists(CONFIG_FILE):
        print(f"Config file '{CONFIG_FILE}' not found. Creating with default values.")
    
    with open(CONFIG_FILE, 'w') as config_file:
        json.dump(config, config_file, indent=4)

# update config file via arg and val
def update_config_param(arg: str, value, parent: str = ""):
    """
    Update a param in the config file.
    @arg: The field to be updated
    @value: The value for the field to be set to 
    @parent: Nested/Subfield that the given arg is in 
    """
    # Ensure the config file exists
    if not os.path.exists(CONFIG_FILE):
        print(f"Config file '{CONFIG_FILE}' not found. Creating with default values.")
        save_json_with_header(CONFIG_FILE, default_config)

    # Read the existing config file
    try:
        config = load_json_with_comments(CONFIG_FILE)
    except json.JSONDecodeError as e:
        raise ValueError(f"Error decoding JSON from config file: {e}")

    # Update the configuration
    if not parent and arg in config:
        config[arg] = value
    elif parent and parent in config: # Nested value
        for list_item in config[parent]:
            if arg in list_item:
                list_item[arg] = value
    else:
        raise KeyError(f"Key '{arg}' not found in the configuration file.")

    # Write the updated config back to the file
    save_json_with_header(CONFIG_FILE, config)

    print(f"{arg} -> {value}")

# Enable services and load default config file for SparkLAN modem
def factory_enable():
    """
    Enable voxl-joystick-server service and write default config file
    """
    enable_service(SERVICE_NAME)
    # restart_service(SERVICE_NAME) // Don't start service during voxl-configure-mpa

# Enable services and load default config file for fiber configuration
def factory_enable_fiber():
    """
    Enable voxl-joystick-server service and update config file for fiber use case
    """
    write_default_config()
    update_config_param("transmitter_ip", "10.0.0.2")
    update_config_param("ip", "10.0.0.10", "receiver_configs")
    update_config_param("port", 14550, "receiver_configs")
    update_config_param("altctl_fm_button", 1)
    update_config_param("posctl_fm_button", 2)
    update_config_param("offboard_fm_button", 3)
    update_config_param("mission_rtl_button", 15)
    enable_service(SERVICE_NAME)

# Handle updating VRX configuration file via wizard
def wizard_update_config() -> bool:
    """
    Update parameters in voxl-joystick-server configuration file via wizard
    """
    options = []
    cases = {}
    invalid = False

    while True:
        try:
            read_file(CONFIG_FILE)
            with open(CONFIG_FILE, 'r') as joystick_server_config:
                joystick_server_config_json = json.load(joystick_server_config)
                options = list(joystick_server_config_json.keys())
                options.insert(0,"Quit")
                options.insert(1,"Back")
                options.insert(2,"Restart Service")

            print(f"{LIT_GRN}Select a field to update:{RESET}")
            for index, option in enumerate(options):
                print(f"[{index + 1}] {option}")
                cases[index + 1] = option
    
            try:
                field = cases.get(int(input()))
                print(f"{CLEAR}", end="")
            except KeyboardInterrupt:
                print()
                exit(0)
            except:
                print(f"{WARN}Invalid answer provided, please select a valid option.{RESET}")
                continue
            
            if field == "Back": break
            if field == "Quit": exit(0)
            if field == "Restart Service": 
                restart_service(SERVICE_NAME)
                continue
            with open(CONFIG_FILE, 'w') as joystick_server_config:
                try:
                    if type(joystick_server_config_json[field]) == bool:
                        value = input(f"Enter a value for {field} (Current: bool, {joystick_server_config_json[field]}): ").lower()
                        if value == "false":  value = False
                        elif value == "true": value = True
                        else:
                            throw_exception = 1/0
                    elif type(joystick_server_config_json[field]) == int:
                        value = int(input(f"Enter a value for {field} (Current: int, {joystick_server_config_json[field]}): "))
                    elif type(joystick_server_config_json[field]) == str:
                        value = input(f"Enter a value for {field} (Current: string, {joystick_server_config_json[field]}): ")
                except KeyboardInterrupt:
                    print()
                    invalid = True
                    continue
                except:
                    print(f"{WARN}Invalid answer provided, please select a valid option.{RESET}")
                    invalid = True
                    continue

                print(f"{field}: {joystick_server_config_json[field]} -> {value}")
                joystick_server_config_json[field] = value
                json.dump(joystick_server_config_json, joystick_server_config, indent=4)
        except FileNotFoundError as e:
            print(f"{RED}{BOLD}[ERROR] Unable to locate configuration file, please run voxl-joystick-server at least once so that it generates a configuration file.{RESET}")
            print(str(e))
            return False
        except Exception as e:
            with open(CONFIG_FILE, 'w') as joystick_server_config: json.dump(joystick_server_config_json, joystick_server_config, indent=4)
            if not invalid: 
                print(f"{RED}{BOLD}[ERROR] An unexpected error occurred while uploading configuration. {str(e)}{RESET}")
                return False
            invalid = False
    return True

# Wizard for handling interactive mode
def wizard() -> int:
    """
    Handle interactive mode, guide user through option menus and handle input.
    Wizard options:
    [1] Quit wizard
    [2] Restart service     -> Restart service after making changes
    [3] Settings            -> Show current settings
    """

    options = ["Quit", "Restart Service", "Settings"]
    cases = {}
    while True:
        print(f"{LIT_GRN}Select an option:{RESET}")
        for index, option in enumerate(options):
            print(f"[{index + 1}] {option}")
            cases[index + 1] = option

        try:
            ans = cases.get(int(input()))
            print(f"{CLEAR}", end="")
        except KeyboardInterrupt:
            print()
            exit(0)
        except:
            print(f"{WARN}Invalid answer provided, please select a valid option.{RESET}")
            continue
            
        # Handle request
        if ans == "Quit":               return 0
        elif ans == "Restart Service":  restart_service(SERVICE_NAME)
        elif ans == "Settings":         wizard_update_config()
        else:
            print(f"{WARN}Invalid answer provided, please select a valid option.{RESET}")
            continue

if __name__ == "__main__":
    args = parse_args()
    
    # Print help message if no args provided
    if all(value is None or value is False for value in args.values()) or args["help"]:	
        print_usage()
        exit(0)

    # Handle disable early (before config file checks)
    if args["disable"]:
        disable_and_stop_service(SERVICE_NAME)
        exit(0)

    if not os.path.exists(CONFIG_FILE) and not args["factory_enable"]:
        print(f"{WARN}[WARNING] Unable to locate configuration file! Generating default configuration file now.{RESET}")
        write_default_config()

    if os.path.exists(CONFIG_FILE) and not os.path.getsize(CONFIG_FILE)>0: 
        print(f"{WARN}[WARNING] Configuration file may be empty (file size: {os.path.getsize(CONFIG_FILE)})! Generating default configuration file now.{RESET}")
        write_default_config()

    # Check for use wizard (no args provided)
    if args["wizard"]: exit(wizard())

    # Handle enable/disable and setup
    if args["factory_enable"]: factory_enable()
    if args["factory_enable_fiber"]: factory_enable_fiber()
    