import os
import shutil
import re
import xml.etree.ElementTree as ET

original_xml = r"{{ start_menu_xml }}"
backup_xml = r"{{ backup_xml }}"

# This script assumes backup and copying already handled by Ansible

START_MENU_PATHS = {
    "%ALLUSERSPROFILE%\\Microsoft\\Windows\\Start Menu\\Programs": os.path.expandvars(r"%ALLUSERSPROFILE%\Microsoft\Windows\Start Menu\Programs")
}

def extract_version(filename):
    match = re.search(r'(\d+\.\d+\.\d+|\d+)', filename)
    if match:
        return tuple(map(int, match.group(0).split('.')))
    return (0,)

def adjust_app_name(app_name):
    return f"{app_name} " if len(app_name) == 1 else app_name

def remove_bracketed_text(filename):
    return re.sub(r"[\(\[].*?[\)\]]", "", filename).strip()

def folder_has_standalone_version(folder_name):
    return bool(re.search(r'\b\d+(?:\.\d+){1,3}\b', folder_name))

def clean_folder_name(folder_name):
    return re.sub(r'\b\d+(?:\.\d+){1,3}\b', '', folder_name).strip()

def remove_version_number(app_name):
    if re.search(r'\d+\.\d+(\.\d+)*$', app_name):
        return re.sub(r'(\s*\d+\.\d+(\.\d+)*\s*)$', '', app_name).strip()
    return app_name

def is_extra_word_present(filename, base_name):
    name_without_ext = os.path.splitext(filename)[0].lower()
    return name_without_ext != base_name.lower()

def version_sort_key(filename, base_name):
    version = extract_version(filename)
    cleaned_name = remove_bracketed_text(filename)
    padded_version = version + (0,) * (3 - len(version))
    return (
        is_extra_word_present(filename, base_name),
        len(cleaned_name),
        -padded_version[0],
        -padded_version[1] if len(padded_version) > 1 else 0,
        -padded_version[2] if len(padded_version) > 2 else 0
    )

def get_clean_base_name(app_name):
    return re.sub(r"(\s*\d+(\.\d+)*\s*)+$", "", app_name).strip()

def find_best_match(app_name, available_files):
    base_name = get_clean_base_name(app_name)
    adjusted_name = adjust_app_name(base_name)
    start_matches = [file for file in available_files if file.lower().startswith(adjusted_name.lower())]
    exact_match = [f for f in start_matches if os.path.splitext(f)[0].lower() == adjusted_name.lower()]
    if exact_match:
        print(f"Exact match for {app_name}: {exact_match[0]}")
        return exact_match[0]
    best = min(start_matches, key=lambda f: version_sort_key(f, adjusted_name)) if start_matches else None
    print(f"Best match for {app_name}: {best}")
    return best

def convert_to_env_path(abs_path):
    for env_path, abs_root in START_MENU_PATHS.items():
        if abs_path.startswith(abs_root):
            return abs_path.replace(abs_root, env_path)
    return abs_path

def find_shortcut_paths(app_names):
    app_paths = {}
    all_lnk_files = {}

    for _, start_menu in START_MENU_PATHS.items():
        for root, _, files in os.walk(start_menu):
            for file in files:
                if file.endswith(".lnk"):
                    full_path = os.path.join(root, file)
                    all_lnk_files[file] = full_path

    for app_name in app_names:
        best_match = find_best_match(app_name, all_lnk_files.keys())
        if best_match:
            env_path = convert_to_env_path(all_lnk_files[best_match])
            app_paths[app_name] = env_path
            print(f"Mapped {app_name} to {env_path}")
        else:
            print(f"No match found for {app_name}")

    return app_paths

def extract_app_names_and_paths(xml_file):
    tree = ET.parse(xml_file)
    root = tree.getroot()
    apps = {}

    for tile in root.findall(".//{*}DesktopApplicationTile"):
        link_path = tile.get("DesktopApplicationLinkPath")
        if link_path:
            folder_name = os.path.basename(os.path.dirname(link_path))
            lnk_name = os.path.splitext(os.path.basename(link_path))[0]

            if folder_has_standalone_version(folder_name):
                app_name = clean_folder_name(folder_name)
            else:
                app_name = remove_bracketed_text(lnk_name)

            app_name = remove_version_number(app_name)
            app_name = adjust_app_name(app_name)
            apps[app_name] = link_path
            print(f"Extracted {app_name} from {link_path}")

    return apps

def update_xml_paths(xml_file, old_new_paths):
    with open(xml_file, "r", encoding="utf-8") as file:
        content = file.read()

    for old_path, new_path in old_new_paths:
        content = content.replace(old_path, new_path)
        print(f"Replaced {old_path} with {new_path}")

    with open(xml_file, "w", encoding="utf-8") as file:
        file.write(content)

    print("Start menu XML updated successfully.")

if __name__ == "__main__":
    xml_file = original_xml
    extracted_apps_with_paths = extract_app_names_and_paths(xml_file)
    app_shortcuts = find_shortcut_paths(extracted_apps_with_paths.keys())

    old_new_paths = []
    for app_name, old_path in extracted_apps_with_paths.items():
        new_path = app_shortcuts.get(app_name)
        if new_path:
            old_new_paths.append((old_path, new_path))
        else:
            print(f"Warning: No new path found for {app_name}. Keeping old path.")

    update_xml_paths(xml_file, old_new_paths)