#!/usr/bin/env python3
import os
import errno
import subprocess
from subprocess import check_output, STDOUT, CalledProcessError
import time
import argparse
import configparser
import shutil
import sys
import re


ENVIRONMENT_VARIABLES_TO_PRESERVE = ["PYTHONPATH", "TEST", "TERM"]


def run(command, verbose=False):
    """
    Utility function to execute given command and return its output.
    """

    try:
        output = check_output(command, shell=True, env=os.environ, stderr=STDOUT, encoding="utf-8")
        return output.strip("\n") if not verbose else (output.strip("\n"), 0, None)
    except CalledProcessError as error:
        return error.output if not verbose else (error.output, error.returncode, error)


def get_environment_dictionary():
    """
    Targetting the gnome-session-binary process to get all environment variables.
    Returns environment as dictionary.
    """

    environment_dictionary = {}

    gnome_session_binary_pid = None

    try:
        gnome_session_binary_pid = run("pgrep -u $USER gnome-session").split("\n")[0]
    except Exception as error:
        print(f"headless: Failed to retrieve gnome-session-binary pid from pgrep: {error}")

    if gnome_session_binary_pid:
        environment_process_path = f"/proc/{gnome_session_binary_pid}/environ"

        # Verify that the environ file can be opened and load environment variables to dictionary.
        try:
            with open(environment_process_path, "r") as environ_file:
                for item in environ_file.read().split("\x00"):
                    if "=" in item:
                        key, value = item.split("=", 1)
                        environment_dictionary[key] = value
        except IOError as error:
            print(f"headless: Environment file manipulation failed on: '{error}'")

        # Preserving wanted environment variables
        for environment_variable in ENVIRONMENT_VARIABLES_TO_PRESERVE:
            if environment_variable in os.environ:
                environment_dictionary[environment_variable] = os.environ[environment_variable]

        # Set TERM as xterm
        print("headless: Setting environment variable TERM as 'xterm'") # REMOVE OR NOT?
        environment_dictionary["TERM"] = "xterm"

        return environment_dictionary

    print("headless: Can't find our environment!")
    return None


def is_binary_existing_and_executable(path):
    """
    Test if given binary file exists.
    """

    if (path.startswith(os.path.sep) or
            path.startswith(os.path.join(".", "")) or
            path.startswith(os.path.join("..", ""))):
        if not os.path.exists(path):
            raise IOError(errno.ENOENT, "No such file", path)

        if not os.access(path, os.X_OK):
            raise IOError(errno.ENOEXEC, "Permission denied", path)

    return True

# remove?
def troubleshoot():
    """
    Test parts of the system for correct configuration.
    """

    # Print running at-spi processes that are required.
    print("Fetching AT-SPI processes")
    at_spi_command = "ps -o pid,tname,command | grep at-spi | grep -v grep"
    at_spi_result = run(at_spi_command)
    print("".join((
        "\nExpecting running processes to be\n",
        "'/usr/libexec/at-spi-bus-launcher'\n",
        "'/usr/libexec/at-spi2-registryd'\n"
    )))
    print(f"Actual running processes:\n'{at_spi_result}'")
    
    # Important variables
    display = os.getenv("DISPLAY")
    xauthority = os.getenv("XAUTHORITY")
    dbus_session_bus_address = os.getenv("DBUS_SESSION_BUS_ADDRESS")
    print("".join((
        f"\nDISPLAY = '{display}'"
        f"\nXAUTHORITY = '{xauthority}'"
        f"\nDBUS_SESSION_BUS_ADDRESS = '{dbus_session_bus_address}'\n"
    )))


def parse():
    parser = argparse.ArgumentParser(prog="$ qecore-headless",
                                     description="Adjusted headless script.",
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("script",
                        nargs="?",
                        default="bash",
                        help="Script to be executed, if not provided 'bash' will be used.")
    parser.add_argument("--session-type",
                        required=False,
                        choices=("xorg", "wayland"),
                        help="Choose which session type will be used.")
    parser.add_argument("--session-desktop",
                        required=False,
                        choices=("gnome", "gnome-classic", "gnome-classic-wayland"),
                        help="Choose which session desktop will be used.")
    parser.add_argument("--dont-start",
                        required=False,
                        action="store_true",
                        help="Use the system as is. Does not have to be under display manager")
    parser.add_argument("--dont-kill",
                        required=False,
                        action="store_true",
                        help="Do not kill the session when script exits.")
    parser.add_argument("--restart",
                        required=False,
                        action="store_true",
                        help="Restart previously running display manager session before script execution.")
    parser.add_argument("--keep",
                        required=False,
                        help="Number of tests to keep gdm running")
    parser.add_argument("--disable-a11y",
                        required=False,
                        action="store_true",
                        help="Disable accessibility technologies on script (not session) exit.")
    parser.add_argument("--force",
                        required=False,
                        action="store_true",
                        help="Will check if the wanted protocol was used. Exit upon fail.")
    parser.add_argument("--debug",
                        required=False,
                        action="store_true",
                        help="Will print debug messages to the file.")
    return parser.parse_args()


class DisplayManager:
    def __init__(self,
                 session_type=None,
                 session_desktop=None,
                 enable_start=True,
                 enable_stop=True,
                 gdm_restart=False):
        self.enable_start = enable_start
        self.enable_stop = enable_stop
        self.gdm_restart = gdm_restart

        self.display_manager = "gdm"
        self.session_type = session_type # xorg, wayland, None -> respect setting of the system
        self.session_desktop = session_desktop # gnome gnome-classic, None -> respect setting of the system
        self.session_started_indicator = "/usr/libexec/gnome-session-binary"
        
        self.user = run("whoami")
        self.user_id = run(f"id -u {self.user}")

        self.config_file = "/etc/gdm/custom.conf"
        self.temporary_config_file = f"/tmp/{os.path.basename(self.config_file)}"

        self.restart_needed = False


    def restore_config(self): # Not used, but implemented if needed.
        # Restore configuration file.
        shutil.copy(self.config_file, self.temporary_config_file)

        config_parser = configparser.ConfigParser()
        config_parser.optionxform = str
        config_parser.read(self.temporary_config_file)

        config_parser.remove_option("daemon", "AutomaticLoginEnable")
        config_parser.remove_option("daemon", "AutomaticLogin")
        config_parser.remove_option("daemon", "WaylandEnable")

        with open(self.temporary_config_file, "w") as _file:
            config_parser.write(_file)


    def handling_config_setup(self):
        # Handling config setup.
        run(f"cp -f {self.config_file} {self.temporary_config_file}")

        config_parser = configparser.ConfigParser()
        config_parser.optionxform = str
        config_parser.read(self.temporary_config_file)

        if not config_parser.has_section("daemon"): # Section does not exist.
            config_parser.add_section("daemon")

        config_parser.set("daemon", "AutomaticLoginEnable", "true")
        config_parser.set("daemon", "AutomaticLogin", self.user)

        if self.session_type == "xorg":
            config_parser.set("daemon", "WaylandEnable", "false")
        elif self.session_type == "wayland":
            config_parser.set("daemon", "WaylandEnable", "true")
        elif self.session_type is None: # Respecting system setting, get the session that is to be started.
            if "WaylandEnable" in config_parser.options("daemon"):
                self.session_type = "wayland" if config_parser.get("daemon", "WaylandEnable") == "true" else "xorg"
            else: # "WaylandEnable" not in config_parser.options("daemon")
                self.session_type = "xorg"
        else:
            print("headless: This is not acceptable session type. Fallback to the 'xorg' session type")
            print("headless: Acceptable names for --session-type: ['xorg', 'wayland']")
            self.session_type = "xorg"
            config_parser.set("daemon", "WaylandEnable", "false")

        with open(self.temporary_config_file, "w") as _file:
            config_parser.write(_file)

        if not os.path.isfile(self.temporary_config_file):
            print("headless: Temporary config file was not found, waiting a bit...")
            time.sleep(1)

        run(f"sudo mv -f {self.temporary_config_file} {self.config_file}")
        run(f"sudo rm -f {self.temporary_config_file}")


    def handling_account_setup(self):
        # Handling account setup.

        # Get all defined desktop file name for xorg and wayland.
        acceptable_x_desktop_names = run("ls /usr/share/xsessions").split("\n")
        acceptable_wayland_desktop_names = run("ls /usr/share/wayland-sessions").split("\n")
        acceptable_desktop_file_names = None

        # Get acceptable desktop file names for xorg.
        if self.session_type == "xorg":
            acceptable_desktop_file_names = [x.strip("desktop").strip(".") for x in \
                acceptable_x_desktop_names if x]

        # Get aceptable desktop file names for wayland.
        elif self.session_type == "wayland":
            acceptable_desktop_file_names = [x.strip("desktop").strip(".") for x in \
                acceptable_wayland_desktop_names if x]

        # Get initial values that we work with.
        interface = f"org.freedesktop.Accounts /org/freedesktop/Accounts/User{self.user_id} org.freedesktop.Accounts.User"
        saved_session_desktop = run(f"busctl get-property {interface} Session")
        saved_xsession_desktop = run(f"busctl get-property {interface} XSession")

        # Handling result from get-property. Making sure they are equal.
        saved_session_desktop = saved_session_desktop[3:-1]
        saved_xsession_desktop = saved_xsession_desktop[3:-1]
        if saved_session_desktop == "":
            run(f"busctl call {interface} SetSession 's' '{saved_xsession_desktop}'")
            saved_session_desktop = saved_xsession_desktop
        elif saved_xsession_desktop == "":
            run(f"busctl call {interface} SetXSession 's' '{saved_session_desktop}'")
            saved_xsession_desktop = saved_session_desktop

        # Chosen desktop differs from current one.
        if self.session_desktop not in (saved_session_desktop, None) and \
            self.session_desktop in acceptable_desktop_file_names:
            print(f"headless: Changing desktop '{saved_session_desktop}' -> '{self.session_desktop}'")
            run(f"busctl call {interface} SetSession 's' '{self.session_desktop}'")
            run(f"busctl call {interface} SetXSession 's' '{self.session_desktop}'")
            self.restart_needed = True

        # Chosing desktop not found in acceptable desktop file names.
        elif self.session_desktop is not None and \
            self.session_desktop not in acceptable_desktop_file_names:
            print("headless: This is not acceptable session desktop name. Fallback to the 'gnome' session desktop")
            print(f"headless: Acceptable names for '{self.session_type}': {acceptable_desktop_file_names}")
            run(f"busctl call {interface} SetSession 's' 'gnome'")
            run(f"busctl call {interface} SetXSession 's' 'gnome'")
            self.restart_needed = True

        if self.restart_needed:
            print("headless: Restart required")
            run("sudo systemctl restart accounts-daemon")
            run("sudo systemctl restart systemd-logind")


    def start_display_manager(self):
        """
        Starting the display manager - gdm.
        """

        # Stop gdm only if requested by user or required by config change, continue using active one othewise.
        if self.gdm_restart or self.restart_needed:
            run("sudo systemctl stop gdm")

            # First check, loginctl.
            preexisting_session = run(f"sudo loginctl | grep {self.user} | grep seat0")
            if preexisting_session:
                preexisting_session_number = preexisting_session.strip().strip(" ")[0]
                run(f"sudo loginctl kill-session --signal=9 {preexisting_session_number}")

            # Second check, loginctl.
            preexisting_session = run(f"sudo loginctl | grep {self.user} | grep seat0")
            if preexisting_session:
                preexisting_session_number = preexisting_session.strip().strip(" ")[0]
                run(f"sudo loginctl kill-session --signal=9 {preexisting_session_number}")

        # Start gdm if gdm is not active already.
        is_display_manager_active = run("systemctl is-active gdm")
        if is_display_manager_active != "active":
            print("headless: Starting Display Manager")
            run("sudo systemctl start gdm")
            time.sleep(4)
        
        # But the session must be running.
        if not self.wait_until_process_is_running(self.session_started_indicator):
            print("headless: Running session indicator not detected running - restart required")
            print("headless: Attempt to restore headless - stopping gdm")
            run("sudo systemctl stop gdm")
            print("headless: Attempt to restore headless - starting gdm")
            run("sudo systemctl start gdm")
            time.sleep(4)


    def stop_display_manager(self):
        """
        Stopping the display manager - gdm.
        """

        print("headless: Stopping Display Manager")
        run("sudo systemctl stop gdm")
        self.wait_until_process_is_not_running(self.session_started_indicator)

        # First check, loginctl.
        if self.is_process_running(self.session_started_indicator):
            still_open_session = run(f"sudo loginctl | grep {self.user} | grep seat0", verbose=True)
            if still_open_session[1] == 0: # check return code
                still_open_session_number = still_open_session[0].strip().strip(" ")
                run(f"sudo loginctl terminate-session {still_open_session_number}")

        # Second check, process.
        if self.is_process_running("xorg") or self.is_process_running("wayland"):
            still_open_session = run(f"sudo loginctl | grep {self.user} | grep seat0", verbose=True)
            if still_open_session[1] == 0: # Check return code.
                still_open_session_number = still_open_session[0].strip().strip(" ")
                run(f"sudo loginctl kill-session --signal=9 {still_open_session_number}")


    @staticmethod
    def is_process_running(process_to_find):
        active_processes = run("ps axw").split("\n")
        for active_process in active_processes:
            if re.search(process_to_find, str(active_process).lower()):
                return True
        return False


    # wait until process IS running with hard limit of 30 seconds
    def wait_until_process_is_running(self, process_to_find):
        for _ in range(60):
            if not self.is_process_running(process_to_find):
                time.sleep(0.5)
            else:
                return True
        return False

    # wait until process IS NOT running with hard limit of 30 seconds
    def wait_until_process_is_not_running(self, process_to_find):
        for _ in range(60):
            if self.is_process_running(process_to_find):
                time.sleep(0.5)
            else:
                break


class Headless:
    def __init__(self):
        self.display_manager_control = None
        self.environment_control = None
        self.script_control = None

        self.arguments = None
        self.script_as_list_of_arguments = None

        self.enable_start = True
        self.enable_stop = True

        self.gdm_restart = False

        self.disable_accessibility_on_script_exit = None

        self.force = None
        self.session_type = None
        self.session_desktop = None

        self.user_script_process = None
        self.user_script_exit_code = None


    @staticmethod
    def set_accessibility_to(enable_accessibility):
        """
        Using simple gsettings command to enable or disable toolkit-accessibility.
        """

        set_accessibility_value = "true" if enable_accessibility else "false"

        gsetting_get_command = " ".join((
            "dbus-run-session gsettings get",
            "org.gnome.desktop.interface",
            "toolkit-accessibility"
        ))

        gsetting_set_command = " ".join((
            "dbus-run-session gsettings set",
            "org.gnome.desktop.interface",
            f"toolkit-accessibility {set_accessibility_value}"
        ))

        accessibility_value = run(gsetting_get_command)
        if accessibility_value != set_accessibility_value:
            print(f"headless: Changing a11y value from '{accessibility_value}' to '{set_accessibility_value}'")
            run(gsetting_set_command)


    @staticmethod
    def adjust_gsettings_values():
        """
        Using simple gsettings command to adjust values of delay, repeat and repeat-interval.
        """

        user = run("whoami")
        scheme = "org.gnome.desktop.peripherals.keyboard"

        gsetting_get_delay_command = f"sudo -Hu {user} gsettings get {scheme} delay"
        gsetting_set_delay_command = f"sudo -Hu {user} gsettings set {scheme} delay 'uint32 500'"

        gsetting_get_repeat_command = f"sudo -Hu {user} gsettings get {scheme} repeat"
        gsetting_set_repeat_command = f"sudo -Hu {user} gsettings set {scheme} repeat true"

        gsetting_get_repeat_interval_command = f"sudo -Hu {user} gsettings get {scheme} repeat-interval"
        gsetting_set_repeat_interval_command = f"sudo -Hu {user} gsettings set {scheme} repeat-interval 'uint32 30'"

        delay_result = run(gsetting_get_delay_command)
        repeat_result = run(gsetting_get_repeat_command)
        repeat_interval_result = run(gsetting_get_repeat_interval_command)

        if delay_result != "uint32 500":
            run(gsetting_set_delay_command)
            print(f"Value of gsettings delay was '{delay_result}' - changing to 'uint32 500'")

        if repeat_result != "true":
            run(gsetting_set_repeat_command)
            print(f"Value of gsettings repeat was '{repeat_result}' - changing to 'true'")

        if repeat_interval_result != "uint32 30":
            run(gsetting_set_repeat_interval_command)
            print(f"Value of gsettings repeat-interval was "f"'{repeat_interval_result}' - changing to 'uint32 30'")


    def handle_keep_logic(self, keep_value):
        keep_file = "/tmp/qecore_keep"
        keep_from_argument = int(keep_value)

        try:
            with open(keep_file, "r") as _file:
                keep_from_file = int(_file.read())
        except Exception:
            keep_from_file = 1

        self.enable_stop = False

        if keep_from_file == 1:
            self.gdm_restart = True

        if keep_from_file >= keep_from_argument:
            self.enable_stop = True
            keep_from_file = 0

        with open(keep_file, "w") as _file:
            _file.write(f"{keep_from_file + 1}")


    def handle_arguments(self):
        """
        Makes all neccessary steps for arguments passed along the headless script.
        """

        # Parse arguments of headless
        self.arguments = parse()

        # Parse arguments of given script
        self.script_as_list_of_arguments = self.arguments.script.split()

        # Handle headless debug variable
        if self.arguments.debug:
            os.environ["DOGTAIL_DEBUG"] = "true"

        # Handle keep argument, check value of /tmp/qecore_keep
        if self.arguments.keep:
            self.handle_keep_logic(self.arguments.keep)

        # Handle headless don't start variable
        if self.arguments.dont_start:
            self.enable_start = False

        # Handle headless don't kill variable
        if self.arguments.dont_kill:
            self.enable_stop = False

        # Handle headless restart variable
        if self.arguments.restart:
            self.gdm_restart = True

        # Handle headless disable a11y variable
        if self.arguments.disable_a11y:
            self.disable_accessibility_on_script_exit = True

        # Handle headless force variable
        if self.arguments.force:
            self.force = True

        # Handle session type variable
        if self.arguments.session_type:
            self.session_type = self.arguments.session_type

        # Handle session desktop variable
        if self.arguments.session_desktop:
            self.session_desktop = self.arguments.session_desktop


    def check_what_desktop_and_type_is_running(self):
        """
        Retrieve information about running process and prints it before user script start.
        """

        if not self.enable_start:
            return

        current_type = "xorg" if os.environ["XDG_SESSION_TYPE"] == "x11" else os.environ["XDG_SESSION_TYPE"]
        current_desktop = os.environ["XDG_SESSION_DESKTOP"]
        print(f"headless: Running '{current_type}' with desktop '{current_desktop}'")


    def verify_that_correct_session_was_started(self):
        """
        Verifies that correct session type as started, terminate on mismatch.
        """

        if not self.enable_start:
            return

        current_type = "xorg" if os.environ["XDG_SESSION_TYPE"] == "x11" else os.environ["XDG_SESSION_TYPE"]
        if self.display_manager_control.session_type and\
            self.display_manager_control.session_type != current_type:
            print("".join((
                "headless: Script requires session of type: ",
                f"'{self.display_manager_control.session_type}'\n",
                "headless: Script was started under session of type: ",
                f"'{current_type}'\n"
            )))
            print("Exitting the headless script.")
            sys.exit(1)

        current_desktop = os.environ["XDG_SESSION_DESKTOP"]
        if self.display_manager_control.session_desktop and\
            self.display_manager_control.session_desktop != current_desktop:
            print("".join((
                "headless: Script requires session with desktop: ",
                f"'{self.display_manager_control.session_desktop}'\n",
                "headless: Script was started under session with desktop: ",
                f"'{current_desktop}'\n"
            )))
            print("Exitting the headless script.")
            sys.exit(1)


    def execute(self):
        """
        Makes all neccessary preparations for the system to start gdm and execute user script.
        """

        # Arguments handling.
        self.handle_arguments()

        # Display manager setup and handling.
        self.display_manager_control = DisplayManager(session_type=self.session_type,
                                                      session_desktop=self.session_desktop,
                                                      enable_start=self.enable_start,
                                                      enable_stop=self.enable_stop,
                                                      gdm_restart=self.gdm_restart)
        self.display_manager_control.handling_config_setup()
        self.display_manager_control.handling_account_setup()

        # Required setup that needs to go through in its entirety or it will fail.
        if self.enable_start or self.display_manager_control.restart_needed:
            self.display_manager_control.start_display_manager()

        # Environment handling.
        environment_dictionary = get_environment_dictionary()
        if environment_dictionary:
            os.environ = environment_dictionary
        else:
            print(f"headless: Display Manager Status is: '{run('systemctl status gdm')}'")
            self.display_manager_control.stop_display_manager()

            troubleshoot()
            sys.exit(1)

        # Force xorg/wayland setting - terminate upon error.
        if self.force:
            self.verify_that_correct_session_was_started()

        # Check xorg/wayland setting - print what the test will run under.
        self.check_what_desktop_and_type_is_running()

        # Accessibility - has to be started after gdm/dbus is running.
        self.set_accessibility_to(True)

        # Make a gsettings that sometimes get set too high and is uncomfortable to work with.
        self.adjust_gsettings_values()

        # User script handling.
        if is_binary_existing_and_executable(self.script_as_list_of_arguments[0]):
            self.user_script_process = subprocess.Popen(self.script_as_list_of_arguments, env=os.environ)
            print(f"headless: Started the script with PID {self.user_script_process.pid}")

            self.user_script_exit_code = self.user_script_process.wait()
            print(f"headless: The user script finished with return code {self.user_script_exit_code}")

        # Disable accessibility upon script exit.
        if self.disable_accessibility_on_script_exit:
            self.set_accessibility_to(False)

        # Stop display manager unless user specifies otherwise.
        if self.enable_stop:
            self.display_manager_control.stop_display_manager()


def main():
    # Headless is not supposed to be run under user root.
    if os.geteuid() == 0:
        print("headless: Script is not meant to be run under the root user.")
        print("Exitting the headless script.")
        sys.exit(1)

    headless = Headless()
    headless.execute()

    sys.exit(headless.user_script_exit_code)


if __name__ == "__main__":
    main()
