#!/usr/bin/env python3
# Nvidia driver builder for Porteus
# Author: Blaze

import os, sys, subprocess, re, tempfile, shutil, signal, time, fnmatch, urllib.request, threading
import gi
gi.require_version('Gtk', '3.0')
from gi.repository import Gtk, GLib, GdkPixbuf

NVIDIA_BASE_64 = 'https://download.nvidia.com/XFree86/Linux-x86_64'
NVIDIA_US_BASE_64 = 'https://us.download.nvidia.com/XFree86/Linux-x86_64'
NVIDIA_BASE_32 = 'https://download.nvidia.com/XFree86/Linux-x86'
NVIDIA_US_BASE_32 = 'https://us.download.nvidia.com/XFree86/Linux-x86'

IMGS = '/mnt/live/memory/images'
PKGDIR = '/var/lib/pkgtools/packages'
ICON_NAMES = ('nvidia', 'video-display', 'preferences-desktop-display', 'computer')
VER_RE = r'\d+(?:\.\d+)+'

T = {
    'en': {
        'title': 'Nvidia driver builder for Porteus',
        'desc': 'Select a file or download a version.\nThe file will be converted to a Porteus module.',
        'file:': 'Select file:', 'dl': 'Download latest version', 'browse': 'Browse',
        'cancel': 'Cancel', 'ok': 'OK', 'close': 'Close', '_cancel': '_Cancel', '_open': '_Open',
        'no_opt': 'Select file or download!', 'no_file': 'File not found',
        'all': 'All files', 'f_run': 'Nvidia files (*.run)',
        'no_ver': 'Could not determine version from filename.\nRename file to include version, e.g.:\n  NVIDIA-Linux-x86_64-550.120.run',
        'bad_ver': 'Invalid driver version.', 'no_link': 'Could not get link',
        'downloading': 'Downloading', 'building': 'Building', 'preparing': 'Preparing environment...',
        'mkmod': 'Creating XZM module...', 'ready': 'Module ready:', 'fail': 'Build failed',
        'mv_act': 'Move and activate', 'open_dir': 'Open folder',
        'moved': 'moved to', 'act_ok': 'successfully activated', 'act_ng': 'activation failed',
        'mv_fail': 'Move error', 'act_fail': 'Activation failed',
        'action': 'Choose action:', 'more': 'You can build more or close.',
        'no_moddir': 'Modules folder not found', 'low_space': 'Low space in',
        'need': 'Need', 'free': 'Free', 'root': 'Root required!', 'x11': 'X session required!',
        'sel_ver': 'Select version', 'version': 'Version',
        'installed': 'Installed Nvidia version', 'na': 'n/a',
        'dl_fail': 'Download failed', 'show_all': 'Show all versions',
        'no_src': 'Kernel sources not found.\nPlease get and activate module:\n  06-crippled_sources-*.xzm\n(available on Porteus server or USB)',
        'no_tools': 'Build tools not found (gcc/make).\nPlease get and activate module:\n  05-devel-*.xzm\n(available on Porteus server or USB)',
        'no_multilib_tools': '32-bit build tools not found.\nPlease install multilib support (gcc-multilib) or disable 32-bit libs option.',
        'stage_fail': 'Staging build failed. The live system was not modified.',
        'extracting': 'Extracting installer...', 'compiling': 'Compiling kernel modules...',
        'staging': 'Staging files...',
        'reboot_host': 'Recommended: reboot or run /etc/rc.d/rc.nvidia to load Nvidia modules.',
        'multilib': 'Include 32-bit libraries (for Wine/multilib)',
        'multilib_tooltip': 'Requires 64-bit system and 32-bit compiler (multilib)',
    },
    'ru': {
        'title': 'Сборка драйвера Nvidia для Porteus',
        'desc': 'Выберите файл или загрузите версию.\nФайл будет преобразован в модуль Porteus.',
        'file:': 'Выберите файл:', 'dl': 'Загрузить последнюю версию', 'browse': 'Указать файл',
        'cancel': 'Отмена', 'ok': 'ОК', 'close': 'Закрыть', '_cancel': '_Отмена', '_open': '_Открыть',
        'no_opt': 'Выберите файл или загрузку!', 'no_file': 'Файл не найден',
        'all': 'Все файлы', 'f_run': 'Файлы Nvidia (*.run)',
        'no_ver': 'Не удалось определить версию из имени файла.\nПереименуйте файл, включив версию, например:\n  NVIDIA-Linux-x86_64-550.120.run',
        'bad_ver': 'Некорректная версия драйвера.', 'no_link': 'Не удалось получить ссылку',
        'downloading': 'Загрузка', 'building': 'Сборка', 'preparing': 'Подготовка окружения...',
        'mkmod': 'Создание модуля XZM...', 'ready': 'Модуль готов:', 'fail': 'Сборка не удалась',
        'mv_act': 'Переместить и активировать', 'open_dir': 'Открыть папку',
        'moved': 'перемещён в', 'act_ok': 'успешно активирован', 'act_ng': 'активация не удалась',
        'mv_fail': 'Ошибка перемещения', 'act_fail': 'Ошибка активации',
        'action': 'Выберите действие:', 'more': 'Можно собрать ещё или закрыть.',
        'no_moddir': 'Папка модулей не найдена', 'low_space': 'Мало места в',
        'need': 'Нужно', 'free': 'Свободно', 'root': 'Нужен root!', 'x11': 'Нужна X-сессия!',
        'sel_ver': 'Выбор версии', 'version': 'Версия',
        'installed': 'Установленная версия Nvidia', 'na': 'н/д',
        'dl_fail': 'Ошибка загрузки', 'show_all': 'Показывать все версии',
        'no_src': 'Исходники ядра не найдены.\nПолучите и активируйте модуль:\n  06-crippled_sources-*.xzm\n(доступен на сервере Porteus или USB)',
        'no_tools': 'Инструменты сборки не найдены (gcc/make).\nПолучите и активируйте модуль:\n  05-devel-*.xzm\n(доступен на сервере Porteus или USB)',
        'no_multilib_tools': '32-битные инструменты сборки не найдены.\nУстановите поддержку multilib или отключите опцию 32-битных библиотек.',
        'stage_fail': 'Изолированная сборка не удалась. Живая система не была изменена.',
        'extracting': 'Извлечение установщика...', 'compiling': 'Компиляция модулей ядра...',
        'staging': 'Подготовка файлов...',
        'reboot_host': 'Рекомендуется перезагрузка или запуск /etc/rc.d/rc.nvidia.',
        'multilib': '32-битные библиотеки (для Wine/multilib)',
        'multilib_tooltip': 'Требуется 64-битная система и 32-битный компилятор (multilib)',
    }
}

def _(k): return T.get(os.environ.get('LANG', 'en')[:2], T['en']).get(k, k)

def safe_rm(p):
    try: os.remove(p)
    except: pass

def cmd(c, t=60, env=None, input_data=None):
    try:
        kw = dict(capture_output=True, text=True, errors='replace', timeout=t, env=env)
        if input_data is None: kw['stdin'] = subprocess.DEVNULL
        else: kw['input'] = input_data
        return subprocess.run(c, **kw)
    except subprocess.TimeoutExpired:
        return subprocess.CompletedProcess(c, 124, '', f'Timeout after {t}s')
    except Exception as e:
        return subprocess.CompletedProcess(c, 1, '', str(e))

def load_icon(size=48):
    try:
        theme = Gtk.IconTheme.get_default()
        for n in ICON_NAMES:
            try:
                pb = theme.load_icon(n, size, 0)
                if pb: return pb
            except: pass
    except: pass
    return None

def clean_ver(s):
    m = re.fullmatch(VER_RE, (s or '').strip())
    return m.group(0) if m else None

def pver():
    try:
        with open('/etc/porteus-version') as f:
            m = re.search(r'(\d+\.\d+(?:\.\d+)?)', f.read().strip())
        return f'v{m.group(1)}' if m else 'unknown'
    except: return 'unknown'

def is_nvidia_file(f):
    f = os.path.basename(f)
    return f.lower().startswith('nvidia-linux-') and f.endswith('.run')

def find_nvidia_mods():
    r = []
    for d in (IMGS, os.environ.get('MODDIR', '')):
        if d and os.path.isdir(d):
            try: r.extend(os.path.join(d, f) for f in os.listdir(d) if is_nvidia_file(f))
            except: pass
    return r

def nvidia_in_images():
    try: return os.path.isdir(IMGS) and any(is_nvidia_file(f) for f in os.listdir(IMGS))
    except: return False

def deactivate_silent(m):
    if os.path.isfile(m): cmd(['deactivate', m, '-q'], t=30)

def deactivate_all_wait():
    for m in find_nvidia_mods(): deactivate_silent(m)
    for _ in range(15):
        if not nvidia_in_images(): break
        time.sleep(1)

def pkg_ver():
    if not os.path.isdir(PKGDIR): return None, None
    try:
        for f in os.listdir(PKGDIR):
            m = re.match(r'nvidia-driver-(\d+(?:\.\d+)+)', f)
            if m: return m.group(1), f
    except: pass
    return None, None

def space_ok(p, mb=1500):
    try:
        s = os.statvfs(p)
        free = s.f_bavail * s.f_frsize // 1048576
        return free >= mb, free
    except: return True, 0

def fm_open(path):
    try:
        subprocess.Popen(['xdg-open', path], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                         stdin=subprocess.DEVNULL, start_new_session=True)
        return True
    except: return False

def mod_dir():
    d = os.environ.get('MODDIR', '')
    if d and os.path.isdir(d): return d
    for p in (IMGS, '/porteus/modules', '/mnt/sda1/porteus/modules'):
        if os.path.isdir(p): return p
    return None

def nvidia_version_lists(arch='x86_64'):
    latest, all_vs = [], []
    base = NVIDIA_BASE_64 if arch == 'x86_64' else NVIDIA_BASE_32
    for attempt in range(2):
        try:
            html = urllib.request.urlopen(f'{base}/', timeout=30).read().decode()
            vs = [v for v in set(re.findall(r'href=["\']?(?:[^"\'>]*/)?(\d+(?:\.\d+)+)/["\']?', html, re.IGNORECASE)) if clean_ver(v)]
            all_vs = sorted(vs, key=lambda v: [int(x) for x in v.split('.')], reverse=True)
            if all_vs: latest = [all_vs[0]]
            break
        except:
            if attempt == 0: time.sleep(2)
    if not latest:
        try:
            txt = urllib.request.urlopen(f'{base}/latest.txt', timeout=20).read().decode()
            m = re.search(VER_RE, txt)
            if m: latest = [clean_ver(m.group(0))]
        except: pass
    return latest, all_vs

def nvidia_urls(v, is64):
    if is64:
        fn = f'NVIDIA-Linux-x86_64-{v}.run'
        return [f'{NVIDIA_US_BASE_64}/{v}/{fn}', f'{NVIDIA_BASE_64}/{v}/{fn}']
    fn = f'NVIDIA-Linux-x86-{v}.run'
    return [f'{NVIDIA_US_BASE_32}/{v}/{fn}', f'{NVIDIA_BASE_32}/{v}/{fn}']

def inst_ver():
    try:
        r = cmd(['nvidia-smi', '--query-gpu=driver_version', '--format=csv,noheader'], t=5)
        if r.returncode == 0 and r.stdout.strip(): return r.stdout.strip().split('\n')[0]
    except: pass
    try:
        r = cmd(['cat', '/proc/driver/nvidia/version'], t=5)
        if r.returncode == 0:
            m = re.search(r'(\d+(?:\.\d+)+)', r.stdout)
            if m: return m.group(1)
    except: pass
    v, _ = pkg_ver(); return v

def inst_short():
    f = inst_ver(); m = re.match(r'(\d+(?:\.\d+)+)', f) if f else None
    return m.group(1) if m else None

def find_file(root, name):
    for dp, dn, fn in os.walk(root):
        if name in fn: return os.path.join(dp, name)
    return None

def find_files_pattern(root, pattern):
    return [os.path.join(dp, f) for dp, dn, fn in os.walk(root)
            for f in fn if fnmatch.fnmatch(f.lower(), pattern.lower())]

def strip_unwanted(path):
    bad_dirs = {'sdk', 'doc', 'docs', 'manual', 'src', 'source', 'sources', 'dtrace',
                'additions', 'unattendedtemplates', 'init', 'installer', 'temp', 'nv-extract'}
    bad_files = ('*.cmd', '*.o', '*.mod', '*.mod.c', '*.order', '*.symvers', '*.c', '*.h',
                 '*.cpp', '*.cc', '*.s', '*.S', 'Makefile', 'Makefile.*', 'Kbuild', '*.debug',
                 '*.pdb', '*.sym', '*.map', '*.a', '*.la', '*.spec', 'UserManual.pdf',
                 'UserManual.qch', 'UserManual.qhc', 'postinst-common.sh', 'prerm-common.sh',
                 'routines.sh', 'check_module_dependencies.sh', 'uninstall.sh', 'nvidia-uninstall',
                 '.autoreg', '*.log', 'nvidia-installer')
    for root, dirs, files in os.walk(path, topdown=False):
        for d in list(dirs):
            if d.lower() in bad_dirs:
                shutil.rmtree(os.path.join(root, d), ignore_errors=True); dirs.remove(d)
        for f in files:
            if any(fnmatch.fnmatch(f.lower(), p.lower()) for p in bad_files):
                safe_rm(os.path.join(root, f))
        try:
            if not os.listdir(root): os.rmdir(root)
        except: pass

class Prog(Gtk.Dialog):
    def __init__(self, par, t, m):
        super().__init__(title=t, transient_for=par, modal=True)
        self.set_deletable(False); self.set_default_size(340, 110)
        self.set_position(Gtk.WindowPosition.CENTER_ON_PARENT)
        b = self.get_content_area(); b.set_spacing(12)
        b.set_margin_top(20); b.set_margin_bottom(20)
        try: b.set_margin_start(24); b.set_margin_end(24)
        except: b.set_margin_left(24); b.set_margin_right(24)
        self.lbl = Gtk.Label(label=m); self.lbl.set_line_wrap(True)
        self.bar = Gtk.ProgressBar(); self.bar.set_pulse_step(0.08)
        b.pack_start(self.lbl, False, False, 0)
        b.pack_start(self.bar, False, False, 0)
        self.show_all()
        self._tid = GLib.timeout_add(80, self._tk); self._d = False
    def _tk(self):
        if not self._d: self.bar.pulse(); return True
        return False
    def msg(self, m):
        if not self._d: self.lbl.set_text(m)
    def close(self):
        if not self._d:
            self._d = True
            if self._tid: GLib.source_remove(self._tid)
            self.destroy()

class VerDlg(Gtk.Dialog):
    def __init__(self, par, latest, all_vs, iv):
        super().__init__(title=_('sel_ver'), transient_for=par, modal=True)
        self.set_default_size(360, 430)
        self.set_position(Gtk.WindowPosition.CENTER_ON_PARENT)
        b = self.get_content_area(); b.set_spacing(8)
        b.set_margin_top(12); b.set_margin_bottom(12)
        try: b.set_margin_start(16); b.set_margin_end(16)
        except: b.set_margin_left(16); b.set_margin_right(16)
        il = Gtk.Label(); il.set_markup(f'<b>{_("installed")}: {iv or _("na")}</b>')
        il.set_halign(Gtk.Align.START); b.pack_start(il, False, False, 0)
        self.latest, self.all = latest, all_vs
        self.chk_all = Gtk.CheckButton(label=_('show_all'))
        self.chk_all.connect('toggled', self._fill)
        b.pack_start(self.chk_all, False, False, 0)
        sw = Gtk.ScrolledWindow(); sw.set_policy(Gtk.PolicyType.NEVER, Gtk.PolicyType.AUTOMATIC)
        self.st = Gtk.ListStore(str)
        self.tv = Gtk.TreeView(model=self.st)
        self.tv.append_column(Gtk.TreeViewColumn(_('version'), Gtk.CellRendererText(), text=0))
        self.tv.set_headers_visible(True)
        self.tv.get_selection().set_mode(Gtk.SelectionMode.SINGLE)
        sw.add(self.tv); b.pack_start(sw, True, True, 0)
        self._fill()
        self.add_button(_('_cancel'), Gtk.ResponseType.CANCEL)
        self.add_button(_('ok'), Gtk.ResponseType.OK)
        self.show_all()
    def _fill(self, *a):
        self.st.clear()
        for v in (self.all if self.chk_all.get_active() else self.latest):
            self.st.append([v])
        if len(self.st): self.tv.get_selection().select_path(Gtk.TreePath.new_first())
    def chosen(self):
        _, it = self.tv.get_selection().get_selected()
        return self.st[it][0] if it else None

class App(Gtk.Window):
    def __init__(self, cli_file=None):
        super().__init__(title=_('title'))
        self._icon = load_icon(48)
        if self._icon: self.set_icon(self._icon)
        self.set_border_width(12); self.set_resizable(True)
        self.set_default_size(640, 360)
        self.set_position(Gtk.WindowPosition.CENTER)
        self.connect('destroy', Gtk.main_quit)
        raw = os.uname().machine
        self.arch = 'i586' if raw in ('i386', 'i486', 'i586', 'i686') else raw
        self.is64 = self.arch == 'x86_64'
        self.nvidia_libdir = 'lib64' if self.is64 else 'lib'
        self.multilib_enabled = False
        self._prog = None
        self._ui()
        if cli_file and os.path.isfile(cli_file): self._set_path(cli_file)

    def _ui(self):
        vb = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=12); self.add(vb)
        hb = Gtk.Box(spacing=12)
        if self._icon: hb.pack_start(Gtk.Image.new_from_pixbuf(self._icon), False, False, 0)
        t = Gtk.Label(); t.set_markup(f'<span size="x-large" weight="bold">{_("title")}</span>')
        hb.pack_start(t, True, True, 0)
        vb.pack_start(hb, False, False, 0)
        vb.pack_start(Gtk.Separator(), False, False, 0)
        d = Gtk.Label(label=_('desc')); d.set_line_wrap(True); d.set_max_width_chars(80)
        d.set_halign(Gtk.Align.START); vb.pack_start(d, False, False, 0)
        fh = Gtk.Box(spacing=6)
        fh.pack_start(Gtk.Label(label=_('file:')), False, False, 0)
        self.ent = Gtk.Entry(); self.ent.set_editable(False); self.ent.set_width_chars(52)
        fh.pack_start(self.ent, True, True, 0)
        bb = Gtk.Button(label=_('browse')); bb.connect('clicked', self._br)
        fh.pack_start(bb, False, False, 0)
        vb.pack_start(fh, False, False, 0)
        self.chk = Gtk.CheckButton(label=_('dl'))
        self.chk.connect('toggled', lambda w: self.ent.set_sensitive(not w.get_active()))
        vb.pack_start(self.chk, False, False, 0)
        self.chk_multilib = Gtk.CheckButton(label=_('multilib'))
        self.chk_multilib.set_tooltip_text(_('multilib_tooltip'))
        self.chk_multilib.set_sensitive(self.is64)
        vb.pack_start(self.chk_multilib, False, False, 0)
        vb.pack_start(Gtk.Separator(), False, False, 4)
        bh = Gtk.Box(spacing=8, halign=Gtk.Align.END)
        bc = Gtk.Button(label=_('cancel')); bc.connect('clicked', lambda w: Gtk.main_quit())
        bo = Gtk.Button(label=_('ok')); bo.connect('clicked', self._ok)
        bh.pack_end(bo, False, False, 0)
        bh.pack_end(bc, False, False, 0)
        vb.pack_end(bh, False, False, 0)
        self.show_all()

    def _set_path(self, p):
        p = p or ''
        self.ent.set_text(p); self.ent.set_tooltip_text(p)

    def _br(self, w):
        d = Gtk.FileChooserDialog(title=_('browse'), parent=self, action=Gtk.FileChooserAction.OPEN)
        d.add_buttons(_('_cancel'), Gtk.ResponseType.CANCEL, _('_open'), Gtk.ResponseType.OK)
        fa = Gtk.FileFilter(); fa.set_name(_('all')); fa.add_pattern('*'); d.add_filter(fa)
        ff = Gtk.FileFilter(); ff.set_name(_('f_run')); ff.add_pattern('*.run')
        d.add_filter(ff); d.set_filter(ff); d.set_current_folder('/mnt')
        if d.run() == Gtk.ResponseType.OK:
            self._set_path(d.get_filename()); self.chk.set_active(False)
        d.destroy()

    def _ok(self, w):
        if self.chk.get_active(): return self._pick()
        p = self.ent.get_text().strip()
        if not p: self._err(_('no_opt'))
        elif not os.path.exists(p): self._err(_('no_file'))
        else: self._set_path(p); self._build(p, False)

    def _msg(self, mt, tx):
        def f():
            bt = Gtk.ButtonsType.OK if mt == Gtk.MessageType.INFO else Gtk.ButtonsType.CLOSE
            d = Gtk.MessageDialog(transient_for=self, modal=True, message_type=mt, buttons=bt, text=tx)
            d.run(); d.destroy(); return False
        GLib.idle_add(f)

    def _err(self, m): self._msg(Gtk.MessageType.ERROR, m)
    def _info(self, m): self._msg(Gtk.MessageType.INFO, m)

    def _po(self, t, m):
        def f():
            if self._prog: self._prog.close()
            self._prog = Prog(self, t, m); return False
        GLib.idle_add(f)

    def _pm(self, m): GLib.idle_add(lambda: (self._prog and self._prog.msg(m)) or False)

    def _pc(self):
        def f():
            if self._prog: self._prog.close(); self._prog = None
            return False
        GLib.idle_add(f)

    def _thread(self, fn, *args): threading.Thread(target=fn, args=args, daemon=True).start()

    def _pick(self):
        self._po(_('downloading'), _('downloading'))
        def w():
            latest, all_vs = nvidia_version_lists(self.arch)
            GLib.idle_add(self._after_pick, latest, all_vs, inst_short())
        self._thread(w)

    def _after_pick(self, latest, all_vs, iv):
        self._pc()
        if latest or all_vs: self._verdlg(latest, all_vs, iv)
        else: self._err(_('no_link'))

    def _verdlg(self, latest, all_vs, iv):
        d = VerDlg(self, latest, all_vs, iv)
        r = d.run(); v = d.chosen(); d.destroy()
        if r == Gtk.ResponseType.OK and v: self._dlver(v)

    def _dlver(self, ver):
        ver = clean_ver(ver)
        if not ver: self._err(_('bad_ver')); return
        self._po(_('downloading'), _('downloading'))
        def w():
            fn = f'NVIDIA-Linux-{"x86_64" if self.is64 else "x86"}-{ver}.run'
            dest = os.path.join('/tmp', fn)
            dl, last_err = False, None
            if not (os.path.exists(dest) and os.path.getsize(dest) > 1048576):
                self._pm(f'{_("downloading")} {fn}...')
                for url in nvidia_urls(ver, self.is64):
                    try:
                        with urllib.request.urlopen(url, timeout=600) as r, open(dest, 'wb') as f:
                            shutil.copyfileobj(r, f)
                        dl = True; break
                    except Exception as e:
                        safe_rm(dest); last_err = e
                if not dl:
                    GLib.idle_add(self._err, f'{_("dl_fail")}: {last_err}')
                    GLib.idle_add(self._pc); return
            GLib.idle_add(self._pc)
            GLib.idle_add(self._set_path, dest)
            GLib.idle_add(self._build, dest, dl)
        self._thread(w)

    def _build(self, fp, dl=False):
        self.multilib_enabled = bool(self.chk_multilib.get_active()) and self.is64
        self._po(_('building'), _('preparing'))
        def w():
            try: self._do(fp, dl)
            except Exception as e:
                GLib.idle_add(self._pc)
                GLib.idle_add(self._err, f'{_("fail")}: {e}')
        self._thread(w)

    def _file_version(self, fn):
        m = re.search(r'NVIDIA-Linux-(?:x86_64|x86|i586|i686|i386)-(\d+(?:\.\d+)+(?:[-_.][^\s/\\]+)?)', fn)
        if m:
            clean = re.search(r'(\d+(?:\.\d+)+)', m.group(1))
            if clean: return clean.group(1)
        m = re.search(r'(\d+(?:\.\d+)+)', fn)
        return m.group(1) if m else None

    def _mki(self, bd):
        os.makedirs(os.path.join(bd, 'etc/rc.d/init.d'), exist_ok=True)
        os.makedirs(os.path.join(bd, 'etc/rc.d/rc4.d'), exist_ok=True)
        ip = os.path.join(bd, 'etc/rc.d/init.d/rc.nvidia')
        body = (
            '#!/bin/sh\n'
            '/sbin/depmod -a\n'
            '/sbin/modprobe nvidia\n'
            '/sbin/modprobe nvidia-modeset\n'
            '/sbin/modprobe nvidia-drm\n'
            '/sbin/modprobe nvidia-uvm\n'
            'if [ -x /usr/bin/nvidia-modprobe ] && [ ! -c /dev/nvidia0 ]; then\n'
            '    /usr/bin/nvidia-modprobe -c 0 -u 2>/dev/null\n'
            'fi\n'
            '/sbin/ldconfig\n'
        )
        open(ip, 'w').write(body); os.chmod(ip, 0o755)
        cmd(['ln', '-sf', '/etc/rc.d/init.d/rc.nvidia', os.path.join(bd, 'etc/rc.d/rc4.d', 'S99nvidia')])
        os.makedirs(os.path.join(bd, 'etc/modprobe.d'), exist_ok=True)
        with open(os.path.join(bd, 'etc/modprobe.d/nvidia.conf'), 'w') as f:
            f.write('blacklist nouveau\noptions nouveau modeset=0\noptions nvidia-drm modeset=1\n')
        os.makedirs(os.path.join(bd, 'etc/X11/xorg.conf.d'), exist_ok=True)
        with open(os.path.join(bd, 'etc/X11/xorg.conf.d/10-nvidia.conf'), 'w') as f:
            f.write(f'Section "OutputClass"\n    Identifier "nvidia"\n    MatchDriver "nvidia-drm"\n'
                    f'    Driver "nvidia"\n    Option "AllowEmptyInitialConfiguration"\n'
                    f'    Option "PrimaryGPU" "yes"\n'
                    f'    ModulePath "/usr/{self.nvidia_libdir}/xorg/modules"\nEndSection\n')
        os.makedirs(os.path.join(bd, 'etc/profile.d'), exist_ok=True)
        ps = os.path.join(bd, 'etc/profile.d/nvidia.sh')
        with open(ps, 'w') as f:
            f.write('#!/bin/sh\n[ -c /dev/nvidia0 ] || /etc/rc.d/init.d/rc.nvidia 2>/dev/null\n')
        os.chmod(ps, 0o755)

    def _do(self, fp, dl=False):
        need_mb = 2000 if self.multilib_enabled else 1500
        ok, free = space_ok('/tmp', need_mb)
        if not ok:
            self._pc(); self._err(f'{_("low_space")} /tmp\n{_("need")}: {need_mb} MB, {_("free")}: {free} MB'); return
        if not os.path.exists('/usr/src/linux'):
            self._pc(); self._err(_('no_src')); return
        if not shutil.which('gcc') or not shutil.which('make'):
            self._pc(); self._err(_('no_tools')); return
        if self.multilib_enabled:
            has_32bit_gcc, test_c = False, None
            test_out = '/tmp/.porteus_32bit_test'
            try:
                with tempfile.NamedTemporaryFile(suffix='.c', mode='w', delete=False) as tf:
                    tf.write('int main() { return 0; }'); test_c = tf.name
                r = cmd(['gcc', '-m32', '-x', 'c', test_c, '-o', test_out], t=30)
                if r.returncode == 0 and os.path.exists(test_out): has_32bit_gcc = True
            finally:
                if test_c: safe_rm(test_c)
                safe_rm(test_out)
            if not has_32bit_gcc:
                self._pc(); self._err(_('no_multilib_tools')); return
        fv = self._file_version(os.path.basename(fp))
        if not fv:
            self._pc(); self._err(_('no_ver')); return
        self._do_staging(fp, dl, fv)

    def _ensure_critical_files(self, bd, extract_dir, kpath, fv):
        usr_lib = os.path.join(bd, 'usr', self.nvidia_libdir)
        drv_dir = os.path.join(usr_lib, 'xorg', 'modules', 'drivers')
        os.makedirs(drv_dir, exist_ok=True)
        drv = os.path.join(drv_dir, 'nvidia_drv.so')
        if not os.path.exists(drv):
            src = find_file(extract_dir, 'nvidia_drv.so')
            if src:
                shutil.copy2(src, drv); os.chmod(drv, 0o755)
        ext_dir = os.path.join(usr_lib, 'xorg', 'modules', 'extensions')
        os.makedirs(ext_dir, exist_ok=True)
        copied = set()
        for ext in find_files_pattern(extract_dir, 'libglxserver_nvidia*'):
            if os.path.realpath(ext).startswith(os.path.realpath(extract_dir)):
                n = os.path.basename(ext)
                if n not in copied:
                    shutil.copy2(ext, ext_dir); copied.add(n)
        mp = os.path.join(bd, 'usr', 'bin', 'nvidia-modprobe')
        if not os.path.exists(mp):
            src = find_file(extract_dir, 'nvidia-modprobe')
            if src:
                os.makedirs(os.path.dirname(mp), exist_ok=True)
                shutil.copy2(src, mp); os.chmod(mp, 0o4755)
        json_map = {
            'nvidia_icd.json': 'usr/share/vulkan/icd.d',
            '10_nvidia.json': 'usr/share/glvnd/egl_vendor.d',
            '10_nvidia_wayland.json': 'usr/share/egl/egl_external_platform.d',
            '10_nvidia_drm.json': 'usr/share/egl/egl_external_platform.d',
        }
        for jf in find_files_pattern(extract_dir, '*.json'):
            n = os.path.basename(jf)
            if n in json_map:
                d = os.path.join(bd, json_map[n])
                if not os.path.exists(os.path.join(d, n)):
                    os.makedirs(d, exist_ok=True)
                    shutil.copy2(jf, d)
        fw_dir = os.path.join(bd, 'lib', 'firmware', 'nvidia', fv)
        for fw in find_files_pattern(extract_dir, 'gsp_*.bin'):
            os.makedirs(fw_dir, exist_ok=True)
            dst = os.path.join(fw_dir, os.path.basename(fw))
            if not os.path.exists(dst): shutil.copy2(fw, dst)
        gbm_src = os.path.join(bd, 'usr', 'lib', 'gbm', 'nvidia-drm_gbm.so')
        if self.nvidia_libdir == 'lib64' and os.path.exists(gbm_src):
            gbm_dst = os.path.join(usr_lib, 'gbm')
            os.makedirs(gbm_dst, exist_ok=True)
            if not os.path.exists(os.path.join(gbm_dst, 'nvidia-drm_gbm.so')):
                shutil.copy2(gbm_src, gbm_dst)
        for ko in find_files_pattern(bd, '*.ko*'):
            dest = os.path.join(kpath, os.path.basename(ko))
            if not os.path.exists(dest) and os.path.dirname(ko) != kpath:
                try: shutil.move(ko, dest)
                except: pass

    def _do_staging(self, fp, dl, fv):
        with tempfile.TemporaryDirectory(dir='/tmp', prefix='nvidia-stage-') as td:
            bd = os.path.join(td, 'build'); os.makedirs(bd, exist_ok=True)
            try:
                self._pm(_('extracting'))
                exd = os.path.join(td, 'nv-extract')
                os.chmod(fp, 0o755)
                r = cmd([fp, '--extract-only', '--target', exd], t=600)
                if r.returncode != 0 or not os.path.isdir(exd):
                    self._pc(); self._err(_('stage_fail')); return
                inst = find_file(exd, 'nvidia-installer')
                if not inst:
                    self._pc(); self._err(_('stage_fail')); return
                os.chmod(inst, 0o755)
                self._pm(_('compiling'))
                kver = os.uname().release
                kpath = os.path.join(bd, 'lib', 'modules', kver, 'kernel', 'drivers', 'video')
                os.makedirs(kpath, exist_ok=True)
                usr = os.path.join(bd, 'usr'); os.makedirs(usr, exist_ok=True)
                env = os.environ.copy()
                env['PATH'] = '/usr/sbin:/usr/bin:/sbin:/bin'
                args = [
                    inst, '--silent', '--no-x-check', '--no-nouveau-check', '--no-abi-note',
                    '--kernel-source-path=/usr/src/linux', '--skip-module-load',
                    '--no-distro-scripts', '--no-rpms', '--no-systemd',
                    '--skip-depmod', '--no-rebuild-initramfs',
                    f'--utility-prefix={usr}', f'--opengl-prefix={usr}', f'--x-prefix={usr}',
                    f'--opengl-libdir={self.nvidia_libdir}', f'--utility-libdir={self.nvidia_libdir}',
                    f'--x-module-path={self.nvidia_libdir}/xorg/modules',
                    f'--x-library-path={self.nvidia_libdir}',
                    f'--kernel-install-path={kpath}'
                ]
                if self.is64 and self.multilib_enabled:
                    args.extend([f'--compat32-prefix={usr}', '--compat32-libdir=lib'])
                args.extend([
                    f'--wine-prefix={usr}',
                    f'--documentation-prefix={usr}',
                    f'--xdg-data-dir={usr}/share',
                    f'--gbm-backend-dir={self.nvidia_libdir}/gbm',
                    f'--glvnd-egl-config-path={usr}/share/glvnd/egl_vendor.d',
                    f'--egl-external-platform-config-path={usr}/share/egl/egl_external_platform.d',
                ])
                r = cmd(args, t=3600, env=env)
                if r.returncode != 0:
                    self._pc(); self._err(f'{_("stage_fail")}: {r.stderr[:500] if r.stderr else ""}'); return
                self._pm(_('staging'))
                self._ensure_critical_files(bd, exd, kpath, fv)
                mp = find_file(bd, 'nvidia-modprobe')
                if mp: os.chmod(mp, 0o4755)
                self._mki(bd); strip_unwanted(bd)
                self._package(bd, fv, dl, fp)
            finally:
                for f in ('/tmp/.message', '/tmp/.cdload'): safe_rm(f)

    def _package(self, bd, fv, dl=False, fp=None):
        if not find_files_pattern(bd, 'nvidia.ko*'):
            self._pc(); self._err(_('fail')); return False
        self._pm(_('mkmod'))
        pn = 'nvidia-driver'
        suf = 'x86_64-multilib' if self.multilib_enabled else self.arch
        out = f'/tmp/{pn}-{fv}-porteus-{pver()}-{suf}-1prt.xzm'
        safe_rm(out)
        pd = os.path.join(bd, 'var/lib/pkgtools/packages'); os.makedirs(pd, exist_ok=True)
        nm = f'{pn}-{fv}-porteus-{pver()}-{suf}-1prt'
        with open(os.path.join(pd, nm), 'w') as f:
            f.write(f'PACKAGE NAME: {nm}\nPACKAGE DESCRIPTION:\n'
                    f'{pn}: {pn} (proprietary graphics driver)\n{pn}:\n'
                    f'{pn}: Nvidia proprietary driver for X.org and Wayland.\n'
                    f'{pn}: Includes kernel modules and user-space libraries.\n')
            if self.multilib_enabled: f.write(f'{pn}: (32-bit libs included for Wine/Steam)\n')
            f.write(f'{pn}:\n{pn}: https://www.nvidia.com/\n{pn}:\nFILE LIST:\n')
            for root, dirs, files in os.walk(bd):
                for fn in files:
                    rel = os.path.relpath(os.path.join(root, fn), bd)
                    if not rel.startswith('var/lib/pkgtools/packages'): f.write(rel + '\n')
        r = cmd(['dir2xzm', bd, out], t=600)
        if r.returncode != 0:
            err = (r.stderr or '').strip() or (r.stdout or '').strip() or _('fail')
            self._pc(); self._err(f'{_("fail")}: {err}'); return False
        self._pc()
        if dl and fp: safe_rm(fp)
        GLib.idle_add(self._result, out)
        return True

    def _result(self, out):
        d = Gtk.MessageDialog(transient_for=self, modal=True, message_type=Gtk.MessageType.INFO,
                              buttons=Gtk.ButtonsType.NONE, text=f'{_("ready")}\n{out}')
        d.set_property('secondary-text', f'{_("action")}\n{_("more")}')
        d.add_button(_('mv_act'), 1); d.add_button(_('open_dir'), 2); d.add_button(_('close'), 3)
        r = d.run(); d.destroy()
        if r == 1: self._thread(self._move_act, out)
        elif r == 2: fm_open('/tmp')
        return False

    def _move_act(self, out):
        md = mod_dir()
        if not md: GLib.idle_add(self._err, _('no_moddir')); return
        try:
            mn, dest = os.path.basename(out), os.path.join(md, os.path.basename(out))
            deactivate_all_wait()
            if not os.path.isfile(out):
                GLib.idle_add(self._err, f'{_("mv_fail")}: {out}'); return
            GLib.idle_add(self._po, _('building'), _('mv_act'))
            shutil.move(out, dest); os.sync()
            r = cmd(['activate', dest], t=60)
            GLib.idle_add(self._pc)
            if r.returncode == 0:
                GLib.idle_add(self._info, f'{mn}\n{_("moved")} {md}\n{_("act_ok")}\n\n{_("reboot_host")}')
            else:
                GLib.idle_add(self._err, f'{_("act_fail")}\n{(r.stderr or "").strip() or (r.stdout or "").strip() or _("act_ng")}')
        except Exception as e:
            GLib.idle_add(self._pc)
            GLib.idle_add(self._err, f'{_("mv_fail")}: {e}')

if __name__ == '__main__':
    if os.getuid() != 0: print(_('root')); sys.exit(1)
    if not os.environ.get('DISPLAY'): print(_('x11')); sys.exit(1)
    pix = load_icon(48)
    if pix:
        try: Gtk.Window.set_default_icon(pix)
        except: pass
    signal.signal(signal.SIGINT, lambda *_: Gtk.main_quit())
    GLib.set_prgname('nvidia-builder')
    GLib.set_application_name(_('title'))
    fp = sys.argv[1] if len(sys.argv) > 1 and os.path.isfile(sys.argv[1]) else None
    App(fp).show_all()
    Gtk.main()