"""Cross-platform shell command and script execution helpers."""
import ctypes
import re
import shlex
import asyncio
import tempfile
import os
import uuid
import platform
import sys
import subprocess
import base64
from ..utility.utils import to_type, to_s, fstr, log, run_in_executor

__all__ = [
    "nt_set_pshell_unrestricted",
    "nt_is_admin",
    "mac_arg_escape",
    "nix_arg_escape",
    "win_arg_escape",
    "get_powershell_path",
    "cmd",
    "powershell_encoded_cmd",
    "ps1_exec_trick",
    "is_pshell_restricted",
    "nt_pshell",
    "get_arg_escape_func",
    "run_py_script",
    "is_root",
    "win_uac",
    "ensure_root",
    "allow_windows_firewall",
    "remove_windows_firewall",
]

"""
Windows doesn't always use UTF-8 for string encoding.
Depending on which functions you're using.
This completely screws encoding and decoding bytes / strs.
"""
os.environ["PYTHONIOENCODING"] = "utf-8"

"""
This requires the process to have been run as administrator
or with a UAC prompt. However, by using the technique of
writing self-contained Python files to disk that accomplish
privileged operations and then restart themselves as admin
it becomes trivial to do the whole process at run time.

E.g. use run_py_script with a script that calls
pyuac.runAsAdmin() if it's not an admin and is designed
to complete the privileged call - all as an external,
new, python process, while the original continues to
run with its original privileges.

My function nt_pshell already checks if powershell can
run scripts and if not -- automatically unrestricts
powershell and continues running the script.
"""


async def nt_set_pshell_unrestricted():
    import winreg

    def mk_unrestricted(path):
        """Set the PowerShell ExecutionPolicy registry value to Unrestricted for path."""
        # Open path to key for 'writing.'
        key = winreg.OpenKey(
            winreg.HKEY_LOCAL_MACHINE, path, access=winreg.KEY_SET_VALUE
        )

        # Set the sub key.
        # Create it if it doesn't exist.
        winreg.SetValueEx(key, "ExecutionPolicy", 0, winreg.REG_SZ, "Unrestricted")

        # We want to be certain the registry
        # has been updated before proceeding.
        winreg.FlushKey(key)

        # Cleanup.
        winreg.CloseKey(key)

    base_path = "SOFTWARE\\Microsoft\\PowerShell\\1\\ShellIds"
    pshell_path = base_path + "\\Microsoft.PowerShell"
    nostics_path = base_path + "\\ScriptedDiagnostics"
    mk_unrestricted(nostics_path)
    mk_unrestricted(pshell_path)


def nt_is_admin():
    """Return True if the current process is running with Windows administrator privileges."""
    if sys.platform != "win32":
        return False
    try:
        return ctypes.windll.shell32.IsUserAnAdmin()
    except (OSError, AttributeError):
        return False


# Surrounds with DOUBLE quotes.
def mac_arg_escape(arg):
    """Return arg shell-escaped and wrapped in double quotes for macOS/BSD shells."""
    black_list = '"\\'
    buf = ""
    for ch in arg:
        if ch in black_list:
            buf += "\\" + ch
        else:
            buf += ch

    return '"' + buf + '"'


# Note: this function escapes an argument string
# for Unix shell but surrounds it by single quotes.
# It returns a single quoted string with the result.
# Hence the surrounding quotes APPEAR escaped when
# printing them.

# Surrounds with SINGLE quotes.
def nix_arg_escape(arg):
    """Return arg shell-escaped and wrapped in single quotes for POSIX shells."""
    return shlex.quote(arg)


def win_arg_escape(arg, allow_vars=0):
    """Return arg shell-escaped and wrapped in double quotes for Windows cmd.exe."""
    # Double all the backslashes before a double quote.
    # Then ensure the double quote is escaped.
    # Doubling up neutralizes double quotes that may be unbalanced.
    # Since this is done you ensure quotes are all escaped properly.
    arg = re.sub(r'(\\*)"', r'\1\1\\"', arg)

    # Double up continuous backslashes at the end of the string.
    arg = re.sub(r"(\\*)$", r"\1\1", arg)

    # Replace special characters.
    arg = re.sub("([()%!^<>&|;,])", r"^\1", arg)

    # Surround entire arg with quotes.
    # This avoids spaces breaking a command.
    arg = '"%s"' % (arg)

    return arg


def get_powershell_path():
    """Return the path to the highest-version powershell.exe found on Windows."""
    ps_dir = "%WINDIR%\\System32\\WindowsPowerShell"
    ps_dir = os.path.expandvars(ps_dir)
    dir_list = os.listdir(ps_dir)

    version_dir = None
    version_info = None
    for sub_dir in dir_list:
        p = "v([0-9]+)(?:[.]([0-9]+))?"
        out = re.findall(p, sub_dir)
        if not len(out):
            continue
        out = out[0]

        if version_dir is None:
            version_dir = sub_dir
            version_info = out
            continue

        big_version, little_version = out
        if big_version > version_info[0]:
            version_dir = sub_dir
            version_info = out
            continue

        if big_version == version_info[0]:
            if little_version > version_info[1]:
                version_dir = sub_dir
                version_info = out
                continue

    if version_dir is None:
        raise FileNotFoundError("Cannot find powershell dir.")

    ps_dir = os.path.join(ps_dir, version_dir, "powershell.exe")
    return ps_dir


# There is an issue with create_subprocess_shell on Windows 10
# with the latest Python versions. When you try use this code
# a window will pop up asking you how you want to open this
# file. A hack I've found that works is to pass the command
# you want to execute to powershell.
# Example: 'powershell "route print"'


async def cmd(
    value, io=None, er=True, timeout=10
):
    # Output type.
    out_type = value
    null_out = to_type("", out_type)

    # Setup STDIN pipes.
    in_val = None
    if io is not None:
        in_val = asyncio.subprocess.PIPE

    this_dir = os.path.dirname(__file__)
    try:
        if er is None:
            er = asyncio.subprocess.DEVNULL
        else:
            er = None

        proc = await asyncio.create_subprocess_shell(
            value,
            stdin=in_val,
            stdout=asyncio.subprocess.PIPE,
            stderr=er,
            cwd=this_dir,
            shell=True,
        )

        if timeout:
            try:
                task = proc.communicate(input=io)
                stdout, stderr = await asyncio.wait_for(task, timeout)
            except asyncio.TimeoutError:
                log(fstr("command {0} timed out", (value,)))
                return null_out
        else:
            stdout, stderr = await proc.communicate(input=io)
    except NotImplementedError:
        # Backup function for old Python versions.
        @run_in_executor
        def blocking_cmd(er):
            """Run value as a blocking subprocess in a thread-pool executor and return its stdout."""
            try:
                if er is None:
                    er = subprocess.DEVNULL
                else:
                    er = None

                proc = subprocess.run(
                    value,
                    stdout=subprocess.PIPE,
                    shell=True,
                    stderr=er,
                    timeout=timeout,
                    check=False,
                )
                stdout = proc.stdout
                stderr = proc.stderr
            except subprocess.TimeoutExpired:
                log(fstr("command {0} timed out", (value,)))
                return null_out

            return stdout, stderr

        # Use an executor to run the blocking command so the event loop isn't wrekt.
        stdout, stderr = await blocking_cmd(er)

    # Log any visible errors.
    if stderr is not None and len(stderr):
        log(
            fstr(
                "cmd {0} stderr = {1}",
                (
                    value,
                    stderr,
                ),
            )
        )

    # Return command output.
    if stdout is None:
        return null_out
    return to_type(stdout, out_type)


def powershell_encoded_cmd(ps1):
    """Return ps1 encoded as a Base64 UTF-16LE string suitable for powershell -encodedCommand."""
    unicode_bytes = ps1.encode("utf-16le")
    param = base64.b64encode(unicode_bytes)
    return to_s(param)


async def ps1_exec_trick(ps1):
    param = powershell_encoded_cmd(ps1)
    ps_path = get_powershell_path()
    out = await cmd(
        fstr(
            "{0} -encodedCommand {1}",
            (
                ps_path,
                param,
            ),
        ),
        timeout=None,
    )
    log(out)
    return out


async def is_pshell_restricted():
    out = await cmd("powershell Get-ExecutionPolicy", timeout=None)
    return "Unrestricted" not in out


async def nt_pshell(
    value, timeout=10, is_unrestricted=None
):
    # Allow powershell scripts to be run
    # by modifying registry if needed.
    if is_unrestricted is None:
        if await is_pshell_restricted():
            await nt_set_pshell_unrestricted()

    # Write a temp file into the temp dir
    # with the script to execute.
    tmp_dir = tempfile.gettempdir()
    cmd_path = os.path.join(tmp_dir, str(uuid.uuid4()) + ".ps1")
    with open(cmd_path, "w") as f:
        f.write(value)
        f.flush()
        f.close()

    # Command to tell powershell to run the script.
    cmd_str = "powershell.exe -file %s" % (win_arg_escape(cmd_path))

    # Wait for powershell to execute the script.
    if timeout:
        out = await cmd(cmd_str, timeout=timeout)
    else:
        out = await cmd(cmd_str, timeout=None)

    # Delete the old script file.
    if os.path.exists(cmd_path):
        os.remove(cmd_path)

    return out


def get_arg_escape_func():
    """Return the platform-appropriate shell argument escaping function, or None if unknown."""
    if platform.system() == "Linux":
        return nix_arg_escape

    if platform.system() == "Windows":
        return win_arg_escape

    if platform.system() in ["Darwin", "FreeBSD"]:
        return mac_arg_escape

    return None


async def run_py_script(
    script, root_pw=None, cleanup=False
):
    # Write a temp file into the temp dir
    # with the script to execute.
    out = None
    tmp_dir = tempfile.gettempdir()
    cmd_path = os.path.join(tmp_dir, str(uuid.uuid4()) + ".py")
    escape_arg = get_arg_escape_func()
    with open(cmd_path, "w") as f:
        f.write(script)
        f.flush()
        f.close()

        # Sudo prefix.
        prefix = ""
        if platform.system() != "Windows":
            if root_pw is not None:
                prefix = "{ echo %s; } | sudo -k -S " % (escape_arg(root_pw))

        # Build cmd to execute the new script.
        exe_cmd = "%s%s %s" % (prefix, escape_arg(sys.executable), escape_arg(cmd_path))

        # Execute the command.
        out = await cmd(exe_cmd, timeout=None)

        # Delete the script file.
        if cleanup:
            if os.path.exists(cmd_path):
                os.remove(cmd_path)

    return out


def is_root():
    """Return True if the current process is running as root (or Windows administrator)."""
    if platform.system() == "Windows":
        import ctypes

        if ctypes.windll.shell32.IsUserAnAdmin() == 0:
            return False
    else:
        if os.geteuid() != 0:
            return False

    return True


def win_uac():
    """Re-launch the current process with a Windows UAC elevation prompt."""
    if sys.platform != "win32":
        return
    ctypes.windll.shell32.ShellExecuteW(
        None, "runas", sys.executable, " ".join(sys.argv), None, 1
    )


def ensure_root():
    """Raise an exception if the current process is not running as root."""
    if not is_root():
        raise PermissionError("root required for this code.")


def allow_windows_firewall(rule_name):
    """Add an inbound allow rule for the current Python exe. Silent on any error."""
    if sys.platform != "win32":
        return
    try:
        exe = sys.executable
        if sys.getwindowsversion().major >= 6:
            subprocess.call(
                [
                    "netsh", "advfirewall", "firewall", "add", "rule",
                    "name=" + rule_name,
                    "dir=in", "action=allow",
                    "program=" + exe,
                    "protocol=any",
                ],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
        else:
            subprocess.call(
                [
                    "netsh", "firewall", "add", "allowedprogram",
                    "program=" + exe,
                    "name=" + rule_name,
                    "mode=ENABLE",
                ],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
    except Exception:
        pass


def remove_windows_firewall(rule_name):
    """Remove the inbound allow rule added by allow_windows_firewall. Silent on any error."""
    if sys.platform != "win32":
        return
    try:
        if sys.getwindowsversion().major >= 6:
            subprocess.call(
                [
                    "netsh", "advfirewall", "firewall", "delete", "rule",
                    "name=" + rule_name,
                ],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
        else:
            subprocess.call(
                [
                    "netsh", "firewall", "delete", "allowedprogram",
                    "program=" + sys.executable,
                ],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
    except Exception:
        pass
