#!/usr/bin/python3
import os
import re
import shutil

HOME = os.environ.get("HOME", "/root")
SYSTEM_DIR = "/usr/share/applications"
LOCAL_DIR = "/usr/local/share/applications"
FLATPAK_DIR = os.path.join(HOME, ".local/share/flatpak/exports/share/applications")
USER_DIR = os.path.join(HOME, ".local/share/applications")
MARK = "X-DroidDeck-GPU=1"
SKIP = re.compile(r"^(steam|droiddeck-game-|droiddeck-steam|steamdeck-)", re.I)
CHROMIUM_ENV = "env -u WAYLAND_DISPLAY XDG_SESSION_TYPE=x11 ELECTRON_DISABLE_SANDBOX=1 "
CHROMIUM_FLAGS = " --no-sandbox --ozone-platform=x11"


def enabled():
    return os.environ.get("BL_DESKTOP_GPU", "on").strip().lower() not in ("off", "0", "no", "false")


def main_group(text):
    lines = text.splitlines()
    start = next((i for i, l in enumerate(lines) if l.strip() == "[Desktop Entry]"), None)
    if start is None:
        return None, None, None
    end = next((i for i in range(start + 1, len(lines)) if lines[i].startswith("[")), len(lines))
    return lines[:start + 1], lines[start + 1:end], lines[end:]


def split_program(exec_line):
    if exec_line.startswith('"'):
        end = exec_line.find('"', 1)
        if end < 0:
            return None, None
        return exec_line[:end + 1], exec_line[end + 1:]
    program, _, rest = exec_line.partition(" ")
    return program, (" " + rest) if rest else ""


def chromium(exec_line):
    program, _ = split_program(exec_line)
    if not program:
        return False
    program = program.strip('"')
    path = program if os.path.isabs(program) else shutil.which(program)
    if not path:
        return False
    here = os.path.dirname(os.path.realpath(path))
    return any(os.path.exists(os.path.join(d, "chrome-sandbox")) for d in (here, os.path.dirname(here)))


def chromium_exec(exec_line):
    program, rest = split_program(exec_line)
    return CHROMIUM_ENV + program + CHROMIUM_FLAGS + rest


def gpu_exec(exec_line):
    if exec_line.startswith("/usr/bin/flatpak run ") or exec_line.startswith("flatpak run "):
        exec_line = "/usr/local/bin/droiddeck-flatpak-run " + exec_line.split(" run ", 1)[1]
    return "droiddeck-gpu " + exec_line


def wrap(text, gpu_on=True):
    head, group, rest = main_group(text)
    if group is None:
        return None
    keys = {}
    for l in group:
        if "=" in l and not l.lstrip().startswith("#"):
            k, v = l.split("=", 1)
            keys.setdefault(k.strip(), v)
    if keys.get("Type", "Application") != "Application" or keys.get("NoDisplay", "").lower() == "true":
        return None
    exec_line = keys.get("Exec", "").strip()
    if not exec_line or exec_line.startswith("droiddeck-gpu "):
        return None
    cats = set(filter(None, keys.get("Categories", "").split(";")))
    gpu = gpu_on and bool(cats & {"Game", "Emulator"})
    browser = chromium(exec_line)
    if not gpu and not browser:
        return None
    if browser:
        exec_line = chromium_exec(exec_line)
    if not gpu:
        def fix(lines):
            return ["Exec=" + chromium_exec(l[5:].strip()) if l.startswith("Exec=") and chromium(l[5:].strip()) else l
                    for l in lines]
        out_group = fix(group)
        while out_group and not out_group[-1].strip():
            out_group.pop()
        out_group.append(MARK)
        out_rest = fix(rest)
        if out_rest and out_rest[0].strip():
            out_rest.insert(0, "")
        return "\n".join(head + out_group + out_rest) + "\n"
    out_group = []
    actions = keys.get("Actions", "")
    for l in group:
        k = l.split("=", 1)[0].strip() if "=" in l else None
        if k == "Exec":
            out_group.append("Exec=" + gpu_exec(exec_line))
        elif k == "Actions":
            continue
        else:
            out_group.append(l)
    while out_group and not out_group[-1].strip():
        out_group.pop()
    out_group.append("Actions=" + ("droiddeck-window;" + actions if actions else "droiddeck-window;"))
    out_group.append(MARK)
    window = ["", "[Desktop Action droiddeck-window]", "Name=Run in a desktop window", "Exec=" + exec_line]
    wrapped_rest = []
    for l in rest:
        if l.startswith("Exec=") and not l[5:].startswith("droiddeck-gpu "):
            action = l[5:].strip()
            l = "Exec=" + gpu_exec(chromium_exec(action) if chromium(action) else action)
        wrapped_rest.append(l)
    if wrapped_rest and wrapped_rest[0].strip():
        wrapped_rest.insert(0, "")
    return "\n".join(head + out_group + window + wrapped_rest) + "\n"


def ours(path):
    try:
        with open(path, encoding="utf-8", errors="replace") as f:
            return MARK in f.read()
    except OSError:
        return False


def main():
    on = enabled()
    os.makedirs(USER_DIR, exist_ok=True)
    written = set()
    sources = [d for d in (SYSTEM_DIR, LOCAL_DIR, FLATPAK_DIR) if os.path.isdir(d)]
    for source in sources:
        for name in sorted(os.listdir(source)):
            if not name.endswith(".desktop") or SKIP.match(name):
                continue
            target = os.path.join(USER_DIR, name)
            if os.path.exists(target) and not ours(target):
                continue
            try:
                with open(os.path.join(source, name), encoding="utf-8", errors="replace") as f:
                    text = wrap(f.read(), on)
            except OSError:
                continue
            if text is None:
                continue
            with open(target + ".tmp", "w", encoding="utf-8") as f:
                f.write(text)
            os.replace(target + ".tmp", target)
            written.add(name)
    removed = 0
    for name in os.listdir(USER_DIR):
        path = os.path.join(USER_DIR, name)
        if name.endswith(".desktop") and name not in written and ours(path):
            os.remove(path)
            removed += 1
    print(f"droiddeck-desktop-gpu: {'on' if on else 'off'}, {len(written)} entries on the GPU, {removed} removed")


main()