#!/usr/bin/env python3
#--------------------------------------------------------------------------------------------------------
# Name: Linux Lite - Lite Menu Sorter
# Architecture: all
# Author: Jerry Bezencon
# Website: https://www.linuxliteos.com
# Language: Python/GTK4
# Licence: GPLv2
#--------------------------------------------------------------------------------------------------------
"""Pick categories from the Whisker Menu and sort their apps alphabetically.

Categories themselves are NOT reordered. The top-level <Layout> in
~/.config/menus/xfce-applications.menu (which controls the order of the
category list) is left untouched. For each ticked category, this app
wholesale-replaces only that submenu's <Layout> body with:

    <Merge type="menus" />
    <Filename>app1.desktop</Filename>
    <Filename>app2.desktop</Filename>
    ...
    <Merge type="files" />

— the same shape menulibre's `Sort A-Z` writes (per
MenulibreXmlMenuElementTree.model_to_xml_layout). Sort key is `str.lower`,
verbatim from MenulibreTreeview._sort_iter:1058.

Optional cleanup: removes ~/.config/menus/applications-merged/*.menu files
that contain `<Name>xfce-<X></Name>`. xdg-desktop-menu auto-creates these
when apps install themselves; libxfce4menu treats them as a SEPARATE menu
from the canonical category, producing duplicate entries in Whisker. This
is the same cleanup menulibre does in `_cleanup_applications_merged`.

Alacarte runs (alacarte-made-N.desktop) are never reordered.
"""

import gi
gi.require_version("Gtk", "4.0")
gi.require_version("Adw", "1")
gi.require_version("GMenu", "3.0")
from gi.repository import Gtk, Adw, GMenu, GLib, Gdk

import re
import shutil
import subprocess
import sys
from datetime import datetime
from pathlib import Path

APP_ID = "com.linuxlite.menusorter"
APP_TITLE = "Lite Menu Sorter"
ICON_NAME = "lite-menusorter"
WIDTH = 540
HEIGHT = 640

LINE_FILENAME_RE = re.compile(r"^([ \t]*)<Filename>([^<]+)</Filename>[ \t]*$")
DUPLICATE_NAME_RE = re.compile(r"<Name>xfce-[a-z0-9_-]+</Name>")


# ---------------------------------------------------------------- GMenu walk

def get_categories(sort_mode="generic"):
    """Walk the live menu via GMenu (the same library Whisker reads) and
    return a list of (menu_id, display_name, [(filename, display_name) sorted]).
    Top-level categories only — submenus like Screensavers handled by their
    parent's <Merge type="menus" />.

    `sort_mode` selects the sort key for each app:
      "generic" — sort by GenericName= (falls back to Name= when missing).
                  Groups apps by function: e.g. all "Web Browser" entries
                  sort together regardless of display name.
      "name"    — sort by Name= (the display name). Matches what users see
                  when "Show generic application names" is off in Whisker.

    Skips alacarte-* containers (user-curated order should not be sorted)."""
    flags = (
        GMenu.TreeFlags.SORT_DISPLAY_NAME
    )
    tree = GMenu.Tree.new("xfce-applications.menu", flags)
    if not tree.load_sync():
        return []
    root = tree.get_root_directory()
    if root is None:
        return []

    categories = []
    iter_ = root.iter()
    while True:
        item = iter_.next()
        if item == GMenu.TreeItemType.INVALID:
            break
        if item != GMenu.TreeItemType.DIRECTORY:
            continue
        d = iter_.get_directory()
        menu_id = d.get_menu_id() or ""
        # Skip user-curated alacarte containers.
        if menu_id.startswith("alacarte"):
            continue
        # Skip the duplicate-causing menus injected by applications-merged
        # (id="xfce-accessories", "xfce-graphics", etc. — these show as a
        # second copy of the canonical category in Whisker, and the dedup
        # option in this app is what removes them).
        if menu_id.startswith("xfce-"):
            continue
        display = d.get_name() or menu_id

        apps = []
        sub = d.iter()
        while True:
            t2 = sub.next()
            if t2 == GMenu.TreeItemType.INVALID:
                break
            if t2 != GMenu.TreeItemType.ENTRY:
                continue
            entry = sub.get_entry()
            ai = entry.get_app_info()
            display_name = ai.get_display_name() or entry.get_desktop_file_id()
            generic = ai.get_generic_name()
            if sort_mode == "name":
                sort_key = display_name
            else:
                sort_key = generic or display_name
            apps.append((entry.get_desktop_file_id(), sort_key, display_name))

        if not apps:
            continue
        apps.sort(key=lambda t: t[1].lower())
        # Drop the sort key from what we expose to the rest of the app — keep
        # (filename, display_name) so existing callers see no change.
        apps = [(fid, dn) for fid, _key, dn in apps]
        categories.append((menu_id, display, apps))
    return categories


# ---------------------------------------------------------------- XML rewrite

def detect_indent(body, default="\t\t\t"):
    m = re.search(r"(?m)^([ \t]+)\S", body)
    return m.group(1) if m else default


def has_alacarte(body):
    for ln in body.split("\n"):
        m = LINE_FILENAME_RE.match(ln)
        if m and m.group(2).startswith("alacarte"):
            return True
    return False


def build_layout_body(filenames, indent):
    lines = [f'{indent}<Merge type="menus" />']
    lines.extend(f"{indent}<Filename>{f}</Filename>" for f in filenames)
    lines.append(f'{indent}<Merge type="files" />')
    return "\n" + "\n".join(lines) + "\n"


def rewrite_selected_layouts(text, selected):
    """selected: dict[menu_name -> list of sorted .desktop ids].
    For each <Menu><Name>X</Name>...<Layout>...</Layout></Menu> in the file
    where X is in `selected`, replace the Layout body with a sorted list.

    The root <Menu> (the Xfce container — depth 1 in our nesting stack) is
    NEVER touched, so the category-list order in the parent menu file is
    preserved.

    Layouts whose body contains an alacarte-* Filename are also left alone.
    """
    pattern = re.compile(
        r"<Menu>|</Menu>|<Name>([^<]+)</Name>|(<Layout\b[^>]*>)(.*?)(</Layout>)",
        re.DOTALL,
    )
    stack = []
    out = []
    pos = 0
    for m in pattern.finditer(text):
        out.append(text[pos:m.start()])
        pos = m.end()
        whole = m.group(0)

        if whole == "<Menu>":
            stack.append(None)
            out.append(whole)
            continue
        if whole == "</Menu>":
            if stack:
                stack.pop()
            out.append(whole)
            continue
        if m.group(1) is not None:
            if stack and stack[-1] is None:
                stack[-1] = m.group(1)
            out.append(whole)
            continue
        if m.group(2) is not None:
            layout_open, body, layout_close = m.group(2), m.group(3), m.group(4)
            menu_name = stack[-1] if stack else None
            depth = len(stack)
            skip = (
                not menu_name
                or depth <= 1               # root menu's own Layout — never touch
                or menu_name not in selected
                or has_alacarte(body)
            )
            if skip:
                out.append(whole)
            else:
                indent = detect_indent(body)
                new_body = build_layout_body(selected[menu_name], indent)
                out.append(layout_open + new_body + indent[:-1] + layout_close)
            continue

    out.append(text[pos:])
    return "".join(out)


def cleanup_duplicate_merged(merged_dir, stamp):
    """Move (rename) every applications-merged/*.menu file containing
    <Name>xfce-<X></Name> to <name>.bak-<stamp>. Returns the number moved."""
    if not merged_dir.is_dir():
        return 0
    moved = 0
    for f in sorted(merged_dir.glob("*.menu")):
        try:
            text = f.read_text(encoding="utf-8")
        except OSError:
            continue
        if DUPLICATE_NAME_RE.search(text):
            try:
                f.rename(f.with_name(f"{f.name}.bak-{stamp}"))
                moved += 1
            except OSError:
                pass
    return moved


# ---------------------------------------------------------------- GTK4 UI

class MenuSorterWindow(Adw.ApplicationWindow):
    def __init__(self, app):
        super().__init__(application=app)
        self.set_title(APP_TITLE)
        self.set_default_size(WIDTH, HEIGHT)
        self.set_icon_name(ICON_NAME)
        self.connect("realize", self._on_realize)

        self._categories = []
        self._rows = []  # (menu_id, checkbutton, action_row)
        # Always launch with Generic Name sorting enabled, regardless of the
        # user's Whisker `launcher-show-name` xfconf state. The switch can be
        # toggled off at runtime to fall back to display-name sort.
        self._sort_mode = "generic"

        main = Gtk.Box(orientation=Gtk.Orientation.VERTICAL)
        self.set_content(main)

        header = Adw.HeaderBar()
        main.append(header)

        scrolled = Gtk.ScrolledWindow()
        scrolled.set_policy(Gtk.PolicyType.NEVER, Gtk.PolicyType.AUTOMATIC)
        scrolled.set_vexpand(True)
        main.append(scrolled)

        clamp = Adw.Clamp()
        clamp.set_maximum_size(480)
        clamp.set_margin_top(20)
        clamp.set_margin_bottom(20)
        clamp.set_margin_start(20)
        clamp.set_margin_end(20)
        scrolled.set_child(clamp)

        content = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=16)
        clamp.set_child(content)

        desc = Gtk.Label(
            label="Tick the menu categories whose apps you want sorted "
                  "alphabetically. The category list itself is not reordered."
        )
        desc.set_wrap(True)
        desc.set_xalign(0.5)
        desc.set_justify(Gtk.Justification.CENTER)
        desc.add_css_class("dim-label")
        content.append(desc)

        # Sort key — placed at the top so the user picks before ticking
        # categories. Always launches in Generic mode (LL 8.0 default).
        sort_group = Adw.PreferencesGroup()
        sort_group.set_title("Sort By")
        content.append(sort_group)

        sort_row = Adw.ActionRow()
        sort_row.set_title("Sort by Generic Name")
        sort_row.set_subtitle(
            "ON: group apps by function (e.g. ‘Package Manager’, "
            "‘Web Browser’).  OFF: sort by display name."
        )
        self.sort_switch = Gtk.Switch()
        self.sort_switch.set_valign(Gtk.Align.CENTER)
        self.sort_switch.set_active(True)
        self.sort_switch.connect("notify::active", self._on_sort_toggle)
        sort_row.add_suffix(self.sort_switch)
        sort_row.set_activatable_widget(self.sort_switch)
        sort_group.add(sort_row)

        # Categories
        self.cat_group = Adw.PreferencesGroup()
        self.cat_group.set_title("Categories")
        content.append(self.cat_group)

        sa_row = Adw.ActionRow()
        sa_row.set_title("Select All")
        self.select_all_check = Gtk.CheckButton()
        self.select_all_check.set_valign(Gtk.Align.CENTER)
        self.select_all_check.connect("toggled", self._on_select_all_toggled)
        sa_row.add_suffix(self.select_all_check)
        sa_row.set_activatable_widget(self.select_all_check)
        self.cat_group.add(sa_row)

        self._populate_categories()

        # Options
        opts = Adw.PreferencesGroup()
        opts.set_title("Options")
        content.append(opts)

        dedup_row = Adw.ActionRow()
        dedup_row.set_title("Remove duplicate categories")
        dedup_row.set_subtitle(
            "Cleans up applications-merged files that cause Whisker to show "
            "categories twice. Safe — affected apps re-merge into their "
            "canonical category via Categories= field."
        )
        self.dedup_check = Gtk.CheckButton()
        self.dedup_check.set_valign(Gtk.Align.CENTER)
        self.dedup_check.set_active(True)
        self.dedup_check.connect("toggled", self._update_apply_sensitivity)
        dedup_row.add_suffix(self.dedup_check)
        dedup_row.set_activatable_widget(self.dedup_check)
        opts.add(dedup_row)

        # Apply
        self.apply_btn = Gtk.Button(label="Apply")
        self.apply_btn.add_css_class("suggested-action")
        self.apply_btn.add_css_class("pill")
        self.apply_btn.set_halign(Gtk.Align.CENTER)
        self.apply_btn.set_size_request(200, -1)
        self.apply_btn.connect("clicked", self._on_apply)
        content.append(self.apply_btn)

        self.status = Gtk.Label()
        self.status.set_wrap(True)
        self.status.set_xalign(0.5)
        self.status.set_justify(Gtk.Justification.CENTER)
        self.status.add_css_class("dim-label")
        content.append(self.status)

        self._update_apply_sensitivity()

    def _populate_categories(self):
        try:
            self._categories = get_categories(self._sort_mode)
        except Exception as e:
            self.status.set_label(f"Failed to load menu: {e}")
            return
        for menu_id, display, apps in self._categories:
            row = Adw.ActionRow()
            row.set_title(display)
            row.set_subtitle(f"{len(apps)} app{'s' if len(apps) != 1 else ''}")
            check = Gtk.CheckButton()
            check.set_valign(Gtk.Align.CENTER)
            check.connect("toggled", self._on_check_toggled)
            row.add_suffix(check)
            row.set_activatable_widget(check)
            self._rows.append((menu_id, check, row))
            self.cat_group.add(row)

    def _on_sort_toggle(self, switch, _param):
        self._sort_mode = "generic" if switch.get_active() else "name"
        # Re-resolve the per-category sorted app lists with the new key.
        # The picker UI rows themselves don't change — they show category
        # names, not apps. Apply uses self._categories on click.
        try:
            self._categories = get_categories(self._sort_mode)
        except Exception as e:
            self.status.set_label(f"Failed to re-sort: {e}")

    def _on_select_all_toggled(self, btn):
        active = btn.get_active()
        for _, check, _ in self._rows:
            check.set_active(active)

    def _on_check_toggled(self, btn):
        self._update_apply_sensitivity()

    def _update_apply_sensitivity(self, *_):
        any_cat = any(c.get_active() for _, c, _ in self._rows)
        self.apply_btn.set_sensitive(any_cat or self.dedup_check.get_active())

    def _on_apply(self, button):
        selected = {}
        names = []
        for menu_id, check, _ in self._rows:
            if not check.get_active():
                continue
            for mid, dn, apps in self._categories:
                if mid == menu_id:
                    selected[mid] = [a[0] for a in apps]
                    names.append(dn)
                    break

        do_dedup = self.dedup_check.get_active()

        body_lines = []
        if selected:
            body_lines.append("Categories to sort alphabetically:")
            body_lines.extend(f"  • {n}" for n in names)
        if do_dedup:
            if body_lines:
                body_lines.append("")
            body_lines.append("Plus: remove duplicate-category files from "
                              "~/.config/menus/applications-merged/.")
        body_lines.append("")
        body_lines.append("A timestamped .bak file will be created before "
                          "any change.")

        dialog = Adw.AlertDialog()
        dialog.set_heading("Apply Changes?")
        dialog.set_body("\n".join(body_lines))
        dialog.add_response("cancel", "Cancel")
        dialog.add_response("apply", "Apply")
        dialog.set_response_appearance("apply", Adw.ResponseAppearance.SUGGESTED)
        dialog.set_default_response("cancel")
        dialog.set_close_response("cancel")
        dialog.connect("response", self._on_confirm, selected, do_dedup)
        dialog.present(self)

    def _on_confirm(self, dialog, response, selected, do_dedup):
        if response != "apply":
            return

        results = []
        try:
            menus_dir = Path.home() / ".config" / "menus"
            stamp = datetime.now().strftime("%Y%m%d-%H%M%S")

            if do_dedup:
                count = cleanup_duplicate_merged(
                    menus_dir / "applications-merged", stamp,
                )
                if count > 0:
                    results.append(
                        f"Removed {count} duplicate-category file"
                        f"{'s' if count != 1 else ''}"
                    )
                else:
                    results.append("No duplicate-category files found")

            if selected:
                menu_file = menus_dir / "xfce-applications.menu"
                if menu_file.is_file():
                    text = menu_file.read_text(encoding="utf-8")
                    new_text = rewrite_selected_layouts(text, selected)
                    if new_text != text:
                        backup = menu_file.with_name(
                            f"{menu_file.name}.bak-{stamp}")
                        shutil.copy2(menu_file, backup)
                        menu_file.write_text(new_text, encoding="utf-8")
                        results.append(
                            f"Sorted {len(selected)} categor"
                            f"{'ies' if len(selected) != 1 else 'y'} "
                            f"(backup: {backup.name})"
                        )
                    else:
                        results.append("No layout changes were needed")
                else:
                    results.append(f"Menu file not found: {menu_file}")

            # Restart xfce4-panel asynchronously so the GTK main loop is never
            # blocked waiting for the panel to come back. Synchronous run() here
            # caused xfwm4 to log XSync timeout warnings during Apply because
            # the window stopped processing events for 1-3s. start_new_session
            # detaches the child so it isn't left as a zombie if we exit early.
            try:
                subprocess.Popen(
                    ["xfce4-panel", "-r"],
                    stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                    start_new_session=True,
                )
                results.append("Whisker Menu refreshed")
            except Exception:
                pass

            self.status.set_label(" • ".join(results))
        except Exception as e:
            self.status.set_label(f"Error: {e}")

    # ---------- window centering (LL convention)

    def _on_realize(self, *_):
        GLib.timeout_add(200, self._center_window)

    def _center_window(self):
        try:
            display = Gdk.Display.get_default()
            if not display:
                return False
            monitors = display.get_monitors()
            if monitors.get_n_items() == 0:
                return False
            geometry = monitors.get_item(0).get_geometry()
            x = geometry.x + (geometry.width - WIDTH) // 2
            y = geometry.y + (geometry.height - HEIGHT) // 2
            subprocess.run(
                ["wmctrl", "-r", APP_TITLE, "-e",
                 f"0,{x},{y},{WIDTH},{HEIGHT}"],
                capture_output=True, timeout=2,
            )
        except Exception:
            pass
        return False


class MenuSorterApp(Adw.Application):
    def __init__(self):
        super().__init__(application_id=APP_ID)

    def do_activate(self):
        win = self.props.active_window
        if not win:
            win = MenuSorterWindow(self)
        win.present()


def main():
    app = MenuSorterApp()
    return app.run(sys.argv)


if __name__ == "__main__":
    sys.exit(main())
