system_parser.py

# system_parser.py
import datetime
import fnmatch
import os
import pwd
import shutil
import re
import json
import secrets
import string
import subprocess
import tempfile
import time
import uuid

FSTAB_PATH = "/etc/fstab"

# Formatting tools required for each native-pool filesystem option, and the
# Arch/CachyOS package that provides them. AppImage packaging does NOT bundle
# these - they're host prerequisites, same as mergerfs itself, since kernel
# module coupling (zfs, btrfs, ntfs3) and boot-time mounting (fstab/systemd)
# both require the real system package, not a copy frozen inside the bundle.
FILESYSTEM_TOOLS = {
    "ext4": {"binary": "mkfs.ext4", "package_arch": "e2fsprogs"},
    "xfs": {"binary": "mkfs.xfs", "package_arch": "xfsprogs"},
    "btrfs": {"binary": "mkfs.btrfs", "package_arch": "btrfs-progs"},
    "f2fs": {"binary": "mkfs.f2fs", "package_arch": "f2fs-tools"},
    "ntfs": {"binary": "mkfs.ntfs", "package_arch": "ntfs-3g"},
    "zfs": {"binary": "zpool", "package_arch": "zfs-utils (AUR, requires DKMS module)"},
}


def check_filesystem_tool_available(fs_key):
    """Returns True if the mkfs/format binary for fs_key is present on PATH."""
    info = FILESYSTEM_TOOLS.get(fs_key)
    if not info:
        return False
    return shutil.which(info["binary"]) is not None

def generate_pool_mount_name():
    """Returns an unused /mnt/mfs<8 random lowercase letters/digits> pool path, so
    a new pool never defaults to the name of one that already exists. The random
    part starts with a letter so it is also usable as a ZFS pool name."""
    existing = set(parse_fstab_pools().keys())
    alphabet = string.ascii_lowercase + string.digits
    while True:
        token = secrets.choice(string.ascii_lowercase) + "".join(secrets.choice(alphabet) for _ in range(7))
        path = f"/mnt/mfs{token}"
        if path not in existing and not os.path.exists(path):
            return path


# Filesystem label length limits enforced by mkfs (XFS 12, ext4 16); the others
# allow far more than a sensible name needs.
LABEL_LIMITS = {"xfs": 12, "ext4": 16}
NAME_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")


def validate_name(name, what, fs_key=None):
    """Returns an error string if name can't be used as a pool/disk name (which
    becomes a /mnt/<name> folder and, when formatting, a filesystem label)."""
    if not name or not NAME_PATTERN.match(name):
        return f"{what} must start with a letter or number and contain only letters, numbers, '.', '_' and '-'."
    limit = LABEL_LIMITS.get(fs_key)
    if limit and len(name) > limit:
        return f"{what} '{name}' is {len(name)} characters; {fs_key.upper()} labels allow at most {limit}."
    if fs_key == "zfs" and not name[0].isalpha():
        return f"{what} '{name}' must start with a letter to be used as a ZFS pool name."
    return None


def mount_path_in_use(path):
    """True if path is currently a mount point or already the mount point of an
    /etc/fstab entry."""
    path = path.rstrip("/")
    if os.path.ismount(path):
        return True
    try:
        with open(FSTAB_PATH, "r") as f:
            for line in f:
                parts = line.split()
                if len(parts) >= 2 and not parts[0].startswith("#") and parts[1].rstrip("/") == path:
                    return True
    except OSError:
        pass
    return False


def existing_zpool_names():
    """Names of the ZFS pools currently imported on this host."""
    if not shutil.which("zpool"):
        return set()
    res = subprocess.run(["zpool", "list", "-H", "-o", "name"], stdout=subprocess.PIPE,
                         stderr=subprocess.PIPE, text=True)
    return set(res.stdout.split()) if res.returncode == 0 else set()


def disk_name_prefix(pool_path):
    """Prefix for generated disk names: the pool mount's last path segment, minus
    the "mfs" that generate_pool_mount_name() puts in front of its random token
    (so /mnt/mfsab12cd34 -> ab12cd34)."""
    base = os.path.basename(pool_path.strip().rstrip("/"))
    if re.fullmatch(r"mfs[a-z][a-z0-9]{7}", base):
        base = base[len("mfs"):]
    base = re.sub(r"[^A-Za-z0-9._-]", "", base).lstrip("._-")
    return base or "pool"


def next_free_disk_name(prefix, pools):
    """<prefix>-<n> for the smallest n that name_conflict() accepts as a disk name."""
    n = 1
    while name_conflict(f"{prefix}-{n}", "/mnt", pools) or existing_folder_notice(f"{prefix}-{n}", "/mnt"):
        n += 1
    return f"{prefix}-{n}"


def next_free_folder_name(pools):
    """.mfspoolpart, else .mfspoolpart.1, .mfspoolpart.2, ... - the first custom
    pool folder name folder_name_conflict() accepts."""
    candidate, n = ".mfspoolpart", 0
    while folder_name_conflict(candidate, pools):
        n += 1
        candidate = f".mfspoolpart.{n}"
    return candidate


def implied_disk_name(meta):
    """The /mnt folder name an existing filesystem's label implies (never relabels it)."""
    label = re.sub(r"[^A-Za-z0-9._-]+", "_", meta.get("fs_label", "")).lstrip("_.-").rstrip("_")
    return label or f"disk-{meta['uuid'][:8]}"


def mount_spec(fs_type, uid, gid):
    """(fstab fs type, mount options) for a filesystem. lsblk reports the on-disk
    format as "ntfs" regardless of driver; write "ntfs3" - the modern in-kernel
    driver - rather than the legacy ntfs-3g FUSE driver, with the uid/gid/umask
    options NTFS needs and native Linux filesystems don't."""
    if "ntfs" in fs_type:
        return "ntfs3", f"defaults,noatime,lazytime,uid={uid},gid={gid},umask=000,nofail"
    return fs_type, "defaults,noatime,nofail"


def format_disk(dev_path, label, fs_key):
    """Formats dev_path (as a ZFS pool named label for zfs). Returns the new
    filesystem UUID ("" for zfs). Raises on failure, including when the new
    filesystem's UUID can't be read."""
    if fs_key == "zfs":
        subprocess.run(["sudo", "zpool", "create", "-f", label, dev_path], check=True)
        subprocess.run(["sudo", "zfs", "set", f"mountpoint=/mnt/{label}", label], check=True)
    else:
        format_cmds = {
            "ext4": ["sudo", "mkfs.ext4", "-F", "-L", label, dev_path],
            "xfs": ["sudo", "mkfs.xfs", "-f", "-L", label, dev_path],
            "btrfs": ["sudo", "mkfs.btrfs", "-f", "-L", label, dev_path],
            "f2fs": ["sudo", "mkfs.f2fs", "-f", "-l", label, dev_path],
            # mkfs.ntfs comes from ntfs-3g. -f performs a quick format (skip the
            # bad-sector scan); the result is then mounted with the ntfs3 kernel
            # driver, not ntfs-3g's FUSE driver.
            "ntfs": ["sudo", "mkfs.ntfs", "-f", "-L", label, dev_path],
        }
        subprocess.run(format_cmds[fs_key], check=True)

    if fs_key == "zfs":
        return ""

    # The new UUID isn't visible until udev has processed the fresh filesystem;
    # reading it immediately after mkfs can come back empty. Wait for it rather
    # than ever writing a placeholder into fstab.
    subprocess.run(["udevadm", "settle", "--timeout=10"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    for _ in range(20):
        res = subprocess.run(["lsblk", dev_path, "-d", "-o", "UUID", "--noheadings"],
                             stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        if res.returncode == 0 and res.stdout.strip():
            return res.stdout.strip()
        time.sleep(0.5)
    raise RuntimeError(f"{dev_path} was formatted, but no filesystem UUID appeared after 10 seconds, "
                       "so it can't be safely written to fstab.")


def get_target_user():
    """(uid, gid) of the person the pool should belong to: the user who launched
    the app, even when it runs elevated via sudo/pkexec. Falls back to 1000."""
    try:
        if os.environ.get("SUDO_UID"):
            uid = int(os.environ["SUDO_UID"])
            return uid, int(os.environ.get("SUDO_GID", uid))
        if os.environ.get("PKEXEC_UID"):
            uid = int(os.environ["PKEXEC_UID"])
            return uid, pwd.getpwuid(uid).pw_gid
        if os.getuid() != 0:
            return os.getuid(), os.getgid()
    except (ValueError, KeyError):
        pass
    return 1000, 1000


# A pool folder ends up inside an fstab field and a mergerfs ":"-separated
# branch list, so whitespace, ":" and glob characters are all off limits.
FOLDER_NAME_PATTERN = re.compile(r"^[A-Za-z0-9._-]+$")


def validate_pool_folder_name(name):
    """Returns an error string if name can't safely be a pool data folder, else None."""
    if not name:
        return "Enter a folder name."
    if name in (".", "..") or not FOLDER_NAME_PATTERN.match(name) or len(name) > 255:
        return "Folder name may only contain letters, numbers, '.', '_' and '-' (no spaces or slashes)."
    return None


def _list_root_entries(root):
    try:
        return os.listdir(root)
    except PermissionError:
        res = subprocess.run(["sudo", "ls", "-1A", "--", root], stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE, text=True)
        if res.returncode != 0:
            raise OSError(res.stderr.strip() or f"cannot list {root}")
        return res.stdout.splitlines()


def list_device_root_entries(dev_path, mount_point):
    """Names at the root of the filesystem on dev_path. Uses mount_point when the
    device is already mounted, otherwise mounts it read-only on a temp dir for
    the duration of the listing. Raises OSError if it can't be read."""
    if mount_point:
        return _list_root_entries(mount_point)

    tmp = tempfile.mkdtemp(prefix="mergerfs-gui-")
    try:
        res = subprocess.run(["sudo", "mount", "-o", "ro", dev_path, tmp],
                             stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        if res.returncode != 0:
            raise OSError(f"could not mount {dev_path} to check it: {res.stderr.strip()}")
        try:
            return _list_root_entries(tmp)
        finally:
            subprocess.run(["sudo", "umount", tmp], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    finally:
        try:
            os.rmdir(tmp)
        except OSError:
            pass


def folder_exists_on_device(dev_path, mount_point, folder_name):
    """True if folder_name already exists (case-insensitively, as NTFS would treat
    it) at the root of the filesystem on dev_path. Raises OSError if the check
    can't be performed."""
    wanted = folder_name.lower()
    return any(e.lower() == wanted for e in list_device_root_entries(dev_path, mount_point))


def find_pool_folders(entries, pattern):
    """Entries matching a pool folder spec (a plain name or one with '*'), case-insensitively."""
    return [e for e in entries if fnmatch.fnmatchcase(e.lower(), pattern.lower())]


def _split_branch(branch, known_mounts):
    """Splits a mergerfs branch like /mnt/disk1/PoolPart.* into (mount, folder)
    using the longest known disk mount that prefixes it; folder is "" when the
    branch is the disk root."""
    best = None
    for mount in known_mounts:
        if mount and (branch == mount or branch.startswith(mount + "/")):
            if best is None or len(mount) > len(best):
                best = mount
    if best is None:
        return None, ""
    return best, branch[len(best):].strip("/")


def parse_fstab_pools():
    """
    Parses /etc/fstab strictly looking for commented mergerfs block structures.
    Extracts each pool's member disks: their mount point, UUID, and the data
    folder the pool uses inside that disk ("" when the pool uses the disk root).
    A member's folder comes from the block's mergerfs source string, so a pool
    layered onto a disk that another block mounted still lists that disk.
    """
    pools = {}
    fstab_path = FSTAB_PATH
    if not os.path.exists(fstab_path):
        return pools

    blocks = []
    current = None

    try:
        with open(fstab_path, "r") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue

                # Detect the end of a managed block
                if line.startswith("# --- END MERGERFS POOL:"):
                    if current:
                        blocks.append(current)
                    current = None
                    continue

                # Detect the start of a managed block
                if line.startswith("# --- BEGIN MERGERFS POOL:"):
                    match = re.search(r"POOL:\s*(.+)\s*---", line)
                    if match:
                        current = {"name": match.group(1).strip(), "disks": [], "sources": []}
                    continue

                # Parse rows when inside our managed block boundaries
                if current:
                    if line.startswith("#"):
                        continue

                    parts = line.split()
                    if len(parts) < 2:
                        continue

                    if "mergerfs" in line or "fuse.mergerfs" in line:
                        current["sources"] = [b.rstrip("/") for b in parts[0].split(":")]
                        continue

                    # Extract disk mount point from Column 2 (Index 1)
                    mount_point = parts[1].rstrip('/')

                    # Extract the UUID from Column 1 (Index 0), when present.
                    # UUID is the robust identifier per fstab(5) - it survives a
                    # disk being unmounted, renamed, or temporarily disconnected
                    # (relevant here since every managed line uses "nofail").
                    member_uuid = None
                    if parts[0].startswith("UUID="):
                        member_uuid = parts[0].split("=", 1)[1].strip().upper()

                    fs_type = parts[2] if len(parts) > 2 else ""
                    current["disks"].append({"mount": mount_point, "uuid": member_uuid, "fs_type": fs_type})
    except Exception as e:
        print(f"Error parsing fstab blocks: {e}")
        return pools

    # A block may reuse a disk mounted by an earlier block, so resolve branches
    # against every disk mount seen anywhere, not just the block's own.
    known_mounts = {}
    known_fs = {}
    for block in blocks:
        for disk in block["disks"]:
            known_mounts.setdefault(disk["mount"], disk["uuid"])
            known_fs.setdefault(disk["mount"], disk["fs_type"])

    for block in blocks:
        members = {}
        for branch in block["sources"]:
            mount, folder = _split_branch(branch, known_mounts)
            if mount:
                members[mount] = {"mount": mount, "uuid": known_mounts[mount], "folder": folder,
                                  "fs_type": known_fs[mount]}
        for disk in block["disks"]:
            members.setdefault(disk["mount"], {"mount": disk["mount"], "uuid": disk["uuid"], "folder": "",
                                               "fs_type": disk["fs_type"]})

        pool_path = block["name"]
        try:
            usage = shutil.disk_usage(pool_path)
            total, free, used = usage.total, usage.free, usage.used
        except OSError:
            total, free, used = 0, 0, 0

        member_list = sorted(members.values(), key=lambda m: m["mount"])
        pools[pool_path] = {
            "source_branches": [m["mount"] for m in member_list],
            "members": member_list,
            "member_uuids": {m["uuid"] for m in member_list if m["uuid"]},
            "total_bytes": total,
            "used_bytes": used,
            "free_bytes": free
        }
    return pools


def existing_fs_labels():
    """Filesystem labels currently present on any block device."""
    res = subprocess.run(["lsblk", "-a", "-n", "-o", "LABEL"], stdout=subprocess.PIPE,
                         stderr=subprocess.PIPE, text=True)
    return {line.strip() for line in res.stdout.splitlines() if line.strip()} if res.returncode == 0 else set()


def _list_dir(path):
    try:
        return os.listdir(path)
    except OSError:
        return []


def _fstab_mount_points():
    mounts = []
    try:
        with open(FSTAB_PATH, "r") as f:
            for line in f:
                parts = line.split()
                if len(parts) >= 2 and not parts[0].startswith("#"):
                    mounts.append(parts[1].rstrip("/"))
    except OSError:
        pass
    return mounts


def name_conflict(name, parent, pools, check_labels=True, own_path=None):
    """Why `name` may NOT be given to a new pool mount or disk mount inside
    `parent` (a hard block with no override), or None. Names are never reused: it
    is a conflict if ANY pool or pool-member disk, fstab mount point, live mount,
    link or file in `parent`, filesystem label or ZFS pool already carries it
    (case differences count as the same name).

    A plain, unused folder already sitting at that path is NOT a conflict here -
    it is the user's call, see existing_folder_notice().

    check_labels=False is for an existing filesystem being mounted under its own
    label; own_path exempts the path that disk is already mounted at."""
    lname = name.lower()
    parent = parent.rstrip("/") or "/"
    path = f"{parent}/{name}" if parent != "/" else f"/{name}"

    for pool_path, pool in pools.items():
        if os.path.basename(pool_path).lower() == lname:
            return f"a pool named '{os.path.basename(pool_path)}' already exists ({pool_path})"
        for member in pool["members"]:
            if os.path.basename(member["mount"]).lower() == lname:
                return f"a disk mounted at {member['mount']} (pool {pool_path}) already uses the name '{name}'"

    if path != own_path:
        for mount in _fstab_mount_points():
            if os.path.dirname(mount) == parent and os.path.basename(mount).lower() == lname:
                return f"{mount} is already in /etc/fstab"
        for entry in _list_dir(parent):
            full = f"{parent}/{entry}"
            if entry.lower() == lname and (os.path.islink(full) or os.path.ismount(full) or not os.path.isdir(full)):
                return f"{full} already exists and is a mount point, a link or a file"

    if check_labels and any(label.lower() == lname for label in existing_fs_labels()):
        return f"a filesystem labelled '{name}' already exists"
    if any(zpool.lower() == lname for zpool in existing_zpool_names()):
        return f"a ZFS pool named '{name}' already exists"
    return None


def existing_folder_notice(name, parent, own_path=None):
    """A plain folder (not a mount, link or file) already at parent/name in any
    case, or None. Such a folder isn't blocked - the caller asks the user whether
    to reuse it, empty or not. Returns {"path", "empty", "count", "case_only"}."""
    parent = parent.rstrip("/") or "/"
    for entry in _list_dir(parent):
        full = f"{parent}/{entry}"
        if entry.lower() == name.lower() and full != own_path:
            if os.path.islink(full) or os.path.ismount(full) or not os.path.isdir(full):
                return None  # not reusable at all; name_conflict() reports it
            items = _list_dir(full)
            return {"path": full, "empty": not items, "count": len(items), "case_only": entry != name}
    return None


def describe_existing_folder(notice):
    """One line for a confirmation prompt about a folder that already exists."""
    contents = "it is empty" if notice["empty"] else f"it contains {notice['count']} item(s)"
    case = "; it differs from the requested name only by letter case" if notice["case_only"] else ""
    return f"{notice['path']} already exists ({contents}{case})"


REUSE_FOLDER_WARNING = ("A disk or pool mounted on a folder hides whatever is already inside it for as long as "
                        "it is mounted.")


def folder_name_conflict(name, pools):
    """Why `name` may not be the data folder of a NEW custom pool, or None. Data
    folders are unique per pool: a name already used by any pool (or matched by a
    pool's pattern such as DrivePool's PoolPart.*) can only be reached by
    explicitly adding a drive to that pool."""
    lname = name.lower()
    if lname.startswith("poolpart"):
        return "names starting with 'PoolPart' are reserved for Windows DrivePool folders"
    for pool_path, pool in pools.items():
        for member in pool["members"]:
            folder = member["folder"]
            if folder and find_pool_folders([name], folder):
                return f"'{name}' is already used as a data folder by pool {pool_path} ({folder})"
    return None


def audit_pools(pools):
    """Cross-checks parsed pools for name/disk reuse that should be impossible:
    returns a list of problem descriptions (empty when everything is separate).
    Sharing a disk between pools is fine only through different folders."""
    problems = []
    paths = list(pools)
    for i, a in enumerate(paths):
        for b in paths[i + 1:]:
            if os.path.basename(a).lower() == os.path.basename(b).lower():
                problems.append(f"pools {a} and {b} share the name '{os.path.basename(a)}'")
            elif a.startswith(b + "/") or b.startswith(a + "/"):
                problems.append(f"pools {a} and {b} are nested")

    mounts_by_uuid, pools_by_slot, pools_by_folder = {}, {}, {}
    for pool_path, pool in pools.items():
        for member in pool["members"]:
            mount, folder = member["mount"], member["folder"]
            if pool_path == mount or pool_path.startswith(mount + "/") or mount.startswith(pool_path + "/"):
                problems.append(f"pool {pool_path} overlaps its own disk mount {mount}")
            if member["uuid"]:
                mounts_by_uuid.setdefault(member["uuid"], set()).add(mount)
            pools_by_slot.setdefault((mount.lower(), folder.lower()), set()).add(pool_path)
            if folder and "*" not in folder:
                pools_by_folder.setdefault(folder.lower(), set()).add(pool_path)

    for uuid_, mounts in mounts_by_uuid.items():
        if len(mounts) > 1:
            problems.append(f"filesystem {uuid_} is mounted at several paths: {', '.join(sorted(mounts))}")
    for (mount, folder), owners in pools_by_slot.items():
        if len(owners) > 1:
            where = f"{mount}/{folder}" if folder else mount
            problems.append(f"{where} belongs to several pools: {', '.join(sorted(owners))}")
    for folder, owners in pools_by_folder.items():
        if len(owners) > 1:
            problems.append(f"data folder '{folder}' is used by several pools: {', '.join(sorted(owners))}")
    for pool_path, pool in pools.items():
        for member in pool["members"]:
            if "*" in member["folder"]:
                for folder, owners in pools_by_folder.items():
                    if find_pool_folders([folder], member["folder"]) and (owners - {pool_path}):
                        problems.append(f"folder '{folder}' of {', '.join(sorted(owners))} matches the pattern "
                                        f"'{member['folder']}' used by {pool_path}")
    return problems


POOL_PATH_PATTERN = re.compile(r"^/[A-Za-z0-9._-]+(/[A-Za-z0-9._-]+)*$")
SYSTEM_TOP_LEVEL_DIRS = {"bin", "boot", "dev", "etc", "lib", "lib64", "proc", "run", "sbin", "sys", "usr", "var"}


def validate_pool_mount_path(path, pools, current=None):
    """Returns an error string if path can't be a new pool mount point, else None.
    `current` is the pool's own existing path when renaming (it is allowed to
    disappear, not to be reused)."""
    if not POOL_PATH_PATTERN.match(path) or any(seg in (".", "..") for seg in path.split("/")):
        return ("Enter an absolute path such as /mnt/mypool, using only letters, numbers, '.', '_' and '-' "
                "(no spaces).")
    if current is not None and path == current:
        return "That is already this pool's name."
    if path in PROTECTED_MOUNTPOINTS or path.split("/")[1] in SYSTEM_TOP_LEVEL_DIRS:
        return f"{path} is a system location and can't be a pool."

    related = list(pools) + [m["mount"] for pool in pools.values() for m in pool["members"]]  # incl. the pool's own path: no nesting inside itself
    resolved = os.path.realpath(path)  # a symlink must not smuggle in an existing pool/disk path
    for other in related:
        for candidate in (path, resolved):
            if candidate == other or candidate.startswith(other + "/") or other.startswith(candidate + "/"):
                return f"{path} overlaps {other}, which is already used by a pool."
    if os.path.ismount(path) or mount_path_in_use(path):
        return f"{path} is already a mount point."
    conflict = name_conflict(os.path.basename(path), os.path.dirname(path), pools)
    if conflict:
        return f"That name can't be used: {conflict}."
    return None


def pool_folder_spec(pool_data):
    """The data folder every member of a pool uses ("" = disk root, "PoolPart.*",
    ".mfspoolpart", ...), or None if members disagree or the pool has none."""
    folders = {m["folder"] for m in pool_data["members"]}
    return folders.pop() if len(folders) == 1 else None


def add_eligibility(pool_path, pools, meta):
    """(ok, reason) for adding the unassigned-inventory disk `meta` to pool_path.
    Uses only fstab/inventory data, so it is cheap enough to evaluate per button."""
    pool = pools.get(pool_path)
    if pool is None:
        return False, "Select a pool first."
    if pool_path in meta.get("shared_with", []):
        return False, "This drive is already in the pool."
    if meta["fs_type"] == "zfs_member":
        return False, "ZFS member disks aren't supported here."
    folder = pool_folder_spec(pool)
    if folder is None:
        return False, "This pool's drives use different data folders, so a new drive's folder is ambiguous."
    if re.search(r"[?\[\]]", folder):
        return False, "This pool's folder pattern is too complex to extend automatically."

    others = meta.get("shared_with", [])
    if others:
        # The disk already hosts other pools' folders; it can take another
        # folder, but never this pool's own (that would pool the same data twice)
        # and never the disk root.
        if not folder:
            return False, f"This pool uses whole disks, but this drive already hosts a folder for {', '.join(others)}."
        for other in others:
            for member in pools[other]["members"]:
                if member["mount"] == meta["fstab_mount"] and (
                        member["folder"].lower() == folder.lower() or find_pool_folders([member["folder"]], folder)):
                    return False, f"This drive already belongs to {other} through the same kind of folder."
    return True, ""


def _find_pool_block(lines, pool_path):
    """(begin, end) line indexes of pool_path's managed block, or raises ValueError."""
    begin = end = None
    for i, line in enumerate(lines):
        stripped = line.strip()
        if stripped.startswith("# --- BEGIN MERGERFS POOL:"):
            m = re.search(r"POOL:\s*(.+)\s*---", stripped)
            if m and m.group(1).strip() == pool_path:
                begin = i
        elif begin is not None and stripped.startswith("# --- END MERGERFS POOL:"):
            end = i
            break
    if begin is None or end is None:
        raise ValueError(f"Could not find the fstab block for {pool_path}.")
    return begin, end


def _find_mergerfs_line(lines, begin, end, pool_path):
    for i in range(begin + 1, end):
        parts = lines[i].split()
        if len(parts) >= 4 and not parts[0].startswith("#") and parts[2] == "fuse.mergerfs":
            return i
    raise ValueError(f"The block for {pool_path} has no fuse.mergerfs line.")


def _write_fstab_with_backup(lines, fstab_path):
    """Replaces fstab_path with `lines` atomically, keeping the old file as a
    timestamped backup whose path is returned."""
    backup = f"{fstab_path}.mergerfs-gui.{datetime.datetime.now():%Y%m%d-%H%M%S}.bak"
    shutil.copy2(fstab_path, backup)
    tmp = fstab_path + ".mergerfs-gui.tmp"
    with open(tmp, "w") as f:
        f.writelines(lines)
    shutil.copymode(fstab_path, tmp)
    os.replace(tmp, fstab_path)
    return backup


def add_branch_to_pool_block(pool_path, branch, disk_mount, disk_line=None, fstab_path=None):
    """Adds `branch` to pool_path's mergerfs line (and its requires-mounts-for
    option), and inserts `disk_line` into the block when the disk isn't mounted
    by another block. The old file is kept as a timestamped backup, and the new
    one replaces it atomically. Returns the backup path; raises ValueError if the
    block or its mergerfs line can't be found or the branch is already present."""
    fstab_path = fstab_path or FSTAB_PATH
    with open(fstab_path, "r") as f:
        lines = f.readlines()

    begin, end = _find_pool_block(lines, pool_path)
    merger_idx = _find_mergerfs_line(lines, begin, end, pool_path)

    parts = lines[merger_idx].split()
    sources = parts[0].split(":")
    if branch.rstrip("/") in [b.rstrip("/") for b in sources]:
        raise ValueError(f"{branch} is already part of {pool_path}.")
    parts[0] = ":".join(sources + [branch])
    requirement = f"x-systemd.requires-mounts-for={disk_mount}"
    if requirement not in parts[3].split(","):
        parts[3] += "," + requirement
    lines[merger_idx] = " ".join(parts) + "\n"
    if disk_line:
        lines.insert(merger_idx, disk_line + "\n")

    return _write_fstab_with_backup(lines, fstab_path)


def rename_pool_in_fstab(old_path, new_path, fstab_path=None):
    """Renames a pool - its identity is its mount path - in its BEGIN/END marker
    comments and on its fuse.mergerfs line. Backup and atomic replace as for
    add_branch_to_pool_block; returns the backup path."""
    fstab_path = fstab_path or FSTAB_PATH
    with open(fstab_path, "r") as f:
        lines = f.readlines()

    begin, end = _find_pool_block(lines, old_path)
    merger_idx = _find_mergerfs_line(lines, begin, end, old_path)

    parts = lines[merger_idx].split()
    if parts[1].rstrip("/") != old_path:
        raise ValueError(f"The mergerfs line for {old_path} mounts {parts[1]} instead; refusing to rename it.")
    parts[1] = new_path
    lines[merger_idx] = " ".join(parts) + "\n"
    for i in (begin, end):
        lines[i] = lines[i].replace(f"POOL: {old_path} ---", f"POOL: {new_path} ---")

    return _write_fstab_with_backup(lines, fstab_path)


def restore_fstab_backup(backup, fstab_path=None):
    """Puts a backup made by add_branch_to_pool_block back in place, atomically."""
    fstab_path = fstab_path or FSTAB_PATH
    tmp = fstab_path + ".mergerfs-gui.tmp"
    shutil.copy2(backup, tmp)
    os.replace(tmp, fstab_path)


def add_branch_to_running_pool(pool_path, branch):
    """Appends `branch` to a mounted pool through mergerfs's runtime control file,
    so it takes effect without unmounting. Returns None on success or a short
    reason string on failure (the fstab change still applies at next mount)."""
    control = os.path.join(pool_path, ".mergerfs")
    value = f"+>{branch}"
    try:
        os.setxattr(control, "user.mergerfs.branches", value.encode())
        return None
    except OSError as first_err:
        if shutil.which("setfattr"):
            res = subprocess.run(["sudo", "setfattr", "-n", "user.mergerfs.branches", "-v", value, control],
                                 stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
            if res.returncode == 0:
                return None
            return res.stderr.strip() or str(first_err)
        return str(first_err)


# Partition-type names (from lsblk's PARTTYPENAME, the GPT type GUID's human
# label) that unambiguously mark a partition as part of a boot/OS layout
# rather than data storage. Matched case-insensitively as a substring.
SYSTEM_PARTITION_TYPE_MARKERS = [
    "efi system",
    "microsoft reserved",
    "bios boot",
    "windows recovery",
]

# Mountpoints that belong to a running Linux system's own root filesystem
# hierarchy. Anything mounted here is part of an OS install, not spare
# storage, regardless of its filesystem type.
PROTECTED_MOUNTPOINTS = {
    "/", "/boot", "/boot/efi", "/efi", "/home", "/root", "/srv",
    "/usr", "/var", "/var/cache", "/var/log", "/var/tmp", "/tmp", "/opt",
}


def _is_live_system_partition(node):
    """True if this partition is swap or mounted on a core system path - i.e. in
    use by the running OS right now, whatever else is on its disk."""
    if (node.get("fstype") or "").strip().lower() == "swap":
        return True
    mount = node.get("mountpoint")
    if mount:
        normalized = mount.rstrip("/") or "/"
        if normalized in PROTECTED_MOUNTPOINTS:
            return True
    return False


def _is_boot_or_os_partition(node):
    """True if this single partition is unambiguously part of a boot/OS layout."""
    parttypename = (node.get("parttypename") or "").strip().lower()
    if any(marker in parttypename for marker in SYSTEM_PARTITION_TYPE_MARKERS):
        return True
    return _is_live_system_partition(node)


def _collect_leaf_nodes(node, leaves):
    """Recurse to the leaf partitions/disks - the actual usable devices."""
    children = node.get("children")
    if children:
        for child in children:
            _collect_leaf_nodes(child, leaves)
    else:
        leaves.append(node)


def get_all_system_devices(trusted_uuids=frozenset()):
    """
    Executes lsblk to query every storage partition attached to the kernel.
    Captures filesystems, mount states, sizes, UUIDs, and friendly text labels.

    A whole physical disk is excluded entirely if ANY of its partitions looks
    like part of a boot/OS layout (EFI System, Microsoft Reserved, BIOS boot,
    Windows recovery, swap, or something mounted at a core system path like
    "/" or "/home"). This is done at the whole-disk level, not per-partition,
    because a Windows disk's main data partition ("C:") is plain NTFS with no
    marker of its own - it can only be told apart from a legitimate NTFS
    storage disk by the ESP/MSR/recovery siblings sitting next to it on the
    same physical disk. This keeps Windows and Linux system disks out of the
    pool-candidate list while leaving genuine extra storage disks untouched.

    trusted_uuids are filesystem UUIDs already configured as pool members. A
    Windows-created data disk (DrivePool) carries the same MSR partition as a
    Windows system disk, so its marker partitions are overlooked when one of
    its partitions is a known pool member - unless something on it is swap or
    mounted at a core system path.
    """
    devices = []
    try:
        cmd = ["lsblk", "-p", "-a", "-o", "NAME,FSTYPE,MOUNTPOINT,LABEL,UUID,SIZE,PARTTYPENAME,TYPE", "--json"]
        result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        
        if result.returncode == 0 and result.stdout.strip():
            data = json.loads(result.stdout)
            
            for top_node in data.get("blockdevices", []):
                leaves = []
                _collect_leaf_nodes(top_node, leaves)

                if any(_is_boot_or_os_partition(leaf) for leaf in leaves):
                    is_trusted_pool_disk = any(
                        (leaf.get("uuid") or "").upper() in trusted_uuids for leaf in leaves
                    ) and not any(_is_live_system_partition(leaf) for leaf in leaves)
                    if not is_trusted_pool_disk:
                        continue  # this whole physical disk is an OS/boot disk

                for leaf in leaves:
                    if leaf.get("type") == "rom":
                        continue  # optical drive - never pool storage
                    if _is_boot_or_os_partition(leaf):
                        continue  # e.g. the MSR partition on a trusted DrivePool disk
                    name = leaf.get("name") or ""
                    # lsblk --json always includes these keys, but sets them to
                    # JSON null (Python None) for a disk with no filesystem/label/
                    # mount/UUID - `.get(key, default)` does NOT catch that, since
                    # the key is present, just null. Use `or` to coalesce None too,
                    # otherwise a fully raw disk's fs_type is None (not ""), fails
                    # the allow-list check below, and silently vanishes.
                    fs_type = leaf.get("fstype") or ""
                    mount = leaf.get("mountpoint") or ""
                    label = leaf.get("label") or ""
                    uuid = leaf.get("uuid") or ""
                    size = leaf.get("size") or "0B"

                    if name and not name.startswith("/dev/loop"):
                        # Support raw unformatted partitions or alternative file systems
                        if fs_type in ["", "ntfs", "ntfs3", "ext4", "xfs", "btrfs", "f2fs", "zfs_member"]:
                            devices.append({
                                "dev_path": name,
                                "fs_type": fs_type if fs_type else "Raw / Unformatted",
                                "mount_point": mount if mount else "",
                                "label": label if label else "",
                                "uuid": uuid if uuid else "N/A",
                                "size_str": size
                            })
    except Exception as e:
        print(f"Failed block inquiry utility: {e}")
    return devices




def get_categorized_inventory():
    """
    Categorizes system components into Active Pool Members or Unassigned Disks.

    A disk used by a pool at its root is fully claimed and hidden. A disk that
    every claiming pool uses only through a subfolder (e.g. DrivePool's
    PoolPart.*, or a custom folder) still has room for another pool's folder,
    so it stays listed, tagged with "shared_with" and its existing "fstab_mount".
    """
    pools = parse_fstab_pools()
    trusted_uuids = {u for p in pools.values() for u in p["member_uuids"]}
    all_devices = get_all_system_devices(trusted_uuids)

    # The filesystem actually on each member disk (lsblk) beats the fstab type
    # column, which may name a driver ("ntfs3") rather than the filesystem.
    fs_by_uuid = {d["uuid"].strip().upper(): d["fs_type"] for d in all_devices
                  if d["uuid"] and d["uuid"] != "N/A"}
    for pool_data in pools.values():
        for m in pool_data["members"]:
            m["fs_type"] = fs_by_uuid.get(m["uuid"]) or m.get("fs_type") or "unknown"

    unassigned_drives = []
    for dev in all_devices:
        mnt = dev["mount_point"].rstrip('/')
        dev_uuid = dev["uuid"].strip().upper() if dev["uuid"] and dev["uuid"] != "N/A" else None

        if mnt and mnt in pools:
            continue

        # A disk already configured as a pool member is "claimed" even if it
        # isn't currently mounted (nofail means offline/missing disks are
        # expected). Checking the mountpoint alone would otherwise let an
        # already-pooled-but-unmounted disk show up as "unassigned" and get
        # reformatted from the pool-creation screen.
        claims = [
            (pool_name, m)
            for pool_name, pool_data in pools.items()
            for m in pool_data["members"]
            if (dev_uuid and m["uuid"] == dev_uuid) or (mnt and m["mount"] == mnt)
        ]
        if claims and not all(m["folder"] for _, m in claims):
            continue

        shared_with = sorted({pool_name for pool_name, _ in claims})
        fstab_mount = claims[0][1]["mount"] if claims else ""

        if shared_with:
            status = f"Mounted at {dev['mount_point']}" if dev['mount_point'] else "Not mounted"
            status += f" | folder pool(s): {', '.join(shared_with)}"
        else:
            status = f"Mounted at {dev['mount_point']}" if dev['mount_point'] else "Unmounted Raw Disk"

        display_label = dev['label'].strip() if dev['label'] else "[No Label]"
        label_text = f"🏷️  {display_label} — {dev['dev_path']} ({dev['size_str']}) [{dev['fs_type']}] — {status} | UUID: {dev['uuid']}"

        unassigned_drives.append({
            "label": label_text,
            "uuid": dev["uuid"],
            "dev_path": dev["dev_path"],
            "fs_type": dev["fs_type"],
            "current_mount": dev["mount_point"],
            "fs_label": dev["label"].strip(),
            "shared_with": shared_with,
            "fstab_mount": fstab_mount
        })

    return pools, unassigned_drives