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


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

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


ENVIRONMENT_DICTIONARY = {}
ENVIRONMENT_VARIABLES_TO_PRESERVE = ["PYTHONPATH", "TEST", "TERM"]


def _get_environment_as_dictionary(path_to_process):
    """
    Get environment variables from process and save to the global dictionary.
    Return True or False if the globel environment variable was successfuly modified.
    """

    def is_session_process(path_to_process):
        pid = path_to_process.split("/")[2]
        if pid == "self" or pid == str(os.getpid()):
            return False
        return True

    # Verify path to process is indeed a process
    if is_session_process(path_to_process):
        file_name = path_to_process + "environ"
    else:
        return False

    # Verify the environ file can be opened
    try:
        path_to_environment = open(file_name, "r").read()
    except IOError:
        return False

    # Loading environment to dictionary
    environment_items = path_to_environment.split("\x00")
    for item in environment_items:
        if "=" in item:
            key, value = item.split("=", 1)
            ENVIRONMENT_DICTIONARY[key] = value

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

    return True


def get_environment_dictionary():
    """
    Iterates over the proc files and looks for environment to use.
    Targetting the org.gnome.Shell process to get all environment variables.
    Returns environment as dictionary.
    """

    def get_gs_pid():
        out = run("ps o pid,command,user").split("\n")
        for line in out:
            if "gnome-shell" in line and "test" in line:
                pid = line.strip(" ").split(" ")[0]
                return int(pid)

    process_environ_to_use = get_gs_pid()
    if process_environ_to_use:
        _get_environment_as_dictionary(f"/proc/{process_environ_to_use}/")
        print(f"headless: Setting environment variable TERM as 'xterm'")
        ENVIRONMENT_DICTIONARY["TERM"] = "xterm"
        return ENVIRONMENT_DICTIONARY

    is_display_manager_active = run("systemctl is-active gdm").strip("\n")
    print("\n".join((
        f"headless: Display Manager Status is '{is_display_manager_active}'",
        "headless: Can't find our environment!",
        "headless: Stopping Display Manager"
    )))
    subprocess.Popen(("sudo systemctl stop gdm").split()).wait()
    sys.exit(1)



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


def start_script_and_return_process(list_of_arguments):
    is_binary_existing_and_executable(list_of_arguments[0])
    return subprocess.Popen(list_of_arguments, env=os.environ)


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_binary = "/usr/bin/gnome-shell"
        self.user = run("whoami").strip("\n")

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

        self.account_file = f"/var/lib/AccountsService/users/{self.user}"
        self.temporary_account_file = f"/tmp/{os.path.basename(self.account_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
        subprocess.Popen(f"cp -f {self.config_file} {self.temporary_config_file}", shell=True).wait()

        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)

        subprocess.Popen(\
            f"sudo mv -f {self.temporary_config_file} {self.config_file}", shell=True).wait()
        subprocess.Popen(\
            f"sudo rm -f {self.temporary_config_file}", shell=True).wait()


    def handling_account_setup(self):
        # Handling account setup
        subprocess.Popen(f"cp -f {self.account_file} {self.temporary_account_file}", shell=True).wait()

        account_config_parser = configparser.ConfigParser()
        account_config_parser.optionxform = str
        account_config_parser.read(self.temporary_account_file)

        acceptable_x_desktop_names = [x for x in \
            run("ls /usr/share/xsessions").split("\n")]
        acceptable_wayland_desktop_names = [x for x in \
            run("ls /usr/share/wayland-sessions").split("\n")]

        if self.session_type == "xorg":
            acceptable_desktop_file_names = [x.strip("desktop").strip(".") for x in \
                acceptable_x_desktop_names if x]
        elif self.session_type == "wayland":
            acceptable_desktop_file_names = [x.strip("desktop").strip(".") for x in \
                acceptable_wayland_desktop_names if x]

        if not account_config_parser.has_section("User"): # section does not exist
            account_config_parser.add_section("User")
            account_config_parser.set("User", "Session", "gnome")
            account_config_parser.set("User", "SystemAccount", "false")

        elif account_config_parser.has_section("User"): # section does exist
            if "Session" in account_config_parser.options("User"): # option does exist
                saved_session_desktop = account_config_parser.get("User", "Session")

            elif "Session" not in account_config_parser.options("User"): # option does not exist
                self.session_desktop = self.session_desktop if self.session_desktop else "gnome"
                account_config_parser.set("User", "Session", self.session_desktop)
                saved_session_desktop = account_config_parser.get("User", "Session")
                self.restart_needed = True

            if self.session_desktop not in (saved_session_desktop, None) and \
                self.session_desktop in acceptable_desktop_file_names: # chosen desktop differs from actual
                print(f"headless: Changing desktop '{saved_session_desktop}' -> '{self.session_desktop}'")
                account_config_parser.set("User", "Session", self.session_desktop)
                account_config_parser.set("User", "SystemAccount", "false")
                self.restart_needed = True

            elif self.session_desktop is not None and \
                self.session_desktop not in acceptable_desktop_file_names:  # chosen desktop not found
                print(" ".join((
                    "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}")
                account_config_parser.set("User", "Session", "gnome")
                self.restart_needed = True


        if self.restart_needed:
            print("headless: Restart required")
            with open(self.temporary_account_file, "w") as _file:
                account_config_parser.write(_file)

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

            subprocess.Popen(f"sudo mv -f {self.temporary_account_file} {self.account_file}",\
                             shell=True).wait()
            subprocess.Popen("sudo systemctl restart accounts-daemon", shell=True).wait()
            subprocess.Popen("sudo systemctl restart systemd-logind", shell=True).wait()
            subprocess.Popen(f"sudo rm -f {self.temporary_account_file}", shell=True).wait()


    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:
            subprocess.Popen(("sudo systemctl stop gdm").split()).wait()

            # 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]
                subprocess.Popen(" ".join((
                    f"sudo loginctl kill-session",
                    f"--signal=9 {preexisting_session_number}"
                )), shell=True).wait()

            # 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]
                subprocess.Popen(" ".join((
                    f"sudo loginctl kill-session",
                    f"--signal=9 {preexisting_session_number}"
                )), shell=True).wait()

        is_display_manager_active = run("systemctl is-active gdm").strip("\n")
        if is_display_manager_active != "active":
            print("headless: Starting Display Manager")
            subprocess.Popen(("sudo systemctl start gdm").split()).wait()
            self.wait_until_process_is_running(self.session_binary)
            time.sleep(4)


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

        print("headless: Stopping Display Manager")
        subprocess.Popen(("sudo systemctl stop gdm").split()).wait()

        self.wait_until_process_is_not_running(self.session_binary)

        # first check, loginctl
        if self.is_process_running(self.session_binary):
            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(" ")
                subprocess.Popen(" ".join((
                    f"sudo loginctl terminate-session",
                    f"{still_open_session_number}"
                )), shell=True).wait()

        # 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(" ")
                subprocess.Popen(" ".join((
                    f"sudo loginctl kill-session --signal=9",
                    f"{still_open_session_number}"
                )), shell=True).wait()


    @staticmethod
    def is_process_running(process_to_find):
        active_processes = subprocess.Popen(["ps", "axw"], stdout=subprocess.PIPE)
        for active_process in active_processes.stdout:
            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): # gnome-shell
                time.sleep(0.5)
            else:
                break

    # 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): # gnome-shell
                time.sleep(0.5)
            else:
                break


def parse():
    import argparse
    parser = argparse.ArgumentParser(prog="$ qecore-headless",
                                     description="Adjusted headless script.",
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("script",
                        help="Script to be executed")
    parser.add_argument("--session-type",
                        required=False,
                        help="xorg wayland")
    parser.add_argument("--session-desktop",
                        required=False,
                        help="gnome gnome-classic gnome-classic-wayland")
    parser.add_argument("--dont-start",
                        action="store_true",
                        help="Use the system as is. Does not have to be under display manager")
    parser.add_argument("--dont-kill",
                        action="store_true",
                        help="Do not kill the session when script exits.")
    parser.add_argument("--restart",
                        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",
                        action="store_true",
                        help="Disable accessibility technologies on script (not session) exit.")
    parser.add_argument("--force",
                        action="store_true",
                        help="Will check if the wanted protocol was used. Exit upon fail.")
    parser.add_argument("--debug",
                        action="store_true",
                        help="Will print debug messages to the file.")
    return parser.parse_args()


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


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

        current_a11y_value = run("dbus-launch gsettings get org.gnome.desktop.interface toolkit-accessibility").strip("\n")
        set_a11y_value = "true" if enable_accessibility else "false"

        if current_a11y_value != set_a11y_value:
            print(f"headless: Changing a11y value from '{current_a11y_value}' to '{set_a11y_value}'")
            subprocess.Popen(
                f"dbus-launch gsettings set org.gnome.desktop.interface toolkit-accessibility {set_a11y_value}",
                shell=True, env=os.environ)


    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:
            keep_from_argument = int(self.arguments.keep)
            keep_file = "/tmp/qecore_keep"

            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}")

        # 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 = 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 = os.environ["XDG_SESSION_TYPE"]
        if self.display_manager_control.session_type != current_type:
            print("".join((
                f"headless: Script requires session of type: ",
                f"'{self.display_manager_control.session_type}'\n",
                f"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 != current_desktop:
            print("".join((
                f"headless: Script requires session with desktop: ",
                f"'{self.display_manager_control.session_desktop}'\n",
                f"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
        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
        os.environ = get_environment_dictionary()

        # 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.
        else:
            self.check_what_desktop_and_type_is_running()

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

        # User script handling
        self.user_script_process = start_script_and_return_process(self.script_as_list_of_arguments)
        print(f"headless: Started the script with PID {self.user_script_process.pid}")

        # Get exit code from user script process
        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 = Headless()
    headless.execute()

    sys.exit(headless.user_script_exit_code)


if __name__ == "__main__":
    main()
