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

from __future__ import annotations

import os
import re
import subprocess
from functools import cache
from pathlib import Path
from subprocess import PIPE, Popen, check_output
from sys import exit
from typing import Optional
from uuid import UUID

import config

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()


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):
        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(
                remote_dataset, local_dataset
            )
            end_snapshot = self._get_last_snapshot_with_backup_tag(local_dataset)
            sender = ZFS.send_incremental(start_snapshot, end_snapshot)
            ZFS.receive(sender, self._external_pool.name)

    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 = Manager._get_regex_matching_backup_tags()
        return [
            snapshot for snapshot in dataset.snapshots if snapshot.matches_regex(regex)
        ]

    @staticmethod
    def _get_regex_matching_backup_tags():
        return "@" + ".*".join([config.snapshot_tag, config.snapshot_interval])

    @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, ds1: Dataset, ds2: Dataset) -> Snapshot:
        common_snapshots = [
            local_snap
            for local_snap in list(cls._find_snapshots_with_backup_tag(ds1))
            for remote_snap in list(cls._find_snapshots_with_backup_tag(ds2))
            if local_snap.snapname == remote_snap.snapname
        ]
        common_snapshots.sort()
        return common_snapshots[-1]


class ZFSPath:
    def __init__(self, elements: list[str], snapname: Optional[str] = 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 a valid ZFSPath")


class ZFS:
    @staticmethod
    def get_datasets(name: str) -> list[Dataset]:
        datasets_str = get_output_of_command(
            [
                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
        run_command([config.ZPOOL, "import", "-N", "-d", directory, name])

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

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

    @staticmethod
    def destroy(full_name, *, recursive=False):
        if config.local_pool_to_backup in full_name:
            msg = "ALERT! Tried to delete in local pool!"
            raise Exception(msg)
        cmdline = [config.ZFS, "destroy"]
        if recursive:
            cmdline.append("-r")
        cmdline.append(full_name)
        run_command(cmdline)

    @staticmethod
    def send_incremental(old_snap: Snapshot, new_snap: Snapshot) -> Popen[str]:
        if old_snap.dataset != new_snap.dataset:
            raise ValueError(
                "Cannot send incremental snapshots if start and end snapshot are not based on the same dataset"
            )
        sender = Popen(
            [config.ZFS, "send", "-R", "-I", old_snap, new_snap], stdout=PIPE
        )
        return sender

    @staticmethod
    def receive(sender: Popen[str], pool_name):
        # -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.
        check_output(
            [config.ZFS, "receive", "-d", "-F", "-u", pool_name],
            stdin=sender.stdout,
        )

    @staticmethod
    def scrub(pool: str, *, wait_for_finish=True) -> None:
        run_command([config.ZPOOL, "scrub", pool])
        if wait_for_finish:
            run_command([config.ZPOOL, "wait", "-t", "scrub", pool])


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(f"Already decrypted {self._name}")
        self._verify_mapper_entry_not_in_use()
        self._decrypt_with_cryptsetup()
        self._decrypted = True

    def _decrypt_with_cryptsetup(self):
        run_command([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(f"mapper entry {self._mapper_entry} already exists")
                exit(1)

    @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(f"Not decrypted {self._name}")
        get_output_of_command([config.CRYPTSETUP, "close", self._name])
        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

    # def datasets_to_backup(self, do_not_backup: list[str]) -> list[Dataset]:
    # return [d for d in self.datasets if d.qualified_name not in do_not_backup]


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(f"Already imported {self.name}")
        self._disk.decrypt()
        self._import_from_directory(MAPPER_PATH)
        self._imported = True

    def export(self) -> None:
        if not self._imported:
            raise ValueError(f"Not imported {self.name}")
        self._export_pool()
        self._disk.encrypt()
        self._imported = False

    def _import_from_directory(self, directory: Path) -> None:
        ZFS.import_pool_from_directory(self._name, directory)

    def _export_pool(self) -> None:
        ZFS.export_pool(self._name)

    def scrub(self) -> None:
        ZFS.scrub(self._name)

    @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("Could not find a backup drive")
            exit(1)

    # def datasets_to_backup(self, do_not_backup: list[str]) -> list[Dataset]:
    # def replace_pool_name(old_name: str):
    # "/".join([self._name, old_name.split("/", maxsplit=1)[1]])
    #
    # do_not_backup_correct_name = [replace_pool_name(item) for item in do_not_backup]
    # return [d for d in self.datasets if d.qualified_name not in do_not_backup_correct_name]

    def clean_old_snapshots(self):
        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 = ".*".join(["", config.snapshot_tag, interval, ""])
        snapshots = [
            snapshot for snapshot in self.snapshots if snapshot.matches_regex(regex)
        ]
        if len(snapshots) > max_count:
            snapshots.sort()
            for snap in snapshots[:-max_count]:
                snap.destroy()


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

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

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

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

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

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


def run_command(cmdline: list[str]):
    subprocess.call(cmdline)


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


def verify_running_as_root():
    if os.getuid() > 0:
        print("Run as root.")
        exit(1)


if __name__ == "__main__":
    main()