import os
import ctypes
import time
import ssl
import urllib2
import re
from subprocess import Popen, PIPE
import _winreg as reg  # For Python 2

# Constants
URL = "https://cdn-patchportal-one.comodo.com/portal/packages/spm/AnyDesk/x86/AnyDesk_9.5.11.exe"
UNINSTALL_PATH = '"C:\\Program Files\\AnyDesk\\AnyDesk.exe" --remove'
PROCESS_NAMES = ["AnyDesk.exe", "ad_svc.exe", "AnyDesk_Service.exe"]

# Temp download path
Down_path = os.environ['TEMP']
fileName = URL.split('/')[-1]
DownTo = os.path.join(Down_path, fileName)

# Context manager for disabling filesystem redirection
class disable_file_system_redirection:
    _disable = ctypes.windll.kernel32.Wow64DisableWow64FsRedirection
    _revert = ctypes.windll.kernel32.Wow64RevertWow64FsRedirection

    def __enter__(self):
        self.old_value = ctypes.c_long()
        self.success = self._disable(ctypes.byref(self.old_value))

    def __exit__(self, type, value, traceback):
        if self.success:
            self._revert(self.old_value)

# Execute command with FS redirection off and with manual timeout handling
def ecmd(command, timeout=30):
    with disable_file_system_redirection():
        obj = Popen(command, shell=True, stdout=PIPE, stderr=PIPE)
        start_time = time.time()
        while True:
            # Check if the command has completed
            if obj.poll() is not None:
                break
            # Check if timeout is reached
            if time.time() - start_time > timeout:
                obj.kill()
                return "Command timed out or was forcibly killed."
            time.sleep(0.1)  # Sleep briefly to avoid high CPU usage
        out, err = obj.communicate()
    return out.strip() if out else err.strip()

# Kill all AnyDesk processes
def kill_anydesk():
    print("Terminating AnyDesk processes...")
    killed = False
    for proc in PROCESS_NAMES:
        result = os.system('taskkill /f /im "{}" >nul 2>&1'.format(proc))
        if result == 0:
            print("Terminated:", proc)
            killed = True
    if not killed:
        print("No AnyDesk processes found or already terminated.")

# Download AnyDesk installer
def downloadFile(DownTo, fromURL):
    headers = {'User-Agent': 'Mozilla/5.0'}
    context = ssl._create_unverified_context()
    request = urllib2.Request(fromURL, headers=headers)
    req = urllib2.urlopen(request, context=context)
    try:
        with open(DownTo, 'wb') as f:
            while True:
                chunk = req.read(8192)
                if chunk:
                    f.write(chunk)
                else:
                    break
        if os.path.isfile(DownTo):
            return '{} - {}KB'.format(DownTo, os.path.getsize(DownTo) / 1024)
    except:
        return 'Download failed. Check URL or path.'

# Modify DisplayVersion in registry by cleaning up the version string
def modify_registry_display_version():
    uninstall_paths = [
        (reg.HKEY_LOCAL_MACHINE, r"SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall", reg.KEY_WOW64_64KEY),
        (reg.HKEY_LOCAL_MACHINE, r"SOFTWARE\Wow6432Node\Microsoft\Windows\CurrentVersion\Uninstall", reg.KEY_WOW64_32KEY)
    ]
    for hive, path, view_flag in uninstall_paths:
        try:
            with disable_file_system_redirection():
                main_key = reg.OpenKey(hive, path, 0, reg.KEY_READ | view_flag)
                for i in range(reg.QueryInfoKey(main_key)[0]):
                    try:
                        subkey_name = reg.EnumKey(main_key, i)
                        subkey_path = path + "\\" + subkey_name
                        subkey = reg.OpenKey(hive, subkey_path, 0, reg.KEY_ALL_ACCESS | view_flag)
                        try:
                            display_name, _ = reg.QueryValueEx(subkey, "DisplayName")
                            if "AnyDesk" in display_name:
                                print("Found AnyDesk:", display_name)
                                try:
                                    # Get the existing DisplayVersion
                                    display_version, _ = reg.QueryValueEx(subkey, "DisplayVersion")
                                    print("Original DisplayVersion:", display_version)
                                    
                                    # Remove spaces and alphabetic characters (keep dots)
                                    cleaned_version = re.sub(r"[^0-9.]", "", display_version)
                                    print("Cleaned DisplayVersion:", cleaned_version)
                                    
                                    # Set the cleaned version back to the registry
                                    reg.SetValueEx(subkey, "DisplayVersion", 0, reg.REG_SZ, cleaned_version)
                                    print("Updated DisplayVersion to:", cleaned_version)
                                    
                                    reg.CloseKey(subkey)
                                    return
                                except KeyError:
                                    print("DisplayVersion not found.")
                            reg.CloseKey(subkey)
                        except Exception as e:
                            reg.CloseKey(subkey)
                            continue
                    except Exception:
                        continue
                reg.CloseKey(main_key)
        except Exception as e:
            print("Registry access error:", str(e))

# Install AnyDesk
def spm_install(dir, args):
    with disable_file_system_redirection():
        print("Downloading AnyDesk...")
        print(downloadFile(DownTo, URL))
        print("Installing silently...")
        install_cmd = '"%s" --remove-first --install "C:\\Program Files\\AnyDesk" --create-desktop-icon' % DownTo
        print(ecmd(install_cmd))
        time.sleep(10)
        modify_registry_display_version()

# Uninstall AnyDesk
def spm_uninstall(dir, args):
    with disable_file_system_redirection():
        print("Uninstalling AnyDesk...")
        kill_anydesk()
        print(ecmd(UNINSTALL_PATH))
        time.sleep(15)
        kill_anydesk()

# Update AnyDesk
def spm_update(dir, args):
    spm_uninstall(dir, args)
    time.sleep(10)
    spm_install(dir, args)