#!/usr/bin/env python3
import argparse
from dataclasses import dataclass
import enum
from functools import cache
import json
from nullroute.core import Core, Env
from nullroute.string import fmt_duration, fmt_size_short
import os
import re
from socket import gethostname
import subprocess
import sys
import time

level_colors = {
    "red":      (196,  9),
    "yellow":   (226, 11),
    "green":    ( 76, 10),
    "gray":     (242,  7),
}

levels_small = [
    ( 10, "green"),
    (  5, "yellow"),
    (  0, "red"),
    ( -1, "gray"),
]

levels_big = [
    ( 10, "green"),
    (  2, "yellow"),
    (  0, "red"),
    ( -1, "gray"),
]

class Kind(enum.Enum):
    NONE = 0        # Undetermined
    ROOT = 1        # The root filesystem
    HOME = 2        # The /home filesystem
    MEDIA = 3       # Removable media (cache as host-independent)
    MEMORY = 4      # In-memory tmpfs
    NETWORK = 5     # Network filesystems (not cached)
    SYSTEM = 6      # Real but "internal" filesystems (like /boot)
    KERNEL = 7      # Virtual pseudo-filesystems (like /proc)

fs_kind_priorities = {
    Kind.HOME:      100,
    Kind.MEDIA:     50,
    Kind.NONE:      1,
    Kind.ROOT:      0,
    Kind.SYSTEM:    -1,
    Kind.NETWORK:   -10,
    Kind.MEMORY:    -100,
    Kind.KERNEL:    -1000,
}

fs_kind_groups = {
    Kind.HOME:      1,
    Kind.MEDIA:     4,
    Kind.NONE:      1,
    Kind.ROOT:      1,
    Kind.NETWORK:   5,
    Kind.MEMORY:    2,
    Kind.KERNEL:    3,
}

fs_kind_prefixes = [
    ("//%s/" % gethostname(),
                    Kind.KERNEL),
    ("//",          Kind.NETWORK),
    ("/boot/",      Kind.SYSTEM),
    ("/dev/",       Kind.KERNEL),
    ("/efi/",       Kind.SYSTEM),
    ("/media/",     Kind.MEDIA),
    ("/n/",         Kind.NETWORK),
    ("/net/",       Kind.NETWORK),
    ("/proc/",      Kind.KERNEL),
    ("/run/media/", Kind.MEDIA),
    ("/run/user/",  Kind.MEMORY),
    ("/sys/",       Kind.KERNEL),
    ("/tmp/",       Kind.MEMORY),
    ("/vol3",       Kind.MEDIA),
    ("/vol4",       Kind.MEDIA),
    ("/home",       Kind.HOME),
    ("/",           Kind.ROOT),
]

fs_kind_fstypes = {
    "afs":          Kind.NETWORK,
    "fuse.sshfs":   Kind.NETWORK,
    "nfs":          Kind.NETWORK,
    "nfs4":         Kind.NETWORK,
    "cgroup":       Kind.KERNEL,
    "devpts":       Kind.KERNEL,
    "devtmpfs":     Kind.KERNEL,
    "ecryptfs":     Kind.KERNEL,
    "efivarfs":     Kind.KERNEL,
    "rootfs":       Kind.KERNEL,
    "tmpfs":        Kind.MEMORY,
}

# cached and shown even if not present
fs_cache_kinds = {
    Kind.ROOT,
    Kind.HOME,
}
fs_cond_cache_kinds = {
    Kind.MEDIA,
}

class PersistentDict(dict):
    def __init__(self, path):
        self.path = path
        self.load()

    def load(self):
        try:
            with open(self.path, "r") as fh:
                data = json.load(fh)
            self.update(data)
        except FileNotFoundError as e:
            Core.debug("cache load failed: %s", e)
        except json.decoder.JSONDecodeError as e:
            Core.warn("cache load failed: %s", e)
            os.unlink(self.path)

    def flush(self):
        try:
            with open(self.path, "w") as fh:
                json.dump(self, fh)
        except Exception as e:
            Core.debug("cache save failed (%r)" % e)

class Partition:
    # These filesystems will be hidden from all views, as their 'capacity' is
    # either mostly (for tmpfs) or completely irrelevant.
    IGNORED_FSTYPES = {
        "afs", "cgroup", "devpts", "devtmpfs", "fuse.revokefs-fuse",
        "ecryptfs", "efivarfs", "rootfs", "tmpfs",
    }

    # Exceptions for IGNORED_FSTYPES; e.g. /tmp is user-facing and should be
    # shown in "--all" even if tmpfs, although it remains hidden by default.
    IGNORED_FSTYPES_EXCEPTIONS = {
        "/run/user/%d" % os.getuid(),
        "/tmp",
    }

    def __init__(self, mtpt, fstype, cached=None):
        self.mtpt   = mtpt
        self.fstype = fstype
        self.cached = cached
        self.kind   = fs_kind_from_path(mtpt, fstype)

    def __repr__(self):
        return "Partition(mtpt=%r, fstype=%r, cached=%r, kind=%r)" \
               % (self.mtpt, self.fstype, self.cached, self.kind)

    def can_skip(self, *, verbose=False):
        if self.kind in {Kind.KERNEL}:
            return True
        if not verbose and self.kind in {Kind.SYSTEM, Kind.MEMORY, Kind.NETWORK}:
            return True
        if self.fstype in self.IGNORED_FSTYPES:
            if verbose and self.mtpt in self.IGNORED_FSTYPES_EXCEPTIONS:
                return False
            return True
        if self.fstype == "zfs":
            # For ZFS, only show the whole pool
            if self.cached:
                # If reading from cache, it's either a pool (which we won't
                # skip) or another host's entry (which we can't check).
                return False
            tmp = _get_zfs_dataset_info(self.mtpt, "used,avail,name")
            tmp = tmp[2].split("/")
            if len(tmp) > 1:
                return True
        if self.kind != Kind.NETWORK:
            if os.path.exists("%s/.diskuse.ignore" % self.mtpt):
                return True
        return False

    def can_cache(self):
        if self.kind == Kind.NETWORK:
            return False
        if os.path.exists("%s/.diskuse.nocache" % self.mtpt):
            return False
        if self.kind in fs_cond_cache_kinds:
            if os.path.exists("%s/.diskuse.cache" % self.mtpt):
                return True
            # When I forget to add an explicit cache tag
            if os.path.exists("%s/Attic" % self.mtpt):
                return True
        if self.kind in fs_cache_kinds:
            return True
        return False

    @property
    @cache
    def devno(self):
        return get_fs_devno(self.mtpt)

    def get_fstat(self):
        try:
            st = os.statvfs(self.mtpt)
        except (PermissionError, OSError) as e:
            Core.debug("skipping %r: could not statvfs: %r", self.mtpt, e)
            return None

        if st.f_blocks == 0:
            Core.debug("skipping %r: f_blocks == 0", self.mtpt)
            return None

        if self.fstype == "zfs":
            tmp = _get_zfs_dataset_info(self.mtpt, "used,avail,name")
            used = int(tmp[0])
            avail = int(tmp[1])
            pool = tmp[2].split("/")[0]

            tmp = _get_zfs_dataset_info(pool, "used,avail,name")
            pool_used = int(tmp[0])
            pool_avail = int(tmp[1])

            pool_total = pool_used + pool_avail
            return {"total": pool_total,
                    "free":  pool_total - used,
                    "avail": avail}

        return {"total": st.f_bsize * st.f_blocks,
                "free":  st.f_bsize * st.f_bfree,
                "avail": st.f_bsize * st.f_bavail}

def get_mounts_from_mtab():
    def octal_unescape(val):
        out = ""; acc = 0; state = -1
        for c in val:
            if state < 0:
                if c == "\\":
                    acc = 0; state = 0
                else:
                    out += c
            else:
                acc = (acc << 3) | int(c); state += 1
                if state == 3:
                    out += chr(acc); state = -1
        return out

    with open("/etc/mtab") as fh:
        for line in fh:
            dev, mtpt, fstype, rest = line.strip().split(None, 3)
            mtpt = octal_unescape(mtpt)
            yield dev, mtpt, fstype

def get_mounts_from_findmnt():
    buf = _get_cmd_output("findmnt", "--json", "--list")
    buf = json.loads(buf)
    for fs in buf["filesystems"]:
        yield fs["source"], fs["target"], fs["fstype"]

class Enumerator():
    def __init__(self):
        self.cache = PersistentDict(Env.find_cache_file("diskuse.json"))
        self.cache_fixed_maxage = 14*86400
        self.cache_media_maxage = 90*86400

    def __enter__(self):
        return self

    def __exit__(self, *_):
        self.cache.flush()

    @cache
    def _get_mounts(self):
        #return [*get_mounts_from_findmnt()]
        return [*get_mounts_from_mtab()]

    def _find_mtpt_upwards(self, base):
        base = os.path.abspath(base) + "/"
        candidate = "/"
        for dev, mtpt, fstype in self._get_mounts():
            mtpt = mtpt.rstrip("/") + "/"
            if base == mtpt or base.startswith(mtpt):
                Core.trace("mtpt %r matches prefix %r", mtpt, base)
                if len(mtpt) > len(candidate):
                    candidate = mtpt
        return candidate

    def _enum_partitions(self, show_cached=False):
        for dev, mtpt, fstype in self._get_mounts():
            yield Partition(mtpt, fstype)
        if show_cached:
            to_purge = []
            for mtpt, rest in self.cache.items():
                if fs_kind_from_path(mtpt, rest["type"]) == Kind.MEDIA:
                    Core.debug("mtpt %r is media, using media maxage" % mtpt)
                    maxage = self.cache_media_maxage
                else:
                    maxage = self.cache_fixed_maxage
                age = time.time() - rest.get("time", 0)
                if age <= maxage:
                    yield Partition(mtpt, rest["type"], cached=rest)
                else:
                    Core.debug("cache entry for %r has expired (%s old)",
                               mtpt, fmt_duration(age))
                    to_purge.append(mtpt)
            for mtpt in to_purge:
                del self.cache[mtpt]

    def enum_partitions_system(self, show_cached=False, show_all=False):
        for part in self._enum_partitions(show_cached):
            if part.can_skip(verbose=show_all):
                Core.trace("skipping %r: is marked as skippable", part)
                continue
            else:
                Core.trace("adding %r", part)
                yield part

    def enum_partitions_exact(self, paths):
        """
        Return only the containing mounts for the paths specified in CLI.

        All 'skip' rules are ignored for explicitly given paths.
        """
        for path in paths:
            mtpt = self._find_mtpt_upwards(path)
            yield Partition(mtpt, fstype=None)

    def _add_to_cache(self, part, data):
        Core.debug("caching %r", part)
        data = {**data,
                "type": part.fstype,
                "time": int(time.time())}
        mtpt = part.mtpt
        if part.kind != Kind.MEDIA:
            assert mtpt.startswith("/")
            mtpt = "//" + gethostname() + mtpt
        self.cache[mtpt] = data

    def get_partition_data(self, part):
        if part.cached:
            return part.cached
        else:
            data = part.get_fstat()
            if data and (part.cached or part.can_cache()):
                self._add_to_cache(part, data)
            return data

class Table():
    def __init__(self):
        self.columns = []
        self.rows = []
        self.can_color = True
        self.can_unicode = not (os.getenv("TERM") == "linux")

    def _make_header(self):
        out = ""
        for i, (head, type, width) in enumerate(self.columns):
            out += " %*s" % (width, head)
        return fmt_ansi(out, "1")

    def _make_separator(self):
        cols = 0
        for (head, type, width) in self.columns:
            cols += 1 + abs(width)
        return " " + fmt_ansi("-" * (cols-1), "2")

    def _make_gauge(self, width, level, color):
        bright = color
        dark = darken(color, 1)
        if self.can_unicode:
            if self.can_color:
                bars = "■", "■", "■"
            else:
                bars = "■", "■", "□"
        else:
            bars = "##-"
        return gauge3(width, level, *bars,
                      full_fmt="38;5;%d" % bright,
                      partial_fmt="2;38;5;%d" % dark,
                      empty_fmt="2;38;5;238")

    def _make_row(self, values):
        out = ""
        cols = 0
        for i, (head, type, width) in enumerate(self.columns):
            if type == "string":
                cell_s = "%*s" % (width, values[i])
                cell_w = len(values[i])
            elif type == "gauge":
                percent, color = values[i]
                cell_s = self._make_gauge(abs(width), percent, color)
                cell_w = abs(width)

            out += " " + cell_s
            if cell_w > abs(width):
                out += "\n" + " " * (cols + 1 + abs(width))
            cols += 1 + cell_w
        return out

    def print(self):
        sep = self._make_separator() + "\n"
        out = self._make_header() + "\n"
        for row in self.rows:
            if row is None:
                out += sep
            else:
                out += self._make_row(row) + "\n"
        print(out, end="")

def get_user_name():
    return os.environ.get("LOGNAME", "root")

def get_home_dir():
    return os.path.expanduser("~")

def get_media_dir():
    return "/run/media/%s" % get_user_name()

def _get_cmd_output(*argv):
    Core.trace("calling command: %r", argv)
    ret = subprocess.run(argv,
                         check=True,
                         stdout=subprocess.PIPE)
    return ret.stdout.decode().strip()

@cache
def _get_zfs_dataset_info(path, fields):
    out = _get_cmd_output("zfs", "list", "-H", "-p", "-o", fields, path)
    return out.split()

@cache
def get_dir_fsid(path):
    return _get_cmd_output("stat", "-f", "-c", "%i", path)

@cache
def get_fs_devno(mtpt):
    return _get_cmd_output("mountpoint", "-d", mtpt)

def fs_kind_from_path(path, fstype):
    if not path.endswith("/"):
        path += "/"
    if path == "/":
        home = get_home_dir()
        path_fsid = get_dir_fsid(path)
        home_fsid = get_dir_fsid(home)
        Core.debug("comparing fsid of %r vs %r", path, home)
        if (path_fsid == home_fsid) and (path_fsid != "0"):
            return Kind.HOME
        else:
            return Kind.ROOT

    fallback = Kind.NONE
    for prefix, kind in fs_kind_prefixes:
        if prefix == "/":
            fallback = kind
        elif prefix.endswith("/"):
            if path.startswith(prefix):
                return kind
        else:
            if path == prefix + "/":
                return kind
        # XXX: use a separate flag

    if fstype in fs_kind_fstypes:
        return fs_kind_fstypes[fstype]

    return fallback

def _path_shorten(path):
    path = path.rstrip("/") + "/"
    home = get_home_dir().rstrip("/") + "/"
    media = get_media_dir().rstrip("/") + "/"
    if path == home:
        return "~"
    elif path == "/":
        return "/ (rootfs)"
    elif path.startswith(home):
        return "~/" + path[len(home):-1]
    elif path.startswith(media):
        return path[len(media):-1]
    else:
        return path[:-1]

def path_shorten(path):
    name = _path_shorten(path)
    aliases = {}
    if name.startswith("vol"):
        name = name.split("_")[0]
    if name in aliases:
        name = aliases[name]
    return name

def fmt_percent(n, digits):
    if round(n, digits) < 100:
        return "%*.*f%%" % (digits+2, digits, n)
    else:
        return "%*.*f%%" % (digits+1, digits-1, n)

def fmt_ansi(text, fmt):
    return "\033[%sm%s\033[m" % (fmt, text) if fmt else text

def gauge3(width, level,
           full_char="█", partial_char="▌", empty_char=" ",
           full_fmt="", partial_fmt="", empty_fmt=""):

    cells     = width * level / 100
    n_full    = int(cells)
    n_partial = int(round(cells % 1))
    n_empty   = int(width - n_full - n_partial)

    return fmt_ansi(full_char    * n_full,    full_fmt) \
         + fmt_ansi(partial_char * n_partial, partial_fmt) \
         + fmt_ansi(empty_char   * n_empty,   empty_fmt)

def darken(color, n):
    def rgb_split(color):
        r = (color - 16) // 6 // 6 % 6
        g = (color - 16) // 6 % 6
        b = (color - 16) % 6
        return (r, g, b)

    def rgb_merge(r, g, b):
        return (r * 6 * 6) + (g * 6) + b + 16

    if 0 <= color <= 7:
        return color
    elif 8 <= color <= 15:
        return color - 8 if n > 0 else color
    elif 16 <= color <= 232:
        r, g, b = rgb_split(color)
        r = max(r - n, 0)
        g = max(g - n, 0)
        b = max(b - n, 0)
        return rgb_merge(r, g, b)
    elif 232 <= color <= 255:
        return max(color - 3*n, 232)

def threshold(total, val):
    if total > 100e9:
        levels = levels_big
    else:
        levels = levels_small

    for tmin, tval in levels:
        if val >= tmin:
            return tval
    return tval

def dump_mtpt(part, data, with_quota=False):
    Core.debug("got %r", part)
    Core.debug("  - %r", data)

    total_bytes     = data["total"]
    free_bytes      = data["free"]
    avail_bytes     = data["avail"]

    used_bytes      = total_bytes - free_bytes
    quota_bytes     = used_bytes + avail_bytes

    quota_part      = used_bytes / quota_bytes * 100
    disk_part       = used_bytes / total_bytes * 100

    if part.cached:
        disk_color      = "gray"
        quota_color     = "gray"
    else:
        disk_color      = threshold(total_bytes, -1)
        quota_color     = threshold(quota_bytes, 100 - quota_part)

    Core.debug(" - total bytes=%r, used=%.2f%%, free=%.2f%%, color=%r",
               total_bytes, disk_part, 100-disk_part, disk_color)
    Core.debug(" - quota bytes=%r, used=%.2f%%, free=%.2f%%, color=%r",
               quota_bytes, quota_part, 100-quota_part, quota_color)

    disk_color      = level_colors[disk_color][0]
    quota_color     = level_colors[quota_color][0]

    if with_quota:
        row = (
            path_shorten(part.mtpt),

            fmt_size_short(total_bytes, si=args.si),
            fmt_size_short(free_bytes, si=args.si),
            fmt_size_short(avail_bytes, si=args.si),

            (quota_part, quota_color),
            fmt_percent(quota_part, 1),

            (disk_part, darken(disk_color, 1)),
            fmt_percent(disk_part, 0),
        )
    else:
        row = (
            path_shorten(part.mtpt),

            fmt_size_short(total_bytes, si=args.si),
            fmt_size_short(used_bytes, si=args.si),
            fmt_size_short(avail_bytes, si=args.si),

            (quota_part, quota_color),
            fmt_percent(disk_part, 1),
        )

    return row

parser = argparse.ArgumentParser()
parser.add_argument("-a", "--all",
                    action="store_true",
                    help="show unimportant filesystems")
parser.add_argument("-c", "--cached",
                    action="store_true",
                    help="include nonpresent cached filesystems")
parser.add_argument("--si",
                    action="store_true",
                    help="use SI decimal units, not IEC binary units")
parser.add_argument("-Q", "--quota",
                    action="store_true",
                    help="display quota information")
parser.add_argument("path", nargs="*")
args = parser.parse_args()

t = Table()

if args.quota:
    t.columns = [
        ("PATH",        "string", -16),
        ("TOTAL",       "string", 7),
        ("FREE",        "string", 7),
        ("AVAIL",       "string", 7),
        ("QUOTA USAGE", "gauge", -20),
        ("",            "string", 5),
        ("DISK USAGE",  "gauge", -10),
        ("",            "string", 3),
    ]
else:
    t.columns = [
        ("PATH",    "string", -20),
        ("TOTAL",   "string", 7),
        ("USED",    "string", 7),
        ("AVAIL",   "string", 7),
        ("",        "gauge", -30),
        ("",        "string", 5),
    ]

this_group = 0
seen_mtpts = set()
seen_devnos = set()

with Enumerator() as en:
    if args.path:
        partitions = en.enum_partitions_exact(args.path)
    else:
        partitions = en.enum_partitions_system(show_cached=args.cached,
                                               show_all=args.all)
    partitions = [*partitions]
    partitions.sort(key=lambda x: x.mtpt)
    partitions.sort(key=lambda x: fs_kind_priorities[x.kind], reverse=True)
    for part in partitions:
        if part.mtpt in seen_mtpts:
            continue
        if part.devno in seen_devnos:
            continue
        data = en.get_partition_data(part)
        if not data:
            continue
        group = fs_kind_groups.get(part.kind)
        Core.debug("got %r (priority %r, group %r)", part.mtpt,
                   fs_kind_priorities.get(part.kind, 0), group)
        row = dump_mtpt(part, data, args.quota)
        if row:
            if this_group and group != this_group:
                t.rows.append(None)
            t.rows.append(row)
            this_group = group
            seen_mtpts.add(part.mtpt)
            seen_devnos.add(part.devno)

t.print()
