#!/usr/bin/python
# -*- coding: utf-8 -*-

from __future__ import annotations

from datetime import datetime
from os import getuid
from pathlib import Path
from re import search
from subprocess import PIPE, Popen, call, check_output
from sys import exit
from typing import Any
from uuid import UUID

import config

EXIT_ERROR = 1

ALREADY_DECRYPTED = "Cannot decrypt, already decrypted"
ALREADY_IMPORTED = "Cannot import, already imported"
CANNOT_FIND_BACKUP_DRIVE = "Could not find a backup drive"
DELETE_IN_LOCAL_POOL = "ALERT! Tried to delete in local pool!"
MAPPER_ENTRY_ALREADY_EXISTS = "mapper entry already exists"
NOT_DECRYPTED = "Cannot encrypt, not decrypted"
NOT_IMPORTED = "Cannot export, not imported"
NOT_VALID_ZFS_PATH = "Not a valid ZFSPath"
RUN_AS_ROOT = "Run as root."
SAME_DATASET = "Cannot send incremental snapshots if start and end snapshot are not based on the same dataset"

DISK_BY_UUID = Path("/dev/disk/by-uuid")
MAPPER_PATH = Path("/dev/mapper")

# noinspection SpellCheckingInspection
"""
Erstes Backup muss für alle notwendigen Datasets durchgeführt werden mit
# zfs send -Rv rpool/ROOT@zfs-auto-snap_daily-2018-06-14-1648 | zfs receive -dvF backup1
"""


def main():
    verify_running_as_root()
    local_pool = Pool(config.local_pool_to_backup)
    external_pool = ExternalPool.find_from_dict(config.backup_pools)
    manager = Manager(local_pool, external_pool)
    manager.backup()


# TODO tests
# TODO docstrings
# TODO first backup
# TODO --skip-import
# TODO --skip-post-backup-scrub
# TODO print time estimation while scrub is in progress
# TODO logging


class Manager:
    def __init__(self, local_pool: Pool, external_pool: ExternalPool):
        self._local_pool = local_pool
        self._external_pool = external_pool

    def backup(self):
        self._external_pool.import_()
        try:
            self._external_pool.scrub(skip_if_recently_scrubbed=True)
            self._backup_all_datasets()
            self._external_pool.clean_old_snapshots()
            self._external_pool.scrub(skip_if_recently_scrubbed=False)
        finally:
            self._external_pool.export()

    def _backup_all_datasets(self):
        print("Backing up datasets")
        for local_dataset in self._datasets_to_backup:
            remote_dataset = self._search_matching_dataset_in_remote_pool(
                local_dataset, self._external_pool
            )
            start_snapshot = self._get_last_common_snapshot(
                local_dataset, remote_dataset
            )
            end_snapshot = self._get_last_snapshot_with_backup_tag(local_dataset)
            stream = ZFS.send_incremental(start_snapshot, end_snapshot)
            ZFS.receive(stream, self._external_pool)

    def _get_last_snapshot_with_backup_tag(self, local_dataset):
        snapshots = self._find_snapshots_with_backup_tag(local_dataset)
        snapshots.sort()
        end_snapshot = snapshots[-1]
        return end_snapshot

    @property
    def _datasets_to_backup(self) -> list[Dataset]:
        return [
            dataset
            for dataset in self._local_pool.datasets
            if dataset.qualified_name not in config.do_not_backup
        ]

    @staticmethod
    def _find_snapshots_with_backup_tag(dataset: Dataset) -> list[Snapshot]:
        regex = _get_regex_matching_snapshots_with_tags(
            [config.snapshot_tag, config.snapshot_interval]
        )
        return [
            snapshot for snapshot in dataset.snapshots if snapshot.matches_regex(regex)
        ]

    @staticmethod
    def _search_matching_dataset_in_remote_pool(dataset: Dataset, pool: Pool):
        dataset_to_search = dataset.replace_pool(pool.name)
        for dataset in pool.datasets:
            if dataset == dataset_to_search:
                return dataset

    @classmethod
    def _get_last_common_snapshot(
        cls, local_dataset: Dataset, remote_dataset: Dataset
    ) -> Snapshot:
        common_snapshots = [
            local_snap
            for local_snap in list(cls._find_snapshots_with_backup_tag(local_dataset))
            for remote_snap in list(cls._find_snapshots_with_backup_tag(remote_dataset))
            if local_snap.snapshot_name == remote_snap.snapshot_name
        ]
        common_snapshots.sort()
        return common_snapshots[-1]


# TODO this is accessed from both Manager and Dataset. Refactor?
def _get_regex_matching_snapshots_with_tags(tags: list[str]):
    return "@" + ".*".join(tags)


class ZFSPath:
    def __init__(self, elements: list[str], snapshot_name: str | None = None):
        self._elements = elements
        self._snapshot_name = snapshot_name
        self._is_snapshot = False
        if snapshot_name:
            self._is_snapshot = True
        self._sanity_check()

    def __repr__(self):
        path = "/".join(self._elements)
        if self._is_snapshot:
            return "@".join([path, self._snapshot_name])
        return path

    def __eq__(self, other: ZFSPath):
        return (
            self._elements == other._elements
            and self._is_snapshot == other._is_snapshot
            and self._snapshot_name == other._snapshot_name
        )

    @classmethod
    def from_string(cls, string: str) -> ZFSPath:
        snapshot_name = None
        if "@" in string:
            string, snapshot_name = string.split("@")
        elements: list[str] = string.split("/")
        return ZFSPath(elements, snapshot_name)

    @property
    def snapshot_name(self) -> str:
        return self._snapshot_name

    @property
    def is_snapshot(self) -> bool:
        return self._is_snapshot

    @property
    def pool(self) -> str:
        return self._elements[0]

    @property
    def dataset(self) -> ZFSPath:
        return ZFSPath(self._elements)

    def replace_pool(self, pool: str) -> ZFSPath:
        return ZFSPath([pool] + self._elements[1:].copy(), self.snapshot_name)

    @property
    def name_without_pool(self) -> str:
        name = "/".join(self._elements[1:])
        if self.is_snapshot:
            name = "@".join([name, self.snapshot_name])
        return name

    def _sanity_check(self):
        for item in self._elements:
            if len(item) == 0:
                raise ValueError(NOT_VALID_ZFS_PATH)


def _assert_type(instance: Any, object_type: Any, message: str) -> None:
    if not isinstance(instance, object_type):
        raise TypeError(message)


class ZPOOL:
    _ZPOOL = Path("/usr/bin/zpool")

    class _Subcommands:
        # -N: no mount
        # -d: directory to search the pool in
        import_ = ["import", "-N", "-d"]
        export = ["export"]
        scrub = ["scrub"]
        status = ["status"]
        wait = ["wait", "-t", "scrub"]

    @classmethod
    def import_from_directory(cls, pool: Pool, directory: Path) -> None:
        _assert_type(pool, Pool, "Can only import Pool objects")
        CommandRunner.run(
            [str(cls._ZPOOL), *cls._Subcommands.import_, directory, pool.name]
        )

    @classmethod
    def export(cls, pool: Pool) -> None:
        _assert_type(pool, Pool, "Can only export Pool objects")
        CommandRunner.run([str(cls._ZPOOL), *cls._Subcommands.export, pool.name])

    @classmethod
    def scrub(
        cls,
        pool: Pool,
        *,
        skip_if_recently_scrubbed: bool,
        wait_for_finish: bool = True,
    ) -> None:
        _assert_type(pool, Pool, "Can only scrub pools.")
        cls._scrub_if_necessary(pool, skip_if_recently_scrubbed)
        if wait_for_finish:
            CommandRunner.run([str(cls._ZPOOL), *cls._Subcommands.wait, pool.name])
        if not cls._healthy(pool):
            raise IOError(f"Pool {pool.name} is not healthy.")

    @classmethod
    def _scrub_if_necessary(cls, pool: Pool, skip_if_recently_scrubbed: bool) -> None:
        def do_scrub():
            CommandRunner.run([str(cls._ZPOOL), *cls._Subcommands.scrub, pool.name])

        if cls._scrub_in_progress(pool):
            return
        if not skip_if_recently_scrubbed:
            do_scrub()
        else:
            if not cls._recently_scrubbed(pool):
                do_scrub()

    @classmethod
    def _recently_scrubbed(cls, pool: Pool) -> bool:
        return datetime.now() - cls._last_scrub(pool) < config.recent_scrub_timedelta

    @classmethod
    def _scrub_in_progress(cls, pool: Pool) -> bool:
        return "scrub in progress" in cls._get_pool_status_output(pool)

    @classmethod
    def _healthy(cls, pool: Pool) -> bool:
        return "ONLINE" in cls._get_pool_status_output(pool)

    @classmethod
    def _get_pool_status_output(cls, pool: Pool) -> str:
        return CommandRunner.get_output([ZFS, *cls._Subcommands.status, pool.name])

    @classmethod
    def _last_scrub(cls, pool: Pool) -> datetime:
        return Time.get_last_scrub_from_status_output(cls._get_pool_status_output(pool))


class ZFS:
    """Interface of `zfs` commands"""

    class _Subcommands:
        class List:
            # -H   Scripting mode, omit headers
            # -d1  maximum of 1
            # -o   name output column
            # -t   object type
            # -r   recursive
            _base = ["list", "-H", "-d1", "-o", "name"]
            dataset = _base + ["-t", "filesystem"]
            snapshot = _base + ["-t", "snapshot", "-r"]

        destroy = ["destroy"]
        # -R  Replicate filesystem
        # -I  send all intermediary snapshots
        send = ["send", "-R", "-I"]
        # -d  Discard the first element of the "send" snapshot's file system name
        # -F  Force a rollback of the file system to the most recent snapshot before performing the "receive".
        # -u  File system that is associated with the received stream is not mounted.
        receive = ["receive", "-d", "-F", "-u"]

    _ZFS = Path("/usr/bin/zfs")

    @classmethod
    def get_first_level_datasets(cls, pool: Pool) -> list[Dataset]:
        _assert_type(pool, Pool, "Can only get datasets from Pool objects")
        command = [str(cls._ZFS), *cls._Subcommands.List.dataset, pool.name]
        datasets_str = CommandRunner.get_output(command)
        # we omit the first return value, which is the pool, not the dataset
        datasets_list = datasets_str.split("\n")[1:]
        return [Dataset(ZFSPath.from_string(d)) for d in datasets_list]

    @classmethod
    def get_snapshots(cls, dataset: Dataset) -> list[Snapshot]:
        _assert_type(dataset, Dataset, "Can only get snapshots from Dataset object")
        command = [
            str(cls._ZFS),
            *cls._Subcommands.List.snapshot,
            dataset.qualified_name,
        ]
        snapshot_names: list[str] = CommandRunner.get_output(command).split()
        return [Snapshot(ZFSPath.from_string(name)) for name in snapshot_names]

    @classmethod
    def destroy_snapshot(cls, snapshot: Snapshot):
        _assert_type(snapshot, Snapshot, "Can only destroy snapshots.")
        if config.local_pool_to_backup == snapshot.pool:
            raise Exception(DELETE_IN_LOCAL_POOL)
        command = [str(cls._ZFS), *cls._Subcommands.destroy, str(snapshot)]
        CommandRunner.run(command)

    @classmethod
    def send_incremental(
        cls, old_snapshot: Snapshot, new_snapshot: Snapshot
    ) -> Popen[str]:
        cls._pre_send_sanity_checks(old_snapshot, new_snapshot)
        print(f"  {old_snapshot} -> {new_snapshot}")
        return CommandRunner.send_stream(
            [
                str(cls._ZFS),
                *cls._Subcommands.send,
                str(old_snapshot),
                str(new_snapshot),
            ]
        )

    @classmethod
    def _pre_send_sanity_checks(cls, old_snapshot: Snapshot, new_snapshot: Snapshot):
        for snapshot in [old_snapshot, new_snapshot]:
            _assert_type(
                snapshot,
                Snapshot,
                f"Cannot send stream, object is not a Snapshot: {snapshot}",
            )
        if old_snapshot.dataset != new_snapshot.dataset:
            raise ValueError(SAME_DATASET)
        if old_snapshot > new_snapshot:  # chronological ordering
            raise ValueError("Old snapshot is newer than new snapshot.")
        if old_snapshot == new_snapshot:
            raise ValueError("Cannot send stream, snapshots are identical.")

    @classmethod
    def receive(cls, stream: Popen[str], pool: Pool):
        _assert_type(pool, Pool, "Can only receive into pools.")
        command = [str(cls._ZFS), *cls._Subcommands.receive, pool.name]
        CommandRunner.receive_from_stream(command, stream)


class Time:
    @staticmethod
    def _get_date_format() -> str:
        """date format as used by zpool status

        %a: day of week, short version
        %b: month, short version
        %e: day of month, space-padded
        %H, %M, %S: hour, minute, second, zero-padded
        %Y: year
        """
        return "%a %b %e %H:%M:%S %Y"

    @classmethod
    def _extract_datetime_segment(cls, line_with_datetime: str) -> str:
        return line_with_datetime.split(" errors on ")[-1]

    @classmethod
    def _get_line_with_datetime_from_status_output(cls, output: str) -> str:
        lines = output.split("\n")
        for line in lines:
            if line.startswith("  scan: scrub repaired"):
                return line.strip()
        else:
            raise ValueError("Could not find the correct line")

    @classmethod
    def get_last_scrub_from_status_output(cls, output: str) -> datetime:
        line_with_datetime = cls._get_line_with_datetime_from_status_output(output)
        date_str = cls._extract_datetime_segment(line_with_datetime)
        return datetime.strptime(date_str, cls._get_date_format())


class Disk:
    def __init__(self, uuid: UUID, name: str):
        self.uuid = uuid
        self._decrypted: bool = False
        self._was_already_decrypted: bool = False
        self._name = name

    def decrypt(self) -> None:
        if self._decrypted:
            raise ValueError(ALREADY_DECRYPTED, self._name)
        self._was_already_decrypted = Cryptsetup.decrypt(self._path, self._mapper_entry)
        self._decrypted = True

    @property
    def _path(self) -> Path:
        return DISK_BY_UUID / str(self.uuid)

    @property
    def _mapper_entry(self) -> str:
        return f"crypt-{self._name}"

    def encrypt(self):
        if not self._decrypted:
            raise ValueError(NOT_DECRYPTED, self._name)
        if self._was_already_decrypted:
            return
        Cryptsetup.encrypt(self._mapper_entry)
        self._decrypted = False


class Cryptsetup:
    class _Subcommands:
        close = "close"
        status = "status"
        open = "open"

    _CRYPTSETUP = Path("/usr/bin/cryptsetup")

    @classmethod
    def encrypt(cls, mapper_entry: str) -> None:
        command = [cls._CRYPTSETUP, cls._Subcommands.close, mapper_entry]
        CommandRunner.run(command)

    @classmethod
    def _status(cls, name: str) -> str:
        command = [cls._CRYPTSETUP, cls._Subcommands.status, name]
        return CommandRunner.get_output(command)

    @classmethod
    def decrypt(cls, path: Path, mapper_entry: str) -> bool:
        if cls._mapper_entry_in_use(mapper_entry):
            if cls._already_decrypted(path, mapper_entry):
                return True
            print(MAPPER_ENTRY_ALREADY_EXISTS)
            exit(EXIT_ERROR)
        command = [cls._CRYPTSETUP, cls._Subcommands.open, str(path), mapper_entry]
        CommandRunner.run(command)
        return False

    @classmethod
    def _mapper_entry_in_use(cls, mapper_entry: str) -> bool:
        return (MAPPER_PATH / mapper_entry).exists()

    @classmethod
    def _get_device_from_status(cls, status: str) -> Path:
        for line in status.split("\n"):
            if line.startswith("  device:"):
                device = Path(line.split(":")[-1].strip())
                break
        else:
            print("panic")
            exit(EXIT_ERROR)
        for item in DISK_BY_UUID.iterdir():
            # noinspection PyUnboundLocalVariable
            if item.readlink() == device:
                return item

    @classmethod
    def _already_decrypted(cls, path: Path, mapper_entry: str) -> bool:
        decrypted_disk_path = cls._get_device_from_status(cls._status(mapper_entry))
        return decrypted_disk_path == path


class Pool:
    def __init__(self, name: str):
        self._name = name
        self._path: ZFSPath = ZFSPath.from_string(name)

    @property
    def datasets(self) -> list[Dataset]:
        return ZFS.get_first_level_datasets(self)

    @property
    def name(self) -> str:
        return self._name


class ExternalPool(Pool):
    def __init__(self, name: str, disk: Disk):
        super().__init__(name)
        self._imported: bool = False
        self._disk = disk

    def import_(self) -> None:
        if self._imported:
            raise ValueError(ALREADY_IMPORTED, self.name)
        self._disk.decrypt()
        ZPOOL.import_from_directory(self, MAPPER_PATH)
        self._imported = True

    def export(self) -> None:
        if not self._imported:
            raise ValueError(NOT_IMPORTED, self.name)
        self._export_pool()
        self._disk.encrypt()
        self._imported = False
        print(f"Exported {self.name}. Disk {self._disk.uuid} can be removed.")

    def _export_pool(self) -> None:
        ZPOOL.export(self)

    def scrub(self, skip_if_recently_scrubbed: bool) -> None:
        print(f"Scrubbing {self._name}")
        ZPOOL.scrub(self, skip_if_recently_scrubbed=skip_if_recently_scrubbed)

    @classmethod
    def find_from_dict(cls, backup_pools: dict[str, UUID]) -> ExternalPool:
        present_uuids = [item.name for item in DISK_BY_UUID.iterdir()]
        for pool_name, disk_uuid in backup_pools.items():
            if str(disk_uuid) in present_uuids:
                disk = Disk(disk_uuid, pool_name)
                print(f"Found disk {pool_name}, {disk_uuid}")
                return cls(pool_name, disk)
        else:
            print(CANNOT_FIND_BACKUP_DRIVE)
            exit(EXIT_ERROR)

    def clean_old_snapshots(self):
        print("Cleaning old snapshots")
        for dataset in self.datasets:
            dataset.clean_old_snapshots()


class Dataset:
    def __init__(self, zfs_path: ZFSPath):
        self._path = zfs_path

    def __repr__(self):
        return self.qualified_name

    def __eq__(self, other: Dataset):
        return self._path == other._path

    @classmethod
    def from_string(cls, qualified_name: str) -> Dataset:
        return cls(ZFSPath.from_string(qualified_name))

    @property
    def qualified_name(self) -> str:
        return str(self._path)

    @property
    def name_without_pool(self) -> str:
        return self._path.name_without_pool

    @property
    def snapshots(self) -> list[Snapshot]:
        """get qualified snapshot names that are direct children of the dataset"""
        return ZFS.get_snapshots(self)

    def replace_pool(self, name: str) -> Dataset:
        return Dataset(self._path.replace_pool(name))

    def clean_old_snapshots(self):
        for interval, number in config.keep_snapshots_per_interval.keys():
            self._clean_old_snapshots_for_interval(interval, number)

    def _clean_old_snapshots_for_interval(self, interval: str, max_count: int):
        regex = _get_regex_matching_snapshots_with_tags([config.snapshot_tag, interval])
        snapshots = [
            snapshot for snapshot in self.snapshots if snapshot.matches_regex(regex)
        ]
        old_snapshots = self._get_n_oldest_snapshots(snapshots, max_count)
        for snapshot in old_snapshots:
            snapshot.destroy()

    @staticmethod
    def _get_n_oldest_snapshots(
        snapshots: list[Snapshot], count: int
    ) -> list[Snapshot]:
        snapshots.sort()
        return snapshots[:count]


class Snapshot:
    def __init__(self, path: ZFSPath):
        self._path = path

    @classmethod
    def from_string(cls, qualified_name: str) -> Snapshot:
        return cls(ZFSPath.from_string(qualified_name))

    def __lt__(self, other: Snapshot) -> bool:
        return self.snapshot_name < other.snapshot_name

    def __repr__(self):
        return str(self._path)

    @property
    def snapshot_name(self) -> str:
        return self._path.snapshot_name

    def matches_regex(self, regex: str) -> bool:
        return bool(search(regex, str(self._path)))

    def destroy(self):
        ZFS.destroy_snapshot(self)

    @property
    def pool(self) -> str:
        return self._path.pool

    @property
    def dataset(self) -> Dataset:
        return Dataset(self._path.dataset)


class CommandRunner:
    @staticmethod
    def _sanity_check(cmdline: list[str]) -> None:
        _assert_type(cmdline, list, "Command line must be a list")
        for item in cmdline:
            _assert_type(item, str, "Only strings allowed in command line")

    @staticmethod
    def _to_list(command: str | list[str]) -> list[str]:
        if isinstance(command, str):
            return [command]
        return command

    @classmethod
    def run(cls, cmdline: str | list[str]) -> None:
        cmdline = cls._to_list(cmdline)
        cls._sanity_check(cmdline)
        call(cmdline)

    @classmethod
    def get_output(cls, cmdline: str | list[str]) -> str:
        cmdline = cls._to_list(cmdline)
        cls._sanity_check(cmdline)
        return check_output(cmdline).strip().decode()

    @classmethod
    def receive_from_stream(cls, cmdline: str | list[str], stream: Popen) -> None:
        cmdline = cls._to_list(cmdline)
        cls._sanity_check(cmdline)
        check_output(cmdline, stdin=stream.stdout)

    @classmethod
    def send_stream(cls, cmdline: str | list[str]) -> Popen:
        cmdline = cls._to_list(cmdline)
        cls._sanity_check(cmdline)
        return Popen(cmdline, stdout=PIPE)


def verify_running_as_root():
    if getuid() > 0:
        print(RUN_AS_ROOT)
        exit(EXIT_ERROR)


if __name__ == "__main__":
    main()