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

from __future__ import annotations

from functools import cache
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 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")

"""
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-decrypt
# TODO --skip-import
# TODO --skip-pre-backup-scrub
# TODO --skip-post-backup-scrub
# TODO print time estimation while scrub is in progress


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()
            self._backup_all_datasets()
            self._external_pool.clean_old_snapshots()
            self._external_pool.scrub()
        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.snapname == remote_snap.snapname
        ]
        common_snapshots.sort()
        return common_snapshots[-1]


def _get_regex_matching_snapshots_with_tags(tags: list[str]):
    return "@" + ".*".join(tags)


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

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

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

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

    @property
    def snapname(self) -> str:
        return self._snapname

    @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.snapname)

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

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


class ZFS:
    @staticmethod
    def get_datasets(name: str) -> list[Dataset]:
        datasets_str = Command.get_output(
            [
                config.ZFS,
                "list",
                "-H",
                "-d",
                "1",
                "-o",
                "name",
                "-t",
                "filesystem",
                name,
            ]
        )
        # 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]

    @staticmethod
    def import_pool_from_directory(name: str, directory: Path) -> None:
        # -N: no mount
        # -d: directory to search the pool in
        Command.run([config.ZPOOL, "import", "-N", "-d", directory, name])

    @staticmethod
    def export_pool(name: str) -> None:
        Command.run([config.ZPOOL, "export", name])

    @staticmethod
    def get_snapshots(name: str) -> list[Snapshot]:
        return [
            Snapshot(ZFSPath.from_string(name))
            for name in Command.get_output(
                [
                    config.ZFS,
                    "list",
                    "-r",
                    "-t",
                    "snapshot",
                    "-o",
                    "name",
                    "-H",
                    "-d",
                    "1",
                    name,
                ]
            ).split()
        ]

    @staticmethod
    def destroy_snapshot(snapshot: Snapshot, *, recursive=False):
        if not isinstance(snapshot, Snapshot):
            raise TypeError("Can only destroy snapshots.")
        if config.local_pool_to_backup == snapshot.pool:
            raise Exception(DELETE_IN_LOCAL_POOL)
        cmdline = [config.ZFS, "destroy"]
        if recursive:
            cmdline.append("-r")
        cmdline.append(str(snapshot))
        Command.run(cmdline)

    @staticmethod
    def send_incremental(old_snap: Snapshot, new_snap: Snapshot) -> Popen[str]:
        if old_snap.dataset != new_snap.dataset:
            raise ValueError(SAME_DATASET)
        if old_snap == new_snap:
            raise ValueError("Cannot send stream, snapshots are identical.")
        print(f"  {old_snap} -> {new_snap}")
        stream = Popen(
            [config.ZFS, "send", "-R", "-I", str(old_snap), str(new_snap)], stdout=PIPE
        )
        return stream

    @staticmethod
    def receive(stream: Popen[str], pool: Pool):
        # -d  Discard the first element of the sent snapshot's file system name
        # -F  Force a rollback of the file system to the most recent snapshot before performing the receive operation.
        # -u  File system that is associated with the received stream is not mounted.
        if not isinstance(pool, Pool):
            raise TypeError("Can only receive into pools.")
        Command.receive_from_stream(
            [config.ZFS, "receive", "-d", "-F", "-u", pool.name],
            stream,
        )

    @classmethod
    def scrub(cls, pool: Pool, *, wait_for_finish: bool = True) -> None:
        if not isinstance(pool, Pool):
            raise TypeError("Can only scrub pools.")
        if not cls._scrub_in_progress(pool):
            Command.run([config.ZPOOL, "scrub", pool.name])
        if wait_for_finish:
            Command.run([config.ZPOOL, "wait", "-t", "scrub", pool.name])
        if not cls._healthy(pool):
            raise IOError(f"Pool {pool.name} is not healthy.")

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

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

    @staticmethod
    def _pool_status_output(pool):
        return Command.get_output([config.ZFS, "status", pool.name])


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

    def decrypt(self) -> None:
        if self._decrypted:
            raise ValueError(ALREADY_DECRYPTED, self._name)
        self._verify_mapper_entry_not_in_use()
        self._decrypt_with_cryptsetup()
        self._decrypted = True

    def _decrypt_with_cryptsetup(self):
        Command.run([config.CRYPTSETUP, "open", self._path, self._mapper_entry])

    def _verify_mapper_entry_not_in_use(self):
        for item in MAPPER_PATH.iterdir():
            if self._mapper_entry == item.name:
                print(MAPPER_ENTRY_ALREADY_EXISTS, self._mapper_entry)
                exit(EXIT_ERROR)

    @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)
        Command.run([config.CRYPTSETUP, "close", self._mapper_entry])
        self._decrypted = False


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

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

    @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()
        ZFS.import_pool_from_directory(self._name, 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:
        # TODO cancel scrub if necessary
        ZFS.export_pool(self._name)

    def scrub(self) -> None:
        print(f"Scrubbing {self._name}")
        ZFS.scrub(self)

    @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.qualified_name)

    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.snapname < other.snapname

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

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

    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 Command:
    @staticmethod
    def run(cmdline: list[str]) -> None:
        call(cmdline)

    @staticmethod
    def get_output(cmdline: str | list[str]) -> str:
        if isinstance(cmdline, str):
            cmdline = [cmdline]
        return check_output(cmdline).strip().decode()

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


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


if __name__ == "__main__":
    main()