from __future__ import print_function

import argparse
import functools
import logging
import operator
import os
import subprocess
from collections import OrderedDict

import boto

from z3.config import get_config


quiet = False


def cached(func):
    @functools.wraps(func)
    def wrapper(self, *a, **kwa):
        # ARGUMENTS AREN'T PART OF CACHE KEY
        # that's safe in this case, but might cause surprising bugs if you reuse it as is
        cache_key = func.__name__ + '_cached_value'
        if not hasattr(self, cache_key):
            val = func(self, *a, **kwa)
            setattr(self, cache_key, val)
        return getattr(self, cache_key)
    return wrapper


class IntegrityError(Exception):
    pass


class S3Snapshot(object):
    CYCLE = 'cycle detected'
    MISSING_PARENT = 'missing parent'
    PARENT_BROKEN = 'parent broken'

    def __init__(self, name, metadata, manager):
        self.name = name
        self._metadata = metadata
        self._mgr = manager
        self._reason_broken = None

    def __repr__(self):
        if self.is_full:
            return "<Snapshot {} [full]>".format(self.name)
        else:
            return "<Snapshot {} [{}]>".format(self.name, self.parent_name)

    @property
    def is_full(self):
        return self._metadata.get('is_full') == 'true'

    @property
    def parent(self):
        parent_name = self._metadata.get('parent')
        return self._mgr.get(parent_name)

    @property
    def parent_name(self):
        return self._metadata.get("parent")

    @cached
    def _is_healthy(self, visited=frozenset()):
        if self.is_full:
            return True
        if self in visited:
            self._reason_broken = self.CYCLE
            return False  # we ended up with a cycle, abort
        if self.parent is None:
            self._reason_broken = self.MISSING_PARENT
            return False  # missing parent
        if not self.parent._is_healthy(visited.union([self])):
            if self.parent.reason_broken == self.CYCLE:
                self._reason_broken = self.CYCLE
            else:
                self._reason_broken = self.PARENT_BROKEN
            return False
        return True

    @property
    def is_healthy(self):
        return self._is_healthy()

    @property
    def reason_broken(self):
        if self.is_healthy:
            return
        return self._reason_broken


class S3SnapshotManager(object):
    def __init__(self, bucket, s3_prefix, snapshot_prefix):
        self.bucket = bucket
        self.s3_prefix = s3_prefix.rstrip('/') + '/'  # make sure we always have a trailing /
        self.snapshot_prefix = snapshot_prefix

    @property
    @cached
    def _snapshots(self):
        prefix = os.path.join(self.s3_prefix, self.snapshot_prefix)
        snapshots = {}
        strip_chars = len(self.s3_prefix)
        for key in self.bucket.list(prefix):
            # print key.key, prefix
            key = self.bucket.get_key(key.key)
            name = key.key[strip_chars:]
            snapshots[name] = S3Snapshot(name, metadata=key.metadata, manager=self)
        # from pprint import pprint
        # pprint(snapshots)
        return snapshots

    def list(self):
        return sorted(self._snapshots.values(), key=operator.attrgetter('name'))

    def get(self, name):
        return self._snapshots.get(name)


class ZFSSnapshot(object):
    def __init__(self, name, metadata, parent=None, manager=None):
        self.name = name
        self.parent = parent

    def __repr__(self):
        return "<Snapshot {} [{}]>".format(self.name, self.parent.name if self.parent else '')


class ZFSSnapshotManager(object):
    def __init__(self, fs_name, snapshot_prefix):
        self._fs_name = fs_name
        self._snapshot_prefix = snapshot_prefix
        self._sorted = None

    def _list_snapshots(self):
        # This is overridden in tests
        # see FakeZFSManager
        return subprocess.check_output(
            ['zfs', 'list', '-Ht', 'snap', '-o',
             'name,used,refer,mountpoint,written'])

    def _parse_snapshots(self):
        """Returns all snapshots grouped by filesystem, a dict of OrderedDict's
        The order of snapshots matters when determining parents for incremental send,
        so it's preserved.
        Data is indexed by filesystem then for each filesystem we have an OrderedDict
        of snapshots.
        """
        try:
            snap = self._list_snapshots()
        except OSError as err:
            logging.error("unable to list local snapshots!")
            return {}
        vols = {}
        for line in snap.splitlines():
            name, used, refer, mountpoint, written = line.split('\t')
            vol_name, snap_name = name.split('@', 1)
            snapshots = vols.setdefault(vol_name, OrderedDict())
            snapshots[snap_name] = {
                'name': name,
                'used': used,
                'refer': refer,
                'mountpoint': mountpoint,
                'written': written,
            }
        return vols

    def _build_snapshots(self, fs_name):
        snapshots = OrderedDict()
        # for fs_name, fs_snaps in self._parse_snapshots().iteritems():
        fs_snaps = self._parse_snapshots().get(fs_name, {})
        parent = None
        for snap_name, data in fs_snaps.iteritems():
            if not snap_name.startswith(self._snapshot_prefix):
                continue
            full_name = '{}@{}'.format(fs_name, snap_name)
            zfs_snap = ZFSSnapshot(
                full_name,
                metadata=data,
                parent=parent,
                manager=self,
            )
            snapshots[full_name] = zfs_snap
            parent = zfs_snap
        return snapshots

    @property
    @cached
    def _snapshots(self):
        return self._build_snapshots(self._fs_name)

    def list(self):
        return self._snapshots.values()

    def get_latest(self):
        return self._snapshots.values()[-1]

    def get(self, name):
        return self._snapshots.get(name)


class CommandExecutor(object):
    @staticmethod
    def shell(cmd, dry_run=False, capture=False):
        if dry_run:
            print(cmd)
        else:
            if capture:
                return subprocess.check_output(
                    cmd, shell=True, stderr=subprocess.STDOUT)
            else:
                return subprocess.check_call(
                    cmd, shell=True)

    @property
    @cached
    def has_pv(self):
        return subprocess.call(
            ['which', 'pv'],
            stderr=subprocess.STDOUT, stdout=subprocess.PIPE) == 0

    def pipe(self, cmd1, cmd2, quiet=False, **kwa):
        """Executes commands"""
        if self.has_pv and not quiet:
            return self.shell("{} | pv | {}".format(cmd1, cmd2), **kwa)
        else:
            return self.shell("{} | {}".format(cmd1, cmd2), **kwa)


class PairManager(object):
    def __init__(self, s3_manager, zfs_manager, command_executor=None):
        self.s3_manager = s3_manager
        self.zfs_manager = zfs_manager
        self._cmd = command_executor or CommandExecutor()

    def list(self):
        pairs = []
        seen = set([])
        for z_snap in self.zfs_manager.list():
            seen.add(z_snap.name)
            pairs.append(
                (self.s3_manager.get(z_snap.name), z_snap))
        for s3_snap in self.s3_manager.list():
            if s3_snap.name not in seen:
                pairs.append((s3_snap, None))
        return pairs

    def _snapshot_to_backup(self, snap_name):
        if snap_name is None:
            z_snap = self.zfs_manager.get_latest()
        else:
            z_snap = self.zfs_manager.get(snap_name)
            if z_snap is None:
                raise Exception('Failed to get the snapshot {}'.format(snap_name))
        return z_snap

    @staticmethod
    def _parse_estimated_size(output):
        try:
            size_line = [line for line in output.splitlines() if len(line)][-1]
            _, size = size_line.split()
            return int(size)
        except:
            logging.error("failed to parse output '%s'", output)
            raise

    def backup_full(self, snap_name=None, dry_run=False):
        """Do a full backup of a snapshot. By default latest local snapshot"""
        z_snap = self._snapshot_to_backup(snap_name)
        estimated_size = self._parse_estimated_size(
            self._cmd.shell(
                "zfs send -nvP '{}'".format(z_snap.name),
                capture=True))
        self._cmd.pipe(
            "zfs send '{}'".format(z_snap.name),
            "pput --estimated {} --meta is_full=true {}{}".format(
                estimated_size, self.s3_manager.s3_prefix, z_snap.name),
            dry_run=dry_run,
        )

    def backup_incremental(self, snap_name=None, dry_run=False):
        """Uploads named snapshot or latest, along with any other snapshots
        required for an incremental backup.
        """
        z_snap = self._snapshot_to_backup(snap_name)
        to_upload = []
        current = z_snap
        while True:
            s3_snap = self.s3_manager.get(current.name)
            if s3_snap is not None:
                if not s3_snap.is_healthy:
                    # abort everything if we run in to unhealthy snapshots
                    raise IntegrityError(
                        "Broken snapshot detected {}, reason: '{}'".format(
                            s3_snap.name, s3_snap.reason_broken
                        ))
                break
            to_upload.append(current)
            if current.parent is None:
                break
            current = current.parent
        for z_snap in reversed(to_upload):
            estimated_size = self._parse_estimated_size(
                self._cmd.shell(
                    "zfs send -nvP -i '{}' '{}'".format(
                        z_snap.parent.name, z_snap.name),
                    capture=True))
            self._cmd.pipe(
                "zfs send -i '{}' '{}'".format(
                    z_snap.parent.name, z_snap.name),
                "pput --estimated {} --meta parent={} {}{}".format(
                    estimated_size, z_snap.parent.name, self.s3_manager.s3_prefix, z_snap.name),
                dry_run=dry_run,
            )

    def restore(self, snap_name, dry_run=False):
        current_snap = self.s3_manager.get(snap_name)
        if current_snap is None:
            raise Exception('no such snapshot "{}"'.format(snap_name))
        to_restore = []
        while True:
            z_snap = self.zfs_manager.get(current_snap.name)
            if z_snap is not None:
                break
            if not current_snap.is_healthy:
                raise IntegrityError(
                    "Broken snapshot detected {}, reason: '{}'".format(
                        current_snap.name, current_snap.reason_broken
                    ))
            to_restore.append(current_snap)
            if current_snap.is_full:
                break
            else:
                current_snap = current_snap.parent
        for s3_snap in reversed(to_restore):
            self._cmd.pipe(
                "z3_get {}".format(
                    os.path.join(self.s3_manager.s3_prefix, s3_snap.name)),
                "zfs recv {}".format(s3_snap.name),
                dry_run=dry_run,
            )


def _get_widths(widths, line):
    for index, value in enumerate(line):
        widths[index] = max(widths[index], len("{}".format(value)))
    return widths


def _prepare_line(s3_snap, z_snap):
    if s3_snap is None:
        snap_type = 'missing'
        health = '-'
        name = z_snap.name.split('@', 1)[1]
        parent_name = '-'
        local_state = 'ok'
    else:
        snap_type = 'full' if s3_snap.is_full else 'incremental'
        health = s3_snap.reason_broken or 'ok'
        parent_name = '' if s3_snap.is_full else s3_snap.parent_name
        name = s3_snap.name
        local_state = 'ok' if z_snap is not None else 'missing'
    return (name, parent_name, snap_type, health, local_state)


def list_snapshots(bucket, s3_prefix, filesystem, snapshot_prefix):
    print("backup status for {}@{}* on {}/{}".format(
        filesystem, snapshot_prefix, bucket.name, s3_prefix))
    prefix = "{}@{}".format(filesystem, snapshot_prefix)
    pair_manager = PairManager(
        S3SnapshotManager(bucket, s3_prefix=s3_prefix, snapshot_prefix=prefix),
        ZFSSnapshotManager(fs_name=filesystem, snapshot_prefix=snapshot_prefix))
    header = ("NAME", "PARENT", "TYPE", "HEALTH", "LOCAL STATE")
    widths = [len(col) for col in header]
    listing = []
    for s3_snap, z_snap in pair_manager.list():
        line = _prepare_line(s3_snap, z_snap)
        listing.append(line)
        widths = _get_widths(widths, line)
    fmt = " | ".join("{{:{w}}}".format(w=w) for w in widths)
    print(fmt.format(*header))
    for line in listing:
        print(fmt.format(*line))


def do_backup(bucket, s3_prefix, filesystem, snapshot_prefix, full, snapshot, dry):
    prefix = "{}@{}".format(filesystem, snapshot_prefix)
    s3_mgr = S3SnapshotManager(bucket, s3_prefix=s3_prefix, snapshot_prefix=prefix)
    zfs_mgr = ZFSSnapshotManager(fs_name=filesystem, snapshot_prefix=snapshot_prefix)
    pair_manager = PairManager(s3_mgr, zfs_mgr)
    snap_name = "{}@{}".format(filesystem, snapshot) if snapshot else None
    if full is True:
        pair_manager.backup_full(snap_name=snap_name, dry_run=dry)
    else:
        pair_manager.backup_incremental(snap_name=snap_name, dry_run=dry)


def restore(bucket, s3_prefix, filesystem, snapshot_prefix, snapshot, dry):
    prefix = "{}@{}".format(filesystem, snapshot_prefix)
    s3_mgr = S3SnapshotManager(bucket, s3_prefix=s3_prefix, snapshot_prefix=prefix)
    zfs_mgr = ZFSSnapshotManager(fs_name=filesystem, snapshot_prefix=snapshot_prefix)
    pair_manager = PairManager(s3_mgr, zfs_mgr)
    snap_name = "{}@{}".format(filesystem, snapshot)
    pair_manager.restore(snap_name, dry_run=dry)


def main():
    cfg = get_config()
    parser = argparse.ArgumentParser(
        description='list z3 snapshots',
    )
    parser.add_argument('--s3-prefix',
                        dest='s3_prefix',
                        default=cfg.get('S3_PREFIX', 'z3-backup/'),
                        help='S3 key prefix, defaults to z3-backup/')
    parser.add_argument('--filesystem', '--dataset',
                        dest='filesystem',
                        default=cfg.get('FILESYSTEM'),
                        help='the zfs dataset/filesystem to operate on')
    parser.add_argument('--snapshot-prefix',
                        dest='snapshot_prefix',
                        default=cfg.get('SNAPSHOT_PREFIX', 'zfs-auto-snap:daily'),
                        help='only operate on snapshots that start with this prefix')
    subparsers = parser.add_subparsers(help='sub-command help', dest='subcommand')

    backup_parser = subparsers.add_parser(
        'backup', help='backup local zfs snapshots to an s3 bucket')
    backup_parser.add_argument('--snapshot', dest='snapshot', default=None,
                               help='Snapshot to backup. Defaults to latest.')
    backup_parser.add_argument('--dry', dest='dry', default=False, action='store_true',
                               help='Snapshot to backup. Defaults to latest.')
    incremental_group = backup_parser.add_mutually_exclusive_group()
    incremental_group.add_argument(
        '--full', dest='full', action='store_true', help='Perform full backup')
    incremental_group.add_argument(
        '--incremental', dest='incremental', default=True, action='store_true',
        help='Perform incremental backup; this is the default')

    restore_parser = subparsers.add_parser('restore', help='not implemented')
    restore_parser.add_argument(
        'snapshot', help='Snapshot to backup. Defaults to latest.')
    restore_parser.add_argument('--dry', dest='dry', default=False, action='store_true',
                                help='Snapshot to backup. Defaults to latest.')
    subparsers.add_parser('status', help='show status of current backups')
    args = parser.parse_args()

    bucket = boto.connect_s3(
        cfg['S3_KEY_ID'], cfg['S3_SECRET']).get_bucket(cfg['BUCKET'])
    if args.subcommand == 'status':
        list_snapshots(bucket, s3_prefix=args.s3_prefix, snapshot_prefix=args.snapshot_prefix,
                       filesystem=args.filesystem)
    elif args.subcommand == 'backup':
        do_backup(bucket, s3_prefix=args.s3_prefix, snapshot_prefix=args.snapshot_prefix,
                  filesystem=args.filesystem, full=args.full, snapshot=args.snapshot,
                  dry=args.dry)
    elif args.subcommand == 'restore':
        restore(bucket, s3_prefix=args.s3_prefix, snapshot_prefix=args.snapshot_prefix,
                filesystem=args.filesystem, snapshot=args.snapshot, dry=args.dry)


if __name__ == '__main__':
    main()
