#!/usr/bin/python3
"""Restricted privileged operations for AnduinOS Driver Center."""

from __future__ import annotations

import os
import re
import subprocess
import sys
import json


XBOX_PACKAGE = "anduinos-xbox-controller-driver"
AUDIO_PACKAGES = ("firmware-sof-anduinos", "alsa-ucm-conf-anduinos")
PRINTING_PACKAGES = (
    "cups",
    "cups-client",
    "cups-core-drivers",
    "cups-filters",
    "cups-filters-core-drivers",
    "cups-ipp-utils",
    "cups-browsed",
    "avahi-daemon",
    "ipp-usb",
    "cups-pk-helper",
    "printer-driver-all",
    "sane-airscan",
)
PRINTING_UNITS = (
    "cups.service",
    "cups.socket",
    "cups.path",
    "cups-browsed.service",
    "ipp-usb.service",
)
PRINTING_AUTOSTART_UNITS = (
    "cups.service",
    "cups.socket",
    "cups.path",
    "cups-browsed.service",
)
PACKAGE_PATTERN = re.compile(r"^[a-z0-9][a-z0-9+.-]+$")
BINARY_PACKAGE_PATTERN = re.compile(r"^[a-z0-9][a-z0-9+.-]+(?::[a-z0-9][a-z0-9-]*)?$")
NVIDIA_DRIVER_PATTERN = re.compile(
    r"^nvidia-driver-(?P<series>[0-9]+)(?P<variant>(?:-[a-z0-9]+)*)$"
)


def run(command: list[str], *, stdin: str | None = None) -> None:
    print("+ " + " ".join(command), flush=True)
    result = subprocess.run(command, input=stdin, text=True, check=False)
    if result.returncode:
        raise RuntimeError(f"{command[0]} exited with status {result.returncode}")


def available_driver_packages() -> set[str]:
    result = subprocess.run(
        ["ubuntu-drivers", "devices"], capture_output=True, text=True, check=False
    )
    packages: set[str] = set()
    for line in result.stdout.splitlines():
        if line.strip().startswith("driver") and ":" in line:
            package = line.split(":", 1)[1].strip().split(" - ", 1)[0]
            if PACKAGE_PATTERN.fullmatch(package):
                packages.add(package)
    return packages


def apt_update() -> None:
    run(["apt-get", "update"])


def package_is_installed(package: str) -> bool:
    result = subprocess.run(
        ["dpkg-query", "-W", "-f=${db:Status-Abbrev}", package],
        capture_output=True,
        text=True,
        check=False,
    )
    return result.returncode == 0 and result.stdout.startswith("ii ")


def installed_nvidia_packages() -> list[str]:
    result = subprocess.run(
        [
            "dpkg-query",
            "-W",
            "-f=${binary:Package}\t${db:Status-Abbrev}\n",
            "nvidia-*",
            "libnvidia-*",
            "linux-modules-nvidia-*",
            "linux-objects-nvidia-*",
            "linux-signatures-nvidia-*",
            "xserver-xorg-video-nvidia-*",
        ],
        capture_output=True,
        text=True,
        check=False,
    )
    packages = []
    for line in result.stdout.splitlines():
        fields = line.split("\t", 1)
        if len(fields) == 2 and fields[1].startswith("ii "):
            package = fields[0]
            if BINARY_PACKAGE_PATTERN.fullmatch(package):
                packages.append(package)
    return sorted(set(packages))


def nvidia_repair_packages(package: str, kernel_release: str) -> list[str]:
    """Return only installed packages belonging to the selected driver family."""
    match = NVIDIA_DRIVER_PATTERN.fullmatch(package)
    if not match:
        raise ValueError("The selected package is not an NVIDIA driver metapackage")
    series = match.group("series")
    family = series + match.group("variant")
    series_token = re.compile(rf"(?:^|-){re.escape(series)}(?:-|$)")
    variant_specific = (
        "nvidia-driver-",
        "nvidia-dkms-",
        "nvidia-kernel-source-",
    )
    kernel_specific = (
        "linux-modules-nvidia-",
        "linux-objects-nvidia-",
        "linux-signatures-nvidia-",
    )
    selected: list[str] = []
    for binary_package in installed_nvidia_packages():
        name = binary_package.split(":", 1)[0]
        if name.startswith(variant_specific):
            if name == package or name.endswith(f"-{family}"):
                selected.append(binary_package)
            continue
        if name.startswith(kernel_specific):
            prefix_matches = any(
                name.startswith(f"{prefix}{family}-") for prefix in kernel_specific
            )
            if prefix_matches and kernel_release in name:
                selected.append(binary_package)
            continue
        if series_token.search(name):
            selected.append(binary_package)
    return sorted(set(selected))


def repair_nvidia_stack(package: str) -> None:
    kernel_release = os.uname().release
    packages = nvidia_repair_packages(package, kernel_release)
    if not packages:
        raise RuntimeError("no installed packages from the selected NVIDIA family were found")
    run(["apt-get", "install", "-y", "--reinstall", *packages])
    run(["depmod", "-a", kernel_release])
    # Static inspection only: loading a graphics module in the live desktop can
    # terminate the compositor. Activation is deliberately deferred to reboot.
    run(["modinfo", "-k", kernel_release, "nvidia"])


def install_driver(package: str) -> None:
    if package not in available_driver_packages():
        raise ValueError("The selected package is not offered by ubuntu-drivers")
    repair_existing = package_is_installed(package)
    apt_update()
    # Let ubuntu-drivers perform family/version switching instead of treating
    # a driver metapackage like an unrelated APT package.
    run(["ubuntu-drivers", "install", package])
    if repair_existing:
        repair_nvidia_stack(package)


def repair_nvidia_driver(package: str) -> None:
    if package not in available_driver_packages():
        raise ValueError("The selected package is not offered by ubuntu-drivers")
    if not package_is_installed(package):
        raise ValueError("The selected NVIDIA driver is not installed")
    apt_update()
    run(["ubuntu-drivers", "install", package])
    repair_nvidia_stack(package)


def install_recommended() -> None:
    apt_update()
    run(["ubuntu-drivers", "install"])


def update_driver(package: str) -> None:
    """Execute only the exact transaction confirmed in the UI."""
    from anduinos_driver_center.driver_update import execute
    if not PACKAGE_PATTERN.fullmatch(package) or package not in available_driver_packages():
        raise ValueError("The selected package is not offered by ubuntu-drivers")
    payload = sys.stdin.read(262145)
    if len(payload) > 262144:
        raise ValueError("Update request is too large")
    result = execute(package, json.loads(payload))
    print(f"{package}: {result['before']} -> {result['after']}", flush=True)


def refresh_driver_info() -> None:
    apt_update()


def install_xbox(reinstall: bool = False) -> None:
    apt_update()
    kernel_release = os.uname().release
    command = ["apt-get", "install", "-y"]
    if reinstall:
        command.append("--reinstall")
    command.extend([f"linux-headers-{kernel_release}", XBOX_PACKAGE])
    run(command)
    # The package's maintainer script historically tolerated DKMS failures.
    # Verify the current running kernel explicitly so a false success can never
    # be reported to the UI.
    run(["dkms", "autoinstall", "-k", kernel_release])
    run(["depmod", "-a", kernel_release])
    run(["modinfo", "-k", kernel_release, "hid-xpadneo"])


def install_audio(reinstall: bool = False) -> None:
    apt_update()
    command = ["apt-get", "install", "-y"]
    if reinstall:
        command.append("--reinstall")
    run([*command, *AUDIO_PACKAGES])


def install_printing_support() -> None:
    apt_update()
    run(["apt-get", "install", "-y", *PRINTING_PACKAGES])


def set_printing_enabled(enabled: bool) -> None:
    if enabled:
        run(["systemctl", "unmask", *PRINTING_UNITS])
        run(["systemctl", "enable", "--now", *PRINTING_AUTOSTART_UNITS])
    else:
        # Mask every activation path, including the static ipp-usb unit. Merely
        # stopping cups.service would allow its socket or path unit to revive it.
        run(["systemctl", "mask", "--now", *PRINTING_UNITS])


def resume_print_queues() -> None:
    result = subprocess.run(
        ["lpstat", "-p"], capture_output=True, text=True, check=False
    )
    if result.returncode and "No destinations added" not in result.stderr:
        raise RuntimeError(f"lpstat exited with status {result.returncode}")
    for line in result.stdout.splitlines():
        match = re.match(r"^printer\s+([A-Za-z0-9][A-Za-z0-9._-]*)\s+(.+)$", line)
        if match and " disabled " in f" {match.group(2).lower()} ":
            run(["cupsenable", match.group(1)])


def main(arguments: list[str]) -> int:
    if os.geteuid() != 0:
        print("This helper must be run through polkit.", file=sys.stderr)
        return 77
    try:
        match arguments:
            case ["intel-inspect"]:
                from anduinos_driver_center.intel_graphics import inspect
                print(json.dumps(inspect().report()))
                return 0
            case ["intel-apply"]:
                from anduinos_driver_center.intel_graphics_settings import apply
                payload = sys.stdin.read(16385)
                if len(payload) > 16384:
                    raise ValueError("Intel graphics request is too large")
                apply(json.loads(payload))
            case ["intel-reset"]:
                from anduinos_driver_center.intel_graphics_settings import apply
                apply(None)
            case ["install", package]:
                install_driver(package)
            case ["install-recommended"]:
                install_recommended()
            case ["update", package]:
                update_driver(package)
            case ["refresh-driver-info"]:
                refresh_driver_info()
            case ["repair-nvidia", package]:
                repair_nvidia_driver(package)
            case ["install-xbox"]:
                install_xbox()
            case ["repair-xbox"]:
                install_xbox(reinstall=True)
            case ["install-audio"]:
                install_audio()
            case ["repair-audio"]:
                install_audio(reinstall=True)
            case ["install-printing-support"]:
                install_printing_support()
            case ["set-printing-enabled", "true"]:
                set_printing_enabled(True)
            case ["set-printing-enabled", "false"]:
                set_printing_enabled(False)
            case ["resume-print-queues"]:
                resume_print_queues()
            case _:
                print("Unknown or incomplete action.", file=sys.stderr)
                return 64
    except (OSError, ValueError, RuntimeError, SystemError) as error:
        print(f"Driver operation failed: {error}", file=sys.stderr)
        return 1
    print("Driver operation completed successfully.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))
