#!/usr/bin/env python3
# subvols -- list btrfs subvolumes
import argparse
import glob
import json
import os
import re
import subprocess
import sys
import pprint

def enum_filesystems():
    cmd = ["findmnt", "-J", "-l", "-t", "btrfs"]
    with subprocess.Popen(cmd, stdout=subprocess.PIPE) as proc:
        data = json.load(proc.stdout)
        for row in data["filesystems"]:
            if "subvol=/" in row["options"].split(","):
                yield row["target"]

def enum_subvolumes(root):
    cmd = ["sudo", "btrfs", "subvolume", "list", "-a", root]
    yield (5, 0, "")
    with subprocess.Popen(cmd, stdout=subprocess.PIPE) as proc:
        for line in proc.stdout:
            line = line.decode().strip()
            m = re.match(r"^ID (\d+) gen \d+ top level (\d+) path (.+)$", line)
            subvol_id = int(m.group(1))
            parent_id = int(m.group(2))
            path = m.group(3)
            if path.startswith("<FS_TREE>/"):
                path = path[len("<FS_TREE>/"):]
            yield subvol_id, parent_id, path

# http://storaged.org/doc/udisks2-api/latest/gdbus-org.freedesktop.UDisks2.Manager.html#gdbus-method-org-freedesktop-UDisks2-Manager.ResolveDevice
# ResolveDevice(devspec = {"path": "/dev/foo"}, options = {})
# /org/freedesktop/UDisks2/block_devices/dm_2d0
# org.freedesktop.UDisks2.Filesystem.BTRFS
# GetSubvolumes(snapshots=False, {})

parser = argparse.ArgumentParser()
parser.add_argument("-o", "--output", metavar="FIELDS",
                        help="output columns")
parser.add_argument("-s", "--sort", action="store_true",
                        help="sort by path instead of ID")
parser.add_argument("-f", "--full-path", action="count", default=0,
                        help="show full paths of each subvolume (-ff for absolute paths)")
parser.add_argument("-v", "--verbose", action="count", default=0,
                        help="show subvol IDs")
parser.add_argument("path", nargs="*",
                        help="Btrfs filesystem root to inspect")
args = parser.parse_args()

mtpts = args.path
if not mtpts:
    mtpts = glob.glob("/run/btrfs/*/")
if not mtpts:
    mtpts = enum_filesystems()
mtpts = [arg.rstrip("/") or "/" for arg in mtpts]

known_columns = ["NAME", "ID", "PARENT", "ABS_ID", "ABS_PARENT"]
if args.verbose:
    visible_columns = ["NAME", "ID", "PARENT"]
    hidden_columns = ["ABS_ID", "ABS_PARENT"]
else:
    visible_columns = ["NAME"]
    hidden_columns = ["ID", "PARENT", "ABS_ID", "ABS_PARENT"]

if arg := args.output:
    if arg.startswith("+"):
        arg = arg[1:]
    else:
        visible_columns = []
    arg = arg.upper().split(",")
    for c in arg:
        if c in known_columns:
            if c in visible_columns:
                # `column` does not like duplicate items in --table-order
                exit("duplicate column: %s" % c)
            if c in hidden_columns:
                hidden_columns.remove(c)
            visible_columns.append(c)
        else:
            exit("unknown column: %s" % c)

for i, mtpt in enumerate(mtpts):
    names = {}
    parents = {}

    for subvol_id, parent_id, subvol_name in enum_subvolumes(mtpt):
        names[mtpt, subvol_id] = subvol_name
        parents[mtpt, subvol_id] = parent_id

    if args.sort:
        # Sort by mtpt then subvol_name
        snames = sorted(names, key=lambda i: (i[0], names[i], i[1]))
    else:
        snames = names

    if i > 0:
        print()

    cmd = [
        "column",
        "--separator=\t",
        "--table",
        "--table-name=subvolumes",
        "--table-columns=NAME,ID,PARENT,ABS_ID,ABS_PARENT",
        "--table-order=%s" % ",".join(visible_columns),
        "--table-hide=%s" % ",".join(hidden_columns),
        "--tree=NAME",
        "--tree-id=ID",
        "--tree-parent=PARENT",
    ]

    with subprocess.Popen(cmd, stdin=subprocess.PIPE) as proc:
        for mtpt, subvol_id in snames:
            subvol_name = names[mtpt, subvol_id]
            parent_id = parents[mtpt, subvol_id]
            if not parent_id:
                assert not subvol_name
                name = mtpt
            elif args.full_path >= 2:
                name = os.path.join(mtpt, subvol_name)
            elif args.full_path == 1:
                name = subvol_name
            else:
                parent_name = names[mtpt, parent_id]
                name = os.path.relpath(os.path.join(mtpt, subvol_name),
                                       os.path.join(mtpt, parent_name))
            line = [
                name,
                subvol_id,
                parent_id,
                "%s@%s" % (mtpt, subvol_id),
                "%s@%s" % (mtpt, parent_id),
            ]
            line = "\t".join(map(str, line)) + "\n"
            line = line.encode()
            proc.stdin.write(line)
        proc.stdin.close()
