#!/usr/bin/env python3
# Govibe Disk Power Management
# Safe disk standby, unmount, power off, rescan, automount, and mount helper for Ubuntu/Linux.

import gi
import json
import os
import shlex
import subprocess
import sys
import threading
import time
from datetime import datetime

gi.require_version('Gtk', '3.0')
from gi.repository import Gtk, Gdk, GLib, Pango

APP_NAME = "Govibe Disk Power Management"
APP_VERSION = "1.8.1"

LSBLK_COLUMNS = "NAME,PATH,TYPE,SIZE,FSTYPE,MOUNTPOINTS,MODEL,SERIAL,TRAN,RM,ROTA,HOTPLUG,STATE,VENDOR"
AUTO_TICK_SECONDS = 30
C_POWER_HELPER = "/usr/lib/govibe-disk-power-management/gv-disk-power-state"
STATE_DIR = os.path.join(os.path.expanduser("~"), ".local", "state", "govibe-disk-power-management")
STATE_FILE = os.path.join(STATE_DIR, "state.json")
ADMIN_HELPER = "/usr/lib/govibe-disk-power-management/gv-disk-admin-helper"
AUTOSTART_DIR = os.path.join(os.path.expanduser("~"), ".config", "autostart")
AUTOSTART_FILE = os.path.join(AUTOSTART_DIR, "govibe-disk-power-management.desktop")


def run_cmd(args, timeout=20):
    try:
        proc = subprocess.run(args, text=True, capture_output=True, timeout=timeout)
        return proc.returncode, proc.stdout.strip(), proc.stderr.strip()
    except subprocess.TimeoutExpired:
        return 124, "", "Command timed out"
    except Exception as exc:
        return 1, "", str(exc)


def has_bin(name):
    return subprocess.run(["/usr/bin/env", "bash", "-lc", f"command -v {shlex.quote(name)} >/dev/null 2>&1"], capture_output=True).returncode == 0


def safe_path(path):
    return isinstance(path, str) and path.startswith('/dev/') and all(c.isalnum() or c in '/_-.:' for c in path)


def flatten_mountpoints(mp):
    if mp is None:
        return []
    if isinstance(mp, list):
        return [str(x) for x in mp if x]
    if isinstance(mp, str):
        if not mp:
            return []
        return [x for x in mp.split('\n') if x]
    return []


def collect_mountpoints(node):
    points = []
    points.extend(flatten_mountpoints(node.get('mountpoints')))
    for child in node.get('children') or []:
        points.extend(collect_mountpoints(child))
    return points


def is_system_mount(mountpoint):
    if not mountpoint:
        return False
    protected_exact = {'/', '/boot', '/boot/efi', '/home', '/usr', '/var', '/opt', '/srv'}
    if mountpoint in protected_exact:
        return True
    protected_prefixes = ('/snap/', '/run/', '/sys/', '/proc/', '/dev/')
    return mountpoint.startswith(protected_prefixes)


def disk_is_system(node):
    return any(is_system_mount(mp) for mp in collect_mountpoints(node))


def disk_identity(node):
    serial = (node.get('serial') or '').strip()
    model = (node.get('model') or '').strip()
    vendor = (node.get('vendor') or '').strip()
    size = (node.get('size') or '').strip()
    path = (node.get('path') or '').strip()
    if serial:
        return 'serial:' + serial
    return 'path:' + '|'.join([path, vendor, model, size])


def read_sys_state(path):
    if not safe_path(path):
        return ''
    name = os.path.basename(path)
    state_path = f'/sys/class/block/{name}/device/state'
    try:
        with open(state_path, 'r', encoding='utf-8', errors='ignore') as fh:
            return fh.read().strip().lower()
    except Exception:
        return ''


def c_power_state(path, strong=False):
    """Real low level power probe.

    Returns active, standby, powered, off, or unknown.
    The bundled C helper opens the block device and asks ATA/SATA disks with
    CHECK POWER MODE, so stale /dev/sdX entries cannot fake a green light.
    """
    if not safe_path(path):
        return 'unknown'
    if os.path.exists(C_POWER_HELPER) and os.access(C_POWER_HELPER, os.X_OK):
        code, out, err = run_cmd([C_POWER_HELPER, path], timeout=6)
        text = (out or err or '').strip().lower().splitlines()
        if text:
            word = text[0].strip()
            if word in ('active', 'idle', 'standby', 'powered', 'off', 'unknown'):
                return 'active' if word == 'idle' else word
    name = os.path.basename(path)
    sys_block = f'/sys/class/block/{name}'
    if not os.path.exists(sys_block):
        return 'off'
    state = read_sys_state(path)
    if state in ('offline', 'blocked', 'quiesce', 'quiesced', 'deleted', 'suspended'):
        return 'off'
    # Fallback only. Do not call this real proof on HDDs because stale kernel
    # entries can exist. Use yellow when the helper cannot verify it.
    if os.path.exists(path):
        return 'unknown'
    return 'off'

def hdparm_power_state(path):
    """Return the best ATA power state string without using kernel 'running' as spindle truth."""
    if not has_bin('hdparm') or not safe_path(path):
        return ''
    code, out, err = run_cmd(['hdparm', '-C', path], timeout=6)
    text = (out + '\n' + err).lower()
    marker = 'drive state is:'
    if marker in text:
        return text.split(marker, 1)[1].strip().splitlines()[0].strip()
    if 'standby' in text:
        return 'standby'
    if 'active/idle' in text or 'active' in text:
        return 'active/idle'
    return ''


def disk_status(node):
    """Strict user visible light state.

    Green is only shown when the disk answers as powered/active.
    Red is shown for a commanded hard power off, a missing disk, a stale
    block device, or an HDD that reports standby/suspended. This avoids the
    false green problem where lsblk/sysfs still show a disk after power off.
    """
    manual = str(node.get('_manual_status') or '').lower().strip()

    if node.get('_missing'):
        return '🔴 Power Off / not detected'

    path = node.get('path') or ''
    if not path:
        return '🔴 Power Off / not detected'

    # If the user clicked Hard Power Off, keep the drive red until a user
    # action such as Wake/Mount marks it active again. Stale /dev/sdX entries
    # must never override this.
    if manual in ('offline', 'off', 'poweroff', 'powered_off', 'forced_off'):
        return '🔴 Power Off / commanded off'

    power = c_power_state(path)

    if power == 'off':
        return '🔴 Power Off / no response'

    if power == 'standby':
        return '🔴 Power Off / standby'

    if power in ('active', 'powered'):
        return '🟢 Power On / detected'

    sys_state = read_sys_state(path)
    if sys_state in ('offline', 'blocked', 'quiesce', 'quiesced', 'deleted', 'suspended'):
        return '🔴 Power Off / offline'

    # On rotating HDDs, unknown is not allowed to become green. This is the
    # important fix: an unmounted/offline HDD can still appear in lsblk.
    if str(node.get('rota')) == '1':
        return '🟡 Unknown / not proven powered'

    # SSD/NVMe/virtual disks often cannot expose ATA power mode. If present,
    # connected, and not rotational, treat as powered.
    if os.path.exists(path):
        return '🟢 Power On / connected'

    return '🔴 Power Off / no device'

def disk_label(node):
    model = (node.get('model') or '').strip()
    vendor = (node.get('vendor') or '').strip()
    size = node.get('size') or ''
    path = node.get('path') or ''
    bits = [path, size]
    name = ' '.join(x for x in [vendor, model] if x).strip()
    if name:
        bits.append(name)
    return '  '.join(bits)


def children_partitions(node):
    parts = []
    for child in node.get('children') or []:
        if child.get('type') in ('part', 'crypt', 'lvm'):
            parts.append(child)
        parts.extend(children_partitions(child))
    return parts


def quote_cmd_path(path):
    if not safe_path(path):
        raise ValueError("Unsafe device path")
    return shlex.quote(path)


def pkexec_bash(command):
    if not has_bin('pkexec'):
        return 1, "", "pkexec is not installed. Install policykit-1 or policykit-desktop-privileges."
    return run_cmd(['pkexec', 'bash', '-lc', command], timeout=180)


def read_block_stat(path):
    if not safe_path(path):
        return None
    name = os.path.basename(path)
    stat_path = f'/sys/class/block/{name}/stat'
    try:
        with open(stat_path, 'r', encoding='utf-8', errors='ignore') as fh:
            return tuple(int(x) for x in fh.read().strip().split())
    except Exception:
        return None


class DiskPowerApp(Gtk.Window):
    def __init__(self):
        super().__init__(title=APP_NAME)
        self.set_default_size(1120, 760)
        self.set_position(Gtk.WindowPosition.CENTER)
        self.set_icon_name('govibe-disk-power-management')
        self.disks = []
        self.known_disks = {}
        self.manual_states = {}
        self.manual_path_states = {}
        self.persistent_poweroff = {}
        self.restore_policy_checked = False
        self.auto_sleep = {}
        self.auto_wake = {}
        self.last_stats = {}
        self.idle_since = {}
        self.auto_busy = set()
        self.selected_node = None
        self.admin_unlocked = False
        self.admin_keepalive_started = False
        self.load_persistent_state()

        self.apply_css()
        self.build_ui()
        self.refresh_disks()
        GLib.timeout_add_seconds(AUTO_TICK_SECONDS, self.auto_manager_tick)
        GLib.idle_add(self.ensure_admin_on_start)
        GLib.timeout_add_seconds(8, self.enforce_persistent_poweroff_once)

    def apply_css(self):
        css = b"""
        window { background: #111827; color: #e5e7eb; }
        .topbar { background: #0b1220; border-bottom: 1px solid #273449; }
        .title { color: #ffffff; font-weight: 800; font-size: 20px; }
        .subtitle { color: #9ca3af; font-size: 12px; }
        treeview { background: #0f172a; color: #e5e7eb; border: 1px solid #334155; }
        treeview:selected { background: #2563eb; color: #ffffff; }
        button { background: #1f2937; color: #ffffff; border: 1px solid #475569; border-radius: 8px; padding: 8px 10px; }
        button:hover { background: #334155; }
        button.suggested-action { background: #14532d; border-color: #22c55e; }
        button.destructive-action { background: #7f1d1d; border-color: #ef4444; }
        button.warning { background: #78350f; border-color: #f59e0b; }
        entry, spinbutton { background: #020617; color: #e5e7eb; border: 1px solid #475569; border-radius: 7px; padding: 5px; }
        label.badge { background: #172554; color: #bfdbfe; padding: 4px 8px; border-radius: 8px; }
        textview { background: #020617; color: #d1d5db; border: 1px solid #334155; }
        .panel { background: #0f172a; border: 1px solid #334155; border-radius: 12px; padding: 10px; }
        """
        provider = Gtk.CssProvider()
        provider.load_from_data(css)
        Gtk.StyleContext.add_provider_for_screen(Gdk.Screen.get_default(), provider, Gtk.STYLE_PROVIDER_PRIORITY_APPLICATION)

    def build_ui(self):
        root = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=0)
        self.add(root)

        top = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL, spacing=12)
        top.set_margin_top(12)
        top.set_margin_bottom(12)
        top.set_margin_start(14)
        top.set_margin_end(14)
        top.get_style_context().add_class('topbar')
        root.pack_start(top, False, False, 0)

        logo = Gtk.Label(label='⚡')
        top.pack_start(logo, False, False, 0)

        titlebox = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=2)
        top.pack_start(titlebox, True, True, 0)
        title = Gtk.Label(label=f'{APP_NAME} {APP_VERSION}')
        title.set_xalign(0)
        title.get_style_context().add_class('title')
        titlebox.pack_start(title, False, False, 0)
        subtitle = Gtk.Label(label='Unmount, sleep, power off, combined unmount+power off, auto sleep, auto wake, rescan, and mount secondary drives from one safe panel')
        subtitle.set_xalign(0)
        subtitle.get_style_context().add_class('subtitle')
        titlebox.pack_start(subtitle, False, False, 0)

        self.status_badge = Gtk.Label(label='Ready')
        self.status_badge.get_style_context().add_class('badge')
        top.pack_end(self.status_badge, False, False, 0)

        main = Gtk.Paned(orientation=Gtk.Orientation.HORIZONTAL)
        root.pack_start(main, True, True, 0)

        left = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=8)
        left.set_margin_top(12)
        left.set_margin_bottom(12)
        left.set_margin_start(12)
        left.set_margin_end(8)
        main.pack1(left, resize=True, shrink=False)

        info = Gtk.Label()
        info.set_markup('<b>Detected drives</b>\nGreen means the disk answers as powered/active. Red means powered off, standby, not detected, or manually powered off. Yellow means Linux sees it but the app cannot prove real HDD power state.')
        info.set_xalign(0)
        info.set_line_wrap(True)
        left.pack_start(info, False, False, 0)

        self.store = Gtk.ListStore(str, str, str, str, str, str, str, bool, object)
        self.tree = Gtk.TreeView(model=self.store)
        columns = [
            ('Device', 0, 135),
            ('Light', 1, 170),
            ('Size', 2, 85),
            ('Type', 3, 150),
            ('Drive', 4, 230),
            ('Mounts', 5, 250),
            ('Kernel', 6, 80),
        ]
        for name, idx, width in columns:
            renderer = Gtk.CellRendererText()
            renderer.set_property('ellipsize', Pango.EllipsizeMode.END)
            col = Gtk.TreeViewColumn(name, renderer, text=idx)
            col.set_resizable(True)
            col.set_min_width(width)
            self.tree.append_column(col)
        self.tree.get_selection().connect('changed', self.on_selection_changed)
        scroller = Gtk.ScrolledWindow()
        scroller.set_policy(Gtk.PolicyType.AUTOMATIC, Gtk.PolicyType.AUTOMATIC)
        scroller.add(self.tree)
        left.pack_start(scroller, True, True, 0)

        right_scroll = Gtk.ScrolledWindow()
        right_scroll.set_policy(Gtk.PolicyType.NEVER, Gtk.PolicyType.AUTOMATIC)
        main.pack2(right_scroll, resize=False, shrink=False)
        right_scroll.set_size_request(390, -1)

        right = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=10)
        right.set_margin_top(12)
        right.set_margin_bottom(12)
        right.set_margin_start(8)
        right.set_margin_end(12)
        right_scroll.add(right)

        action_panel = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=8)
        action_panel.get_style_context().add_class('panel')
        right.pack_start(action_panel, False, False, 0)

        self.selected_label = Gtk.Label(label='No disk selected')
        self.selected_label.set_xalign(0)
        self.selected_label.set_line_wrap(True)
        action_panel.pack_start(self.selected_label, False, False, 0)

        self.btn_refresh = self.make_button('Refresh Status Lights / Drive List', self.refresh_disks)
        action_panel.pack_start(self.btn_refresh, False, False, 0)

        self.btn_mount = self.make_button('Wake / Mount Selected Drive', self.mount_selected, style='suggested-action')
        action_panel.pack_start(self.btn_mount, False, False, 0)

        self.btn_unmount = self.make_button('Unmount Selected Drive Partitions', self.unmount_selected)
        action_panel.pack_start(self.btn_unmount, False, False, 0)

        self.btn_sleep = self.make_button('Sleep Drive Now', self.sleep_selected, style='warning')
        action_panel.pack_start(self.btn_sleep, False, False, 0)

        self.btn_poweroff = self.make_button('Hard Power Off Drive', self.poweroff_selected, style='destructive-action')
        action_panel.pack_start(self.btn_poweroff, False, False, 0)

        self.btn_unmount_poweroff = self.make_button('Unmount + Power Off Drive', self.unmount_poweroff_selected, style='destructive-action')
        action_panel.pack_start(self.btn_unmount_poweroff, False, False, 0)

        self.btn_forget_poweroff = self.make_button('Forget Saved Power-Off Policy', self.forget_saved_poweroff_selected)
        action_panel.pack_start(self.btn_forget_poweroff, False, False, 0)
        self.btn_forget_poweroff.set_sensitive(False)

        self.btn_rescan = self.make_button('Rescan SATA / USB Buses', self.rescan_buses)
        action_panel.pack_start(self.btn_rescan, False, False, 0)

        auto_panel = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=8)
        auto_panel.get_style_context().add_class('panel')
        right.pack_start(auto_panel, False, False, 0)

        auto_title = Gtk.Label()
        auto_title.set_markup('<b>Auto power management</b>')
        auto_title.set_xalign(0)
        auto_panel.pack_start(auto_title, False, False, 0)

        row = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL, spacing=8)
        auto_panel.pack_start(row, False, False, 0)
        row.pack_start(Gtk.Label(label='Idle minutes:'), False, False, 0)
        self.auto_minutes = Gtk.SpinButton.new_with_range(1, 1440, 1)
        self.auto_minutes.set_value(15)
        row.pack_start(self.auto_minutes, True, True, 0)

        self.btn_auto_enable = self.make_button('Enable Auto Sleep For Selected Drive', self.enable_auto_sleep, style='warning')
        auto_panel.pack_start(self.btn_auto_enable, False, False, 0)
        self.btn_auto_disable = self.make_button('Disable Auto Sleep For Selected Drive', self.disable_auto_sleep)
        auto_panel.pack_start(self.btn_auto_disable, False, False, 0)

        self.wake_path_entry = Gtk.Entry()
        self.wake_path_entry.set_placeholder_text('/media/yourname/DriveName or leave blank to use current mount')
        auto_panel.pack_start(self.wake_path_entry, False, False, 0)

        self.btn_wake_enable = self.make_button('Enable Auto Wake / Mount Watch', self.enable_auto_wake, style='suggested-action')
        auto_panel.pack_start(self.btn_wake_enable, False, False, 0)
        self.btn_wake_disable = self.make_button('Disable Auto Wake For Selected Drive', self.disable_auto_wake)
        auto_panel.pack_start(self.btn_wake_disable, False, False, 0)

        self.auto_status_label = Gtk.Label(label='Auto sleep disabled. Auto wake disabled.')
        self.auto_status_label.set_xalign(0)
        self.auto_status_label.set_line_wrap(True)
        auto_panel.pack_start(self.auto_status_label, False, False, 0)

        startup_panel = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=8)
        startup_panel.get_style_context().add_class('panel')
        right.pack_start(startup_panel, False, False, 0)

        startup_title = Gtk.Label()
        startup_title.set_markup('<b>Startup and administrator session</b>')
        startup_title.set_xalign(0)
        startup_panel.pack_start(startup_title, False, False, 0)

        self.admin_status_label = Gtk.Label(label='Administrator helper: checking...')
        self.admin_status_label.set_xalign(0)
        self.admin_status_label.set_line_wrap(True)
        startup_panel.pack_start(self.admin_status_label, False, False, 0)

        self.btn_unlock_admin = self.make_button('Unlock Administrator For This Session', self.unlock_admin_clicked, style='suggested-action')
        startup_panel.pack_start(self.btn_unlock_admin, False, False, 0)

        start_row = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL, spacing=8)
        startup_panel.pack_start(start_row, False, False, 0)
        self.btn_start_login = self.make_button('Start On Login', self.enable_start_on_login, style='suggested-action')
        start_row.pack_start(self.btn_start_login, True, True, 0)
        self.btn_stop_login = self.make_button('Disable Login Start', self.disable_start_on_login)
        start_row.pack_start(self.btn_stop_login, True, True, 0)

        self.startup_status_label = Gtk.Label(label='Start on login: checking...')
        self.startup_status_label.set_xalign(0)
        self.startup_status_label.set_line_wrap(True)
        startup_panel.pack_start(self.startup_status_label, False, False, 0)
        self.update_startup_status()

        notes = Gtk.Label()
        notes.set_markup(
            '<b>Safe meaning</b>\n'
            'Sleep Drive uses hdparm standby. Internal HDDs usually wake without reboot when accessed.\n\n'
            '<b>Status lights</b> now use a real C ATA power probe. Kernel running, lsblk online, mounted state, and stale /dev/sdX alone are never trusted. After Hard Power Off, the drive is saved in a persistent keep-powered-off policy. When this app starts at login after a reboot, it will unmount and power off that drive again automatically until you use Wake/Mount or Forget Saved Power-Off Policy.\n\n'
            '<b>Unmount + Power Off</b> safely unmounts every mounted non-system partition on the selected disk, then powers off the whole drive in one action.\n\n'
            '<b>Auto sleep</b> watches disk I/O and sleeps the drive after the selected idle time. It does not hard power off automatically.\n\n'
            '<b>Auto wake</b> checks a mount path. If it is inaccessible, the app rescans buses and tries to mount the drive. Full internal SATA power off may still require reboot on some controllers.\n\n<b>Administrator helper</b> is limited to disk actions only. It avoids storing your password and allows the app to start on login without repeated prompts. Saved powered-off disks are automatically powered off again after login/reboot.'
        )
        notes.set_xalign(0)
        notes.set_line_wrap(True)
        right.pack_start(notes, False, False, 0)

        log_label = Gtk.Label()
        log_label.set_markup('<b>Operation log</b>')
        log_label.set_xalign(0)
        right.pack_start(log_label, False, False, 0)

        self.log_buffer = Gtk.TextBuffer()
        self.log_view = Gtk.TextView(buffer=self.log_buffer)
        self.log_view.set_editable(False)
        self.log_view.set_cursor_visible(False)
        self.log_view.set_wrap_mode(Gtk.WrapMode.WORD_CHAR)
        log_scroll = Gtk.ScrolledWindow()
        log_scroll.set_policy(Gtk.PolicyType.AUTOMATIC, Gtk.PolicyType.AUTOMATIC)
        log_scroll.set_min_content_height(160)
        log_scroll.add(self.log_view)
        right.pack_start(log_scroll, True, True, 0)

        foot = Gtk.Label(label='govibe.org')
        foot.set_xalign(1)
        foot.get_style_context().add_class('subtitle')
        root.pack_end(foot, False, False, 8)

        self.set_action_buttons(False)

    def make_button(self, label, callback, style=None):
        b = Gtk.Button(label=label)
        if style:
            b.get_style_context().add_class(style)
        b.connect('clicked', lambda *_: callback())
        return b

    def set_action_buttons(self, enabled):
        for b in [self.btn_mount, self.btn_unmount, self.btn_sleep, self.btn_poweroff, self.btn_unmount_poweroff, self.btn_auto_enable, self.btn_auto_disable, self.btn_wake_enable, self.btn_wake_disable]:
            b.set_sensitive(enabled)
        if hasattr(self, 'btn_forget_poweroff') and self.selected_node is None:
            self.btn_forget_poweroff.set_sensitive(False)

    def log(self, msg):
        ts = datetime.now().strftime('%H:%M:%S')
        end_iter = self.log_buffer.get_end_iter()
        self.log_buffer.insert(end_iter, f"[{ts}] {msg}\n")
        mark = self.log_buffer.create_mark(None, self.log_buffer.get_end_iter(), False)
        self.log_view.scroll_to_mark(mark, 0.0, True, 0.0, 1.0)

    def set_status(self, text):
        self.status_badge.set_text(text)

    def clean_node_for_persistence(self, node):
        try:
            clean = json.loads(json.dumps(node))
        except Exception:
            return {}
        if isinstance(clean, dict):
            for key in list(clean.keys()):
                if str(key).startswith('_'):
                    clean.pop(key, None)
            for child in clean.get('children') or []:
                if isinstance(child, dict):
                    for key in list(child.keys()):
                        if str(key).startswith('_'):
                            child.pop(key, None)
        return clean

    def load_persistent_state(self):
        try:
            if not os.path.exists(STATE_FILE):
                return
            with open(STATE_FILE, 'r', encoding='utf-8') as fh:
                data = json.load(fh)
            kd = data.get('known_disks') or {}
            if isinstance(kd, dict):
                for identity, node in kd.items():
                    if isinstance(node, dict):
                        self.known_disks[identity] = self.clean_node_for_persistence(node)
            pp = data.get('persistent_poweroff') or {}
            if isinstance(pp, dict):
                for identity, info in pp.items():
                    if isinstance(identity, str) and identity:
                        if not isinstance(info, dict):
                            info = {'path': '', 'label': ''}
                        self.persistent_poweroff[identity] = {
                            'path': str(info.get('path') or ''),
                            'label': str(info.get('label') or ''),
                            'saved_at': str(info.get('saved_at') or ''),
                        }
                        self.manual_states[identity] = 'forced_off'
                        path = str(info.get('path') or '')
                        if path:
                            self.manual_path_states[path] = 'forced_off'
            legacy_manual = data.get('manual_states') or {}
            if isinstance(legacy_manual, dict):
                for identity, state in legacy_manual.items():
                    if str(state) == 'forced_off' and identity not in self.persistent_poweroff:
                        self.persistent_poweroff[identity] = {'path': '', 'label': '', 'saved_at': str(data.get('saved_at') or '')}
                        self.manual_states[identity] = 'forced_off'
        except Exception as exc:
            self.log(f'Could not load saved disk state: {exc}') if hasattr(self, 'log_buffer') else None

    def save_persistent_state(self):
        try:
            os.makedirs(STATE_DIR, exist_ok=True)
            clean_known = {}
            for identity, node in self.known_disks.items():
                clean_known[identity] = self.clean_node_for_persistence(node)
            data = {
                'known_disks': clean_known,
                'persistent_poweroff': self.persistent_poweroff,
                'saved_at': datetime.now().isoformat(timespec='seconds'),
                'note': 'Disks listed in persistent_poweroff are automatically unmounted and powered off again when the app starts at login after reboot.',
            }
            tmp = STATE_FILE + '.tmp'
            with open(tmp, 'w', encoding='utf-8') as fh:
                json.dump(data, fh, indent=2, sort_keys=True)
            os.replace(tmp, STATE_FILE)
        except Exception as exc:
            self.log(f'Could not save disk state: {exc}') if hasattr(self, 'log_buffer') else None

    def is_persistent_poweroff(self, identity, path=''):
        if identity in self.persistent_poweroff:
            return True
        if path:
            for info in self.persistent_poweroff.values():
                if isinstance(info, dict) and info.get('path') == path:
                    return True
        return False

    def save_poweroff_policy(self, node):
        identity = disk_identity(node)
        path = node.get('path') or ''
        self.persistent_poweroff[identity] = {
            'path': path,
            'label': disk_label(node),
            'saved_at': datetime.now().isoformat(timespec='seconds'),
        }
        self.manual_states[identity] = 'forced_off'
        if path:
            self.manual_path_states[path] = 'forced_off'
        self.save_persistent_state()

    def clear_poweroff_policy(self, identity, path=''):
        removed = False
        if identity in self.persistent_poweroff:
            self.persistent_poweroff.pop(identity, None)
            removed = True
        if path:
            for key, info in list(self.persistent_poweroff.items()):
                if isinstance(info, dict) and info.get('path') == path:
                    self.persistent_poweroff.pop(key, None)
                    removed = True
        self.manual_states.pop(identity, None)
        if path:
            self.manual_path_states.pop(path, None)
        if removed:
            self.save_persistent_state()
        return removed

    def admin_action_noprompt(self, args, timeout=180):
        if not isinstance(args, list) or not args:
            return 1, '', 'Invalid administrator action.'
        if os.path.exists(ADMIN_HELPER) and os.access(ADMIN_HELPER, os.X_OK):
            return run_cmd(['sudo', '-n', ADMIN_HELPER] + args, timeout=timeout)
        return 1, '', 'Administrator helper is missing, so automatic login restore cannot run.'

    def helper_ping(self):
        if not os.path.exists(ADMIN_HELPER) or not os.access(ADMIN_HELPER, os.X_OK):
            return False, 'Administrator helper is missing.'
        code, out, err = run_cmd(['sudo', '-n', ADMIN_HELPER, 'ping'], timeout=8)
        if code == 0:
            return True, out or 'Administrator helper is ready.'
        return False, err or out or 'Administrator helper needs unlock.'

    def start_sudo_keepalive(self):
        if self.admin_keepalive_started:
            return
        self.admin_keepalive_started = True
        def keepalive():
            while True:
                try:
                    subprocess.run(['sudo', '-n', '-v'], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=8)
                except Exception:
                    pass
                time.sleep(60)
        threading.Thread(target=keepalive, daemon=True).start()

    def set_admin_status(self, text):
        if hasattr(self, 'admin_status_label'):
            self.admin_status_label.set_text(text)

    def prompt_admin_password(self):
        dialog = Gtk.Dialog(title='Administrator unlock', parent=self, flags=0)
        dialog.add_button('Cancel', Gtk.ResponseType.CANCEL)
        dialog.add_button('Unlock', Gtk.ResponseType.OK)
        dialog.set_default_response(Gtk.ResponseType.OK)
        box = dialog.get_content_area()
        label = Gtk.Label(label='Type your Linux administrator password one time for this app session. The password is not saved.')
        label.set_xalign(0)
        label.set_line_wrap(True)
        box.pack_start(label, False, False, 8)
        entry = Gtk.Entry()
        entry.set_visibility(False)
        entry.set_activates_default(True)
        entry.set_placeholder_text('Administrator password')
        box.pack_start(entry, False, False, 8)
        dialog.show_all()
        response = dialog.run()
        password = entry.get_text()
        dialog.destroy()
        if response != Gtk.ResponseType.OK or not password:
            return False
        try:
            proc = subprocess.run(['sudo', '-S', '-p', '', '-v'], input=password + '\n', text=True, capture_output=True, timeout=20)
        finally:
            password = ''
        if proc.returncode == 0:
            self.admin_unlocked = True
            self.start_sudo_keepalive()
            return True
        self.log((proc.stderr or proc.stdout or 'Administrator unlock failed.').strip())
        return False

    def ensure_admin(self, show_dialog=True):
        ok, message = self.helper_ping()
        if ok:
            self.admin_unlocked = True
            self.set_admin_status('Administrator helper: ready. Disk actions should not ask for a password.')
            return True
        code, out, err = run_cmd(['sudo', '-n', '-v'], timeout=8)
        if code == 0:
            self.admin_unlocked = True
            self.start_sudo_keepalive()
            self.set_admin_status('Administrator session: unlocked for this app session.')
            return True
        self.set_admin_status('Administrator helper: needs unlock. Use the unlock button once for this session.')
        if not show_dialog:
            return False
        if self.prompt_admin_password():
            ok, message = self.helper_ping()
            if ok:
                self.set_admin_status('Administrator session: unlocked for this app session.')
                return True
            self.set_admin_status('Administrator session unlocked, but the disk helper did not respond correctly.')
            return True
        self.set_admin_status('Administrator session: locked.')
        return False

    def ensure_admin_on_start(self):
        self.ensure_admin(show_dialog=False)
        return False

    def unlock_admin_clicked(self):
        self.ensure_admin(show_dialog=True)

    def admin_action(self, args, timeout=180):
        if not isinstance(args, list) or not args:
            return 1, '', 'Invalid administrator action.'
        if os.path.exists(ADMIN_HELPER) and os.access(ADMIN_HELPER, os.X_OK):
            code, out, err = run_cmd(['sudo', '-n', ADMIN_HELPER] + args, timeout=timeout)
            if code == 0:
                self.admin_unlocked = True
                self.set_admin_status('Administrator helper: ready. Disk actions should not ask for a password.')
                return code, out, err
            auth_text = (out + '\n' + err).lower()
            if 'password' in auth_text or 'sudo' in auth_text or 'authentication' in auth_text or 'a password is required' in auth_text:
                if self.ensure_admin(show_dialog=True):
                    return run_cmd(['sudo', '-n', ADMIN_HELPER] + args, timeout=timeout)
            return code, out, err
        if self.ensure_admin(show_dialog=True):
            return run_cmd(['sudo', '-n'] + args, timeout=timeout)
        return 1, '', 'Administrator session is locked.'

    def autostart_desktop_text(self):
        return '''[Desktop Entry]\nType=Application\nName=Govibe Disk Power Management\nComment=Start Govibe Disk Power Management at login\nExec=govibe-disk-power-management\nIcon=govibe-disk-power-management\nTerminal=false\nX-GNOME-Autostart-enabled=true\nCategories=System;Utility;GTK;\n'''

    def update_startup_status(self):
        enabled = os.path.exists(AUTOSTART_FILE)
        if hasattr(self, 'startup_status_label'):
            self.startup_status_label.set_text('Start on login: enabled' if enabled else 'Start on login: disabled')
        if hasattr(self, 'btn_start_login'):
            self.btn_start_login.set_sensitive(not enabled)
        if hasattr(self, 'btn_stop_login'):
            self.btn_stop_login.set_sensitive(enabled)

    def enable_start_on_login(self):
        try:
            os.makedirs(AUTOSTART_DIR, exist_ok=True)
            with open(AUTOSTART_FILE, 'w', encoding='utf-8') as fh:
                fh.write(self.autostart_desktop_text())
            os.chmod(AUTOSTART_FILE, 0o644)
            self.log('Start on login enabled for this Linux user.')
        except Exception as exc:
            self.log(f'Could not enable start on login: {exc}')
        self.update_startup_status()

    def disable_start_on_login(self):
        try:
            if os.path.exists(AUTOSTART_FILE):
                os.remove(AUTOSTART_FILE)
            self.log('Start on login disabled for this Linux user.')
        except Exception as exc:
            self.log(f'Could not disable start on login: {exc}')
        self.update_startup_status()

    def run_background(self, label, func, after=None):
        self.set_status(label)
        self.log(label)
        def worker():
            try:
                result = func()
            except Exception as exc:
                result = (1, '', str(exc))
            def finish():
                code, out, err = result if isinstance(result, tuple) else (0, str(result), '')
                if out:
                    self.log(out)
                if err:
                    self.log(err)
                self.set_status('Done' if code == 0 else 'Error')
                if after:
                    after()
                return False
            GLib.idle_add(finish)
        threading.Thread(target=worker, daemon=True).start()

    def refresh_disks(self):
        old_selected = disk_identity(self.selected_node) if self.selected_node else None
        self.store.clear()
        self.selected_node = None
        self.selected_label.set_text('No disk selected')
        self.set_action_buttons(False)
        code, out, err = run_cmd(['lsblk', '-J', '-o', LSBLK_COLUMNS], timeout=20)
        if code != 0:
            self.log(err or 'Could not read lsblk output')
            return
        try:
            data = json.loads(out)
        except Exception as exc:
            self.log(f'Could not parse lsblk JSON: {exc}')
            return
        self.disks = [n for n in data.get('blockdevices', []) if n.get('type') == 'disk']
        current_ids = set()

        for node in self.disks:
            identity = disk_identity(node)
            current_ids.add(identity)
            path = node.get('path') or ''
            if self.is_persistent_poweroff(identity, path):
                node['_manual_status'] = 'forced_off'
                node['_persistent_poweroff'] = True
                self.manual_states[identity] = 'forced_off'
                if path:
                    self.manual_path_states[path] = 'forced_off'
                if identity not in self.persistent_poweroff:
                    self.persistent_poweroff[identity] = {
                        'path': path,
                        'label': disk_label(node),
                        'saved_at': datetime.now().isoformat(timespec='seconds'),
                    }
            elif identity in self.manual_states:
                node['_manual_status'] = self.manual_states[identity]
            elif path in getattr(self, 'manual_path_states', {}):
                node['_manual_status'] = self.manual_path_states[path]
            # Persistent power-off must survive reboot. Do not clear it only
            # because Linux or the desktop automounter powered the disk back on.
            if node.get('_manual_status') in ('offline', 'off', 'poweroff', 'powered_off', 'forced_off') and collect_mountpoints(node) and not self.is_persistent_poweroff(identity, path):
                node['_manual_status'] = 'active'
                self.manual_states[identity] = 'active'
                self.manual_path_states[path] = 'active'
            self.known_disks[identity] = json.loads(json.dumps(node))
            self.append_disk_row(node)

        missing_count = 0
        for identity, old_node in list(self.known_disks.items()):
            if identity in current_ids:
                continue
            missing = json.loads(json.dumps(old_node))
            missing['_missing'] = True
            missing['_manual_status'] = 'offline'
            self.append_disk_row(missing)
            missing_count += 1

        msg = f'Refreshed. {len(self.disks)} active disk drive(s) detected.'
        if missing_count:
            msg += f' {missing_count} previously detected drive(s) are now offline.'
        self.log(msg)
        self.set_status('Ready')
        self.save_persistent_state()
        self.update_auto_status()

        if old_selected:
            self.select_identity(old_selected)

    def append_disk_row(self, node):
        mounts = collect_mountpoints(node)
        system = disk_is_system(node)
        missing = bool(node.get('_missing'))
        model = ' '.join(x for x in [(node.get('vendor') or '').strip(), (node.get('model') or '').strip()] if x).strip()
        dtype = []
        if str(node.get('rota')) == '1':
            dtype.append('HDD')
        elif str(node.get('rota')) == '0':
            dtype.append('SSD/NVMe')
        if str(node.get('rm')) == '1' or str(node.get('hotplug')) == '1':
            dtype.append('USB/Removable')
        if system:
            dtype.append('SYSTEM PROTECTED')
        if missing:
            dtype.append('OFFLINE')
        state = node.get('state') or ''
        self.store.append([
            node.get('path') or node.get('name') or '',
            disk_status(node),
            node.get('size') or '',
            ', '.join(dtype) or 'Disk',
            model or 'Unknown model',
            'Offline / not detected' if missing else (', '.join(mounts) if mounts else 'Not mounted'),
            'offline' if missing else (state or 'Unknown'),
            system or missing,
            node,
        ])

    def select_identity(self, identity):
        itr = self.store.get_iter_first()
        while itr:
            node = self.store[itr][8]
            if disk_identity(node) == identity:
                self.tree.get_selection().select_iter(itr)
                return
            itr = self.store.iter_next(itr)

    def on_selection_changed(self, selection):
        model, treeiter = selection.get_selected()
        if treeiter is None:
            self.selected_node = None
            self.selected_label.set_text('No disk selected')
            self.set_action_buttons(False)
            if hasattr(self, 'btn_forget_poweroff'):
                self.btn_forget_poweroff.set_sensitive(False)
            return
        node = model[treeiter][8]
        self.selected_node = node
        protected = disk_is_system(node) or bool(node.get('_missing'))
        mounts = collect_mountpoints(node)
        identity = disk_identity(node)
        label = f"Selected: {disk_label(node)}\nStatus: {disk_status(node)}\nMounts: {', '.join(mounts) if mounts else 'not mounted'}"
        if identity in self.auto_sleep:
            label += f"\nAuto sleep: enabled after {self.auto_sleep[identity]['minutes']} minute(s) idle"
        if identity in self.auto_wake:
            label += f"\nAuto wake watch: {self.auto_wake[identity]['path']}"
        if self.is_persistent_poweroff(identity, node.get('path') or ''):
            label += "\nSaved reboot policy: keep this disk powered off after login/reboot"
        if hasattr(self, 'btn_forget_poweroff'):
            self.btn_forget_poweroff.set_sensitive(True)
        if node.get('_missing'):
            label += "\nThis drive is currently offline or fully powered off. Use Rescan SATA / USB Buses, or reboot if it does not return."
        elif protected:
            label += "\nProtected system disk. Actions disabled."
        self.selected_label.set_text(label)
        self.set_action_buttons(not protected)
        self.update_auto_status()

    def update_auto_status(self):
        if not hasattr(self, 'auto_status_label'):
            return
        parts = []
        if self.auto_sleep:
            parts.append(f'Auto sleep active on {len(self.auto_sleep)} drive(s).')
        else:
            parts.append('Auto sleep disabled.')
        if self.auto_wake:
            parts.append(f'Auto wake active on {len(self.auto_wake)} drive(s).')
        else:
            parts.append('Auto wake disabled.')
        self.auto_status_label.set_text(' '.join(parts))

    def confirm(self, title, text):
        dialog = Gtk.MessageDialog(parent=self, flags=0, message_type=Gtk.MessageType.WARNING, buttons=Gtk.ButtonsType.OK_CANCEL, text=title)
        dialog.format_secondary_text(text)
        response = dialog.run()
        dialog.destroy()
        return response == Gtk.ResponseType.OK

    def require_node(self):
        node = self.selected_node
        if not node:
            self.log('No disk selected.')
            return None
        if node.get('_missing'):
            self.log('Blocked: selected drive is offline. Rescan buses or reboot to bring it back.')
            return None
        if disk_is_system(node):
            self.log('Blocked: selected drive contains protected system mountpoints.')
            return None
        path = node.get('path')
        if not safe_path(path):
            self.log('Blocked: invalid device path.')
            return None
        return node

    def unmount_selected(self):
        node = self.require_node()
        if not node:
            return
        parts = children_partitions(node)
        mounted = []
        for part in parts:
            path = part.get('path')
            mps = flatten_mountpoints(part.get('mountpoints'))
            if path and mps and not any(is_system_mount(mp) for mp in mps):
                mounted.append(path)
        if not mounted:
            self.log('No mounted non system partitions found on selected drive.')
            return
        if not self.confirm('Unmount selected drive partitions?', 'This will unmount: ' + ', '.join(mounted)):
            return
        self.run_background('Unmounting selected partitions', lambda: self.admin_action(['unmount'] + mounted), after=self.refresh_disks)

    def sleep_selected(self):
        node = self.require_node()
        if not node:
            return
        path = node.get('path')
        identity = disk_identity(node)
        if not self.confirm('Sleep selected drive now?', f'{path} will enter standby. It should wake again when accessed.'):
            return
        def after():
            self.manual_states[identity] = 'standby'
            self.manual_path_states[path] = 'standby'
            self.save_persistent_state()
            self.refresh_disks()
        self.run_background('Putting drive into standby', lambda: self.admin_action(['sleep', path]), after=after)

    def poweroff_selected(self):
        node = self.require_node()
        if not node:
            return
        path = node.get('path')
        identity = disk_identity(node)
        msg = (f'{path} will be powered off through udisksctl. For internal SATA drives, Linux may not see it again until rescan or reboot.\n\n'
               'Use Sleep Drive if you only want it quiet but easy to wake.')
        if not self.confirm('Hard power off selected drive?', msg):
            return
        def task():
            return self.admin_action(['poweroff', path], timeout=180)
        def after():
            self.save_poweroff_policy(node)
            self.refresh_disks()
        self.run_background('Powering off selected drive', task, after=after)

    def unmount_poweroff_selected(self):
        node = self.require_node()
        if not node:
            return
        path = node.get('path')
        identity = disk_identity(node)
        parts = children_partitions(node)
        mounted = []
        for part in parts:
            part_path = part.get('path')
            mps = flatten_mountpoints(part.get('mountpoints'))
            if not part_path or not mps:
                continue
            if any(is_system_mount(mp) for mp in mps):
                self.log('Blocked: selected drive has a protected system mountpoint.')
                return
            mounted.append(part_path)

        mount_text = ', '.join(mounted) if mounted else 'no mounted partitions found'
        msg = (f'This will sync data, unmount all mounted partitions on {path}, then hard power off the whole drive.\n\n'
               f'Partitions to unmount: {mount_text}\n\n'
               'For internal SATA drives, Linux may need rescan or reboot to see the disk again.')
        if not self.confirm('Unmount + power off selected drive?', msg):
            return

        def task():
            messages = []
            if mounted:
                # First try normal udisksctl unmount so desktop mounts are handled cleanly.
                failed = []
                for part_path in mounted:
                    code, out, err = run_cmd(['udisksctl', 'unmount', '-b', part_path], timeout=60)
                    if out:
                        messages.append(out)
                    if err:
                        messages.append(err)
                    if code != 0:
                        failed.append(part_path)
                if failed:
                    code, out, err = self.admin_action(['unmount'] + failed, timeout=180)
                    if out:
                        messages.append(out)
                    if err:
                        messages.append(err)
                    if code != 0:
                        return code, '\n'.join(messages), 'Unmount failed. Drive was not powered off.'
            else:
                messages.append('No mounted partitions needed unmounting.')

            code, out, err = self.admin_action(['poweroff', path], timeout=180)
            if out:
                messages.append(out)
            if err:
                messages.append(err)
            return code, '\n'.join(messages), '' if code == 0 else 'Power off failed.'

        def after():
            self.save_poweroff_policy(node)
            self.refresh_disks()
        self.run_background('Unmounting and powering off selected drive', task, after=after)

    def forget_saved_poweroff_selected(self):
        node = self.selected_node
        if not node:
            self.log('No disk selected.')
            return
        identity = disk_identity(node)
        path = node.get('path') or ''
        if self.clear_poweroff_policy(identity, path):
            self.log('Saved power-off policy removed for selected drive. It will no longer be automatically powered off after login/reboot.')
        else:
            self.log('No saved power-off policy existed for selected drive.')
        self.refresh_disks()

    def enforce_persistent_poweroff_once(self):
        if self.restore_policy_checked:
            return False
        self.restore_policy_checked = True
        if not self.persistent_poweroff:
            return False

        def task():
            messages = []
            errors = []
            code_final = 0
            targets = []
            for node in list(self.disks):
                identity = disk_identity(node)
                path = node.get('path') or ''
                if not self.is_persistent_poweroff(identity, path):
                    continue
                if not safe_path(path) or disk_is_system(node) or node.get('_missing'):
                    continue
                targets.append((identity, path, node))
            if not targets:
                return 0, 'Saved power-off policy exists, but no matching powered disk is currently detected.', ''
            for identity, path, node in targets:
                messages.append(f'Restoring saved power-off policy for {path}.')
                mounted = []
                skip = False
                for part in children_partitions(node):
                    part_path = part.get('path')
                    mps = flatten_mountpoints(part.get('mountpoints'))
                    if not part_path or not mps:
                        continue
                    if any(is_system_mount(mp) for mp in mps):
                        errors.append(f'Skipped {path}: protected mountpoint detected.')
                        skip = True
                        break
                    mounted.append(part_path)
                if skip:
                    continue
                if mounted:
                    code, out, err = self.admin_action_noprompt(['unmount'] + mounted, timeout=180)
                    if out:
                        messages.append(out)
                    if err:
                        errors.append(err)
                    if code != 0:
                        code_final = code
                        errors.append(f'Could not unmount {path}; power off skipped for safety.')
                        continue
                code, out, err = self.admin_action_noprompt(['poweroff', path], timeout=180)
                if out:
                    messages.append(out)
                if err:
                    errors.append(err)
                if code != 0:
                    code_final = code
                else:
                    self.manual_states[identity] = 'forced_off'
                    self.manual_path_states[path] = 'forced_off'
            return code_final, '\n'.join(messages), '\n'.join(errors)

        self.run_background('Restoring saved powered-off disks after login/reboot', task, after=self.refresh_disks)
        return False

    def rescan_buses(self):
        self.run_background('Rescanning storage buses', lambda: self.admin_action(['rescan'], timeout=180), after=self.refresh_disks)

    def mount_selected(self):
        node = self.require_node()
        if not node:
            return
        self.mount_node(node)

    def mount_node(self, node):
        identity = disk_identity(node)
        parts = children_partitions(node)
        candidates = []
        for part in parts:
            path = part.get('path')
            fstype = part.get('fstype')
            mps = flatten_mountpoints(part.get('mountpoints'))
            if path and fstype and not mps:
                candidates.append(path)
        if not candidates:
            self.log('No unmounted filesystem partition found. Running a bus rescan instead.')
            self.rescan_buses()
            return
        target = candidates[0]
        def task():
            code, out, err = run_cmd(['udisksctl', 'mount', '-b', target], timeout=120)
            if code == 0:
                return code, out, err
            text = (out + '\n' + err).lower()
            if 'not authorized' in text or 'authentication' in text or 'permission' in text:
                code2, out2, err2 = self.admin_action(['mount', target], timeout=180)
                return code2, (out + '\n' + out2).strip(), (err + '\n' + err2).strip()
            return code, out, err
        def after():
            path = node.get('path') or ''
            self.clear_poweroff_policy(identity, path)
            self.manual_states[identity] = 'active'
            if path:
                self.manual_path_states[path] = 'active'
            self.save_persistent_state()
            self.refresh_disks()
        self.run_background(f'Mounting {target}', task, after=after)

    def enable_auto_sleep(self):
        node = self.require_node()
        if not node:
            return
        if str(node.get('rota')) != '1':
            self.log('Auto sleep is mainly for rotating HDDs. Selected drive is not marked as HDD.')
        identity = disk_identity(node)
        minutes = int(self.auto_minutes.get_value_as_int())
        self.auto_sleep[identity] = {'minutes': minutes, 'path': node.get('path'), 'label': disk_label(node)}
        stat = read_block_stat(node.get('path'))
        self.last_stats[identity] = stat
        self.idle_since[identity] = time.time()
        self.log(f'Auto sleep enabled for {node.get("path")} after {minutes} minute(s) idle.')
        self.update_auto_status()
        self.on_selection_changed(self.tree.get_selection())

    def disable_auto_sleep(self):
        node = self.selected_node
        if not node:
            self.log('No disk selected.')
            return
        identity = disk_identity(node)
        if identity in self.auto_sleep:
            self.auto_sleep.pop(identity, None)
            self.last_stats.pop(identity, None)
            self.idle_since.pop(identity, None)
            self.auto_busy.discard(identity)
            self.log('Auto sleep disabled for selected drive.')
        else:
            self.log('Auto sleep was not enabled for selected drive.')
        self.update_auto_status()
        self.on_selection_changed(self.tree.get_selection())

    def enable_auto_wake(self):
        node = self.require_node()
        if not node:
            return
        identity = disk_identity(node)
        path_text = self.wake_path_entry.get_text().strip()
        if not path_text:
            mounts = collect_mountpoints(node)
            path_text = mounts[0] if mounts else ''
        if not path_text:
            self.log('Auto wake needs a mount path. Mount the drive first or type a path like /media/user/DriveName.')
            return
        self.auto_wake[identity] = {'path': path_text, 'label': disk_label(node)}
        self.log(f'Auto wake enabled for selected drive. Watch path: {path_text}')
        self.update_auto_status()
        self.on_selection_changed(self.tree.get_selection())

    def disable_auto_wake(self):
        node = self.selected_node
        if not node:
            self.log('No disk selected.')
            return
        identity = disk_identity(node)
        if identity in self.auto_wake:
            self.auto_wake.pop(identity, None)
            self.log('Auto wake disabled for selected drive.')
        else:
            self.log('Auto wake was not enabled for selected drive.')
        self.update_auto_status()
        self.on_selection_changed(self.tree.get_selection())

    def get_current_node_by_identity(self, identity):
        for node in self.disks:
            if disk_identity(node) == identity:
                return node
        return None

    def auto_manager_tick(self):
        try:
            self.auto_sleep_tick()
            self.auto_wake_tick()
        except Exception as exc:
            self.log(f'Auto manager error: {exc}')
        return True

    def auto_sleep_tick(self):
        now = time.time()
        for identity, policy in list(self.auto_sleep.items()):
            if identity in self.auto_busy:
                continue
            node = self.get_current_node_by_identity(identity)
            if not node or disk_is_system(node) or node.get('_missing'):
                continue
            path = node.get('path')
            if not safe_path(path):
                continue
            stat = read_block_stat(path)
            old = self.last_stats.get(identity)
            if old is None or stat is None or old != stat:
                self.last_stats[identity] = stat
                self.idle_since[identity] = now
                if old is not None and stat is not None and old != stat:
                    self.manual_states[identity] = 'active'
                continue
            idle_seconds = now - self.idle_since.get(identity, now)
            minutes = int(policy.get('minutes') or 15)
            if idle_seconds < minutes * 60:
                continue
            status = disk_status(node)
            if 'standby' in status.lower():
                continue
            self.auto_busy.add(identity)
            self.log(f'Auto sleep: {path} has been idle for {minutes} minute(s). Sending standby command.')
            def task(identity=identity, path=path):
                return self.admin_action(['sleep', path], timeout=180)
            def after(identity=identity):
                self.manual_states[identity] = 'standby'
                node = self.get_current_node_by_identity(identity)
                if node:
                    self.manual_path_states[node.get('path') or ''] = 'standby'
                self.idle_since[identity] = time.time()
                self.auto_busy.discard(identity)
                self.refresh_disks()
            self.run_background('Auto sleep drive standby', task, after=after)
            break

    def auto_wake_tick(self):
        for identity, policy in list(self.auto_wake.items()):
            watch_path = policy.get('path') or ''
            if not watch_path:
                continue
            node = self.get_current_node_by_identity(identity)
            if node:
                mounts = collect_mountpoints(node)
                # If the watched path is already on one of this drive's mounted filesystems, it is available.
                if any(watch_path == mp or watch_path.startswith(mp.rstrip('/') + '/') for mp in mounts):
                    continue
            if identity in self.auto_busy:
                continue
            self.auto_busy.add(identity)
            self.log(f'Auto wake: watched path is not mounted from this drive: {watch_path}. Rescanning and trying to mount.')
            def task():
                return self.admin_action(['rescan'], timeout=180)
            def after(identity=identity):
                self.auto_busy.discard(identity)
                self.refresh_disks()
                node = self.get_current_node_by_identity(identity)
                if node and not disk_is_system(node):
                    self.mount_node(node)
            self.run_background('Auto wake rescan', task, after=after)
            break


def main():
    if os.geteuid() == 0:
        print('Do not run this GUI as root. It uses pkexec only for selected privileged actions.', file=sys.stderr)
    app = DiskPowerApp()
    app.connect('destroy', Gtk.main_quit)
    app.show_all()
    Gtk.main()


if __name__ == '__main__':
    main()
