#!/usr/bin/env python
import json
import os
import subprocess
import time
from nullroute.core import Core, Env

def remote_popen(host, cmd):
    cmd = ["ssh", host, cmd] if host else ["sh", "-c", cmd]
    Core.debug("calling %r", cmd)
    return subprocess.Popen(cmd, stdout=subprocess.PIPE)

def remote_fstat(host, paths):
    cmd = "findmnt -J -l -b -o TARGET,FSTYPE,SIZE,AVAIL,USED"
    with remote_popen(host, cmd) as proc:
        out = json.load(proc.stdout)
        out = out["filesystems"]
        out = [x for x in out if x["target"] in paths]
    # ZFS df only reports dataset usage, not total usage
    for fs in out:
        if fs["fstype"] == "zfs":
            cmd = "zfs list -H -p -o used,avail '%s'"
            with remote_popen(host, cmd % fs["target"]) as proc:
                tmp = proc.stdout.read().decode().split()
            used, avail = map(int, tmp)
            fs["used"] = used
            fs["avail"] = avail
            fs["size"] = used + avail
    return out

cache_path = Env.find_cache_file("diskuse.json")
config_path = Env.find_config_file("diskuse-autocache.json")
media_dir = "/run/media/%(USER)s" % os.environ

with open(cache_path, "r") as fh:
    local = json.load(fh)

with open(config_path, "r") as fh:
    all_disks = json.load(fh)

for host, paths in all_disks.items():
    for fs in remote_fstat(host, paths.keys()):
        mtpt = paths[fs["target"]] or fs["target"]
        mtpt = mtpt.format(media=media_dir)
        local[mtpt] = {
            "type":   fs["fstype"],
            "total":  int(fs["size"]),
            "free":   int(fs["size"]) - int(fs["used"]),
            "avail":  int(fs["avail"]),
            "time":   int(time.time()),
            "expire": int(time.time() + 3 * 86400),
        }

with open(cache_path, "w") as fh:
    json.dump(local, fh)
