from datetime import datetime, timezone
from dateutil.relativedelta import relativedelta
from pprint import pprint
import subprocess
import sys


def keep_time(now, snapshots, convert, delta, count):
    keep = set()
    if count < 1:
        return keep
    start = convert(now) - delta*(count-1)
    for snapshot in snapshots:
        if snapshot >= start:
            keep.add(snapshot)
            start = convert(snapshot) + delta
        if start > now:
            break
    return keep


def keep_year(now, snapshots, count):
    return keep_time(now, snapshots, lambda d: d.replace(month=1, day=1, hour=0, minute=0, second=0, microsecond=0), relativedelta(years=1), count)


def keep_month(now, snapshots, count):
    return keep_time(now, snapshots, lambda d: d.replace(day=1, hour=0, minute=0, second=0, microsecond=0), relativedelta(months=1), count)


def keep_day(now, snapshots, count):
    return keep_time(now, snapshots, lambda d: d.replace(hour=0, minute=0, second=0, microsecond=0),  relativedelta(days=1), count)


def keep_hour(now, snapshots, count):
    return keep_time(now, snapshots, lambda d: d.replace(minute=0, second=0, microsecond=0), relativedelta(hours=1),
                     count)


def keep_frequent(now, snapshots, count):
    return keep_time(now, snapshots, lambda d: d.replace(minute=d.minute//15*15, second=0, microsecond=0),
                     relativedelta(minutes=15), count)


def snapshots_to_keep(now, snapshots, yearly=0, monthly=0, daily=0, hourly=0, frequently=0):
    snapshots.sort()
    keep = set()
    keep.update(keep_year(now, snapshots, yearly))
    keep.update(keep_month(now, snapshots, monthly))
    keep.update(keep_day(now, snapshots, daily))
    keep.update(keep_hour(now, snapshots, hourly))
    keep.update(keep_frequent(now, snapshots, frequently))
    return sorted(list(keep))


def parse_int(fields, name, fieldname):
    if fields[fieldname] == '-':
        return 0
    try:
        return int(fields[fieldname])
    except ValueError:
        print(f"Warning: Could not parse value {fields[fieldname]} for fs/zvol {name} prop {fieldname}",
              output=sys.stderr)
        return -1


def zfs_list(fieldnames, args):
    proc = subprocess.run(['/usr/sbin/zfs', 'list', '-H', '-o', ','.join(['name'] + fieldnames)] + args,
                          capture_output=True, check=True, encoding='UTF-8')
    pprint(proc.args)
    if not fieldnames:
        result = list()
    else:
        result = dict()
    for line in proc.stdout.splitlines():
        fields = line.split('\t')
        if not fieldnames:
            result.append(fields[0])
        else:
            result[fields[0]] = {n: v for n, v in zip(fieldnames, fields[1:])}
    return result


def read_snapshot_props():
    periods = ['yearly', 'monthly', 'daily', 'hourly', 'frequently']
    fieldnames = ['org.blankertz:snapshot'] + ['org.blankertz:snapshot:' + period for period in periods]
    raw_list = zfs_list(fieldnames, ['-t', 'filesystem,volume', '-r'])
    result = dict()
    for name, fields in raw_list.items():
        result[name] = {'snapshot': True if fields['org.blankertz:snapshot'].lower() == 'true' else False}
        result[name].update({period: parse_int(fields, name, 'org.blankertz:snapshot:' + period) for period in periods})
    return result


def read_snapshots(fs_to_snap):
    raw_list = zfs_list([], ['-t', 'snapshot'] + fs_to_snap) if fs_to_snap else []
    result = dict()
    for name in raw_list:
        fs, snap = name.split('@')
        print(f'{fs}, {snap}')
        if snap.startswith('zfs-smart-snap'):
            result[fs] = result.get(fs, []) + [snap]
    return result


def find_latest_snapshots(existing_snaps):
    result = dict()
    for fs, snaps in existing_snaps.items():
        snaps.sort()
        result[fs] = snaps[-1]
    return result


def written_since_last(fs, snap):
    proc = subprocess.run(['/usr/sbin/zfs', 'list', '-Hp', '-o', f'written@{snap}', fs],
                          check=True, capture_output=True, encoding="UTF-8")
    return int(proc.stdout) > 0


def main():
    now = datetime.now(timezone.utc).replace(second=0, microsecond=0)
    props = read_snapshot_props()
    pprint(props)
    fs_to_snap = [k for k, v in props.items() if v['snapshot']]
    if fs_to_snap:
        existing_snaps = read_snapshots(fs_to_snap)
        pprint(existing_snaps)
        latest_snaps = find_latest_snapshots(existing_snaps)
        fs_to_really_snap = [fs for fs in fs_to_snap if fs not in latest_snaps or
                             written_since_last(fs, latest_snaps[fs])]
        if fs_to_really_snap:
            snapname = 'zfs-smart-snap-' + now.strftime('%Y-%m-%d-%H%M')
            subprocess.run(['/usr/sbin/zfs', 'snapshot'] + [f'{fs}@{snapname}' for fs in fs_to_really_snap], check=True)

    # TODO: Cleanup expired snapshots


if __name__ == "__main__":
    main()