"""Defines the actual operating classes for Tapestry.py. There is essentially
no executable code in this script - it's just used to hold the working classes
for functions.py.

Classes are to be organized by general purpose - see block comments for guidance.

"""

import bz2
import ftplib
import hashlib
import json
import multiprocessing as mp
import os
import pickle
import shutil
import tarfile
from platform import system
from textwrap import wrap

# Define Exceptions


class RecoveryIndexError(Exception):
    """Standard exception to be raised when the recovery index has an issue
    with the file it is provided."""

    def __init__(self, message):
        self.message = message

    def __str__(self):
        return self.message


# Define Process and Task Classes


class ChildProcess(mp.Process):
    """A simple, logicless worker process which iterates against a queue. This
    expects to find either a callable "task" class in the queue or the value
    None, which indicates the process should exit.
    """

    def __init__(self, queue_tasking, queue_complete, working_directory, locks, debug=False):
        """Initialize the class. Must provide the queue to attach to, the
        working directory to use, and a debug flag.

        :param queue_tasking: mp.Queue object used for multiprocess communication.
        :param queue_complete: mp.Queue object used for communicating completions
        :param working_directory: String to the absolute path of the desired
        working directory.
        :param locks: a dictionary of blockname:lock pairs (optional).
        :param debug: Boolean, activates the debugging statement.
        """
        mp.Process.__init__(self)
        self.queue = queue_tasking
        self.debug = debug
        self.ret = queue_complete
        if not os.path.isdir(working_directory):
            os.mkdir(working_directory)
        os.chdir(working_directory)

        global dict_locks
        dict_locks = locks

    def run(self):
        proc_name = self.name
        while True:
            next_task = self.queue.get()
            if next_task is None:
                if self.debug:
                    self.ret.put('%s: Exiting' % proc_name)
                self.queue.task_done()
                break
            try:
                debug_message = next_task()
                self.ret.put(debug_message)
            except TypeError:
                self.ret.put("Something has gone wrong - unexpected item in queue.")
            self.queue.task_done()
        return


class TaskTarBuild(object):
    """A simple task object which instructs the process to add a particular
    file to a particular tarfile, all fully enumerated.
    """

    def __init__(self, tarf, fid, path, bname):
        """
        Create the tarfile or, otherwise, add a file to it.
        :param tarf: which tarfile to use, relative to the working directory
        :param fid: GUID FID as a string.
        :param path: absolute path to the file in need of backup
        :param bname: block name, used as a key to the locks dictionary.
        """
        self.tarf = tarf
        self.a = fid
        self.b = path
        self.block = bname

    def __call__(self):
        global dict_locks

        f_lock = dict_locks[self.block]  # Aquires the lock indicated in the index value from the master
        f_lock.acquire()
        tar = tarfile.open(name=self.tarf, mode="a:")
        tar.add(self.b, arcname=self.a, recursive=False)
        tar.close()
        f_lock.release()

        return "Added %s to tarfile %s" % (self.b, self.tarf)


class TaskTarUnpack(object):
    """A simple object that describes a file to pull from a particular tarfile
    and puts it back where it belongs. Absolute paths required.
    """
    def __init__(self, tar, fid, category_dir, path_end):
        """Initializing this object gives it all the information it needs to
        unpack one particular file back to its original position in the
        filesystem or to the fallback location.

        :param tar: string describing the absolute path of the relevant tarball
        :param fid: The GUID FID that the file has been packed into the tarball as.
        :param category_dir: full path to the top of the category.
        :param path_end: sub-path from the category including the filename
        """
        self.tar = tar
        self.fid = fid
        self.catdir = category_dir
        self.pathend = path_end

    def __call__(self):
        path_end = self.pathend.strip('~/')
        abs_path_out = os.path.join(self.catdir, path_end)
        placement, name_proper = os.path.split(abs_path_out)
        # split the pathend component into the subpath from the category dir, and the original filename
        # print("Attempting to pull %s" % self.fid)
        unpacking = True
        while unpacking:
            try:  # Issue 13: need to catch the raise from tf.extract() when workers collide
                with tarfile.open(self.tar, "r:*") as tf:
                    # print("Now opened: %s" % self.tar)  # Debugging statement, uncomment to use.
                    tf.extract(self.fid, path=placement)  # the file is now located where it needs to be.
                    # print("Extracted %s" % self.fid)  # Debugging statement, uncomment to use.
                    placed = os.path.join(placement, self.fid)
                    os.rename(placed, abs_path_out)  # and now it's named correctly.
                    unpacking = False
            except FileExistsError:
                unpacking = True
        return "Restored %s to %s" % (self.fid, abs_path_out)


class TaskCompress(object):
    """A simple task that points exactly to a tarfile to compress, and then
    compresses it to a specified level
    """

    def __init__(self, t, lvl):
        """Defines the absolute path to a tarfile and the level (1-9) of
        compression to apply.

        :param t: Absolute path to a tarfile
        :param lvl: Integer between 1 and 9 which determines the compression level
        """
        self.tarf = t
        self.level = lvl

    def __call__(self):
        with open(self.tarf, "rb") as b:
            bz2d = self.tarf+".bz2"
            bz2f = bz2.BZ2File(bz2d, "wb", compresslevel=self.level)
            shutil.copyfileobj(b, bz2f)

        return "Compressed %s to level %s" % (self.tarf, self.level)


class TaskDecompress(object):
    """A simple task that points exactly to a tarfile to decompress. The task
    can determine for itself if the file actually needs decompressing at runtime.
    """

    def __init__(self, t):
        """Defines the absolute path to a file to check for compression.

        :param t: Absolute path to a tarfile
        """
        self.tarf = t

    def __call__(self):
        with open(self.tarf, "rb") as file:
            signature = file.read(3)
            # print(signature)
            if signature.startswith(b"BZh"):  # Microheader that indicates a BZ2 file.
                shutil.copy(self.tarf, (self.tarf + ".temp"))
                with bz2.BZ2File((self.tarf + ".temp"), "rb") as compressed:
                    with open(self.tarf, "wb") as uncompressed:
                        shutil.copyfileobj(compressed, uncompressed)
                os.remove(self.tarf+".temp")  # No sense in cluttering up the drive.
                return "Decompressed %s" % self.tarf
            else:
                return "File is already decompressed."


class TaskEncrypt(object):
    """This task takes an argument for the fingerprint to use, the file to be
    encrypted, and the directory currently used for output.

    """
    def __init__(self, t, fp, out, gpg):
        """This object expects the following arguments in order to signify a
        tarfile which needs to be encrypted. Calling the task will cause the
        tarfile to be encrypted using the gpg object passed to it at runtime.

        :param t: an absolute path denoting the location of the tarfile.
        :param fp: the fingerprint of the PGP key to use.
        :param out: The absolute path to the output directory.
        :param gpg: An gnupg.GPG object provided to allow an interface with
        the local GPG runtime.
        """
        self.tarf = t
        self.fp = fp
        self.out = out
        self.gpg = gpg

    def __call__(self):
        path, tapped = os.path.split(self.tarf)
        tapped = tapped.strip(".bz2")
        tapped = tapped.replace(".tar", ".tap")
        tgt_output = os.path.join(self.out, tapped)
        with open(self.tarf, "rb") as tgt:
            k = self.gpg.encrypt_file(tgt, self.fp, output=tgt_output, armor=True, always_trust=True)
        if k.ok:
            return "Encryption Success for %s." % self.tarf
        elif not k.ok:
            return "Encryption Failed for %s, status: %s" % (self.tarf, k.status)


class TaskDecrypt(object):
    """This task contains both the information and method to take a Tapestry
    blockfile and decrypt it. This task is naive in that it relies on another
    task to do signature verification.

    """
    def __init__(self, block, working_directory, gpg):
        """This object requires the following arguments in order to fetch,
        open, and decrypt a single block, outputting the contents to a
        working directory, based on the operation of a gnupg.GPG object.

        :param block: Absolute path to the block to be decrypted.
        :param working_directory: Absolute path to the output directory.
        :param gpg:
        """
        self.tap_file = block
        self.tap_name = os.path.split(block)[1]
        self.tap_name += ".decrypted"
        self.absolute_output = os.path.join(working_directory, self.tap_name)
        self.gpg = gpg

    def __call__(self):
            with open(self.tap_file, "rb") as tgt:
                k = self.gpg.decrypt_file(tgt, output=self.absolute_output, always_trust=True)
            if k.ok:
                return "Decryption Success for %s." % self.tap_name
            elif not k.ok:
                return "Decryption Failed for %s, status: %s" % (self.tap_name, k.status)


class TaskSign(object):
    """This task takes an argument for the fingerprint to use, the file to be
    signed, the output directory, and a GPG object.

    """
    def __init__(self, t, fp, out, gpg):
        """This object expects the following arguments in order to sign a
        tapfile which needs to be encrypted. Calling the task will cause the
        tapfile to be signed using the gpg object passed to it at runtime.

        :param t: an absolute path denoting the location of the tapfile.
        :param fp: the fingerprint of the PGP key to use.
        :param out: The absolute path to the output directory.
        :param gpg: An gnupg.GPG object provided to allow an interface with
        the local GPG runtime.
        """
        self.file = t
        self.fp = fp
        self.out = out
        self.gpg = gpg

    def __call__(self):
        with open(self.file, "rb") as p:
            path, tapped = os.path.split(self.file)
            absolute_output = os.path.join(self.out, tapped)
            absolute_output += ".sig"
            # print(absolute_output)
            k = self.gpg.sign_file(p, keyid=self.fp, output=absolute_output, detach=True)
            # print(k)
            # print("\n")
            if k.status == "signature created":
                return "Signing Success for %s." % tapped
            else:
                return "Signing Failed for %s, status: %s" % (tapped, k.status)


class TaskCheckIntegrity(object):
    """A task, to be completed by ChildProcess, which contains the information
    and operations needed to check the integrity of a file in the archive
    against it's known-good composition, based on the MD5 hash.
    """

    def __init__(self, tar_file, fid, kg_hash):
        """Provided with a tarfile, an FID found within it, and a known good
        hash, returns True or False if the file matches, as well as a string
        used in debugging.

        :param tar_file: string denoting absolute path to the tarball
        :param fid: GUID file identifier of the file in question.
        :param hash: the output of md5.hexdigest(), as found (e.g) in the RIFF
        """
        self.tarf = tar_file
        self.fid = fid
        self.hash_good = kg_hash

    def __call__(self):
        hasher = hashlib.sha256()  # Had to change this because switched to sha256 from md5 in functions.py.
        with tarfile.open(self.tarf, "r:*") as tarball:
            file_under_test = tarball.extractfile(self.fid)
            if file_under_test is None:
                return [False, "File %s not found in block\n" % self.fid]
            else:
                hasher.update(file_under_test.read())
        if hasher.hexdigest() == self.hash_good:
            return [True, "File %s has a valid hash." % self.fid]
        else:
            print(hasher.hexdigest())
            print(self.hash_good)
            return [False, "File %s has an invalid hash.\n" % self.fid]


# Define Package Overrides


class FTP_TLS(ftplib.FTP_TLS):  # With thanks to hynekcer
    """Explicit FTPS, with shared TLS session"""
    def ntransfercmd(self, cmd, rest=None):
        conn, size = ftplib.FTP.ntransfercmd(self, cmd, rest)
        if self._prot_p:
            conn = self.context.wrap_socket(conn,
                                            server_hostname=self.host,
                                            session=self.sock.session)  # this is the fix
        return conn, size


# Define Utility Objects

class Block(object):
    """Class representation of a tapestry block object. Exposes a method "put"
    to add a file to the block, and another, "meta" method, for adding the
    full-run metadata file to a batch before proceeding with packing.

    """

    def __init__(self, name, max_size, count, smallest):
        """Initialize the block with its preset values.

        :param name: string, output filename
        :param max_size: int in bytes
        """
        self.name = name  # This string will be used as the final output filename
        self.max_size = max_size
        self.size = 0
        self.file_index = {}
        self.remaining = max_size
        self.files = 0  # A simple integer counter
        self.run_metadata = {}
        self.num_block = count
        self.block_metadata = {}
        self.global_index = {}
        self.smallest = smallest
        self.full = False

    def put(self, file_identifier, file_index_object):
        """Accepts an arbitrary file index identifier and the corresponding
        dictionary of metadata, as from a recovery index object"""
        if file_index_object['fsize'] <= self.remaining:  # the new file will fit
            self.file_index.update({file_identifier: file_index_object})
            self.size += file_index_object['fsize']
            self.files += 1
            self.remaining = self.max_size - self.size
            if self.remaining <= self.smallest:
                self.full = True
            return True
        else:  # This file won't fit and has to be placed somewhere else.
            return False

    def meta(self, sum_blocks, sum_size, sum_files, datestamp, comment_string, full_index, drop_dir):
        """Provided these arguments, populate the runMetadata portion of a RIFF,
        then create the corresponding RIFF file.
        """
        meta_value = {}
        meta_value.update({"sumBlock": sum_blocks})
        meta_value.update({"sizeExtraLarge": sum_size})
        meta_value.update({"countFilesSum": sum_files})
        meta_value.update({"dateRec": datestamp})
        if comment_string is None:
            comment_string = "No Comment"
        meta_value.update({"comment": comment_string})
        self.run_metadata = meta_value

        self.block_metadata = {
            "numBlock": self.num_block, "sizeLarge": self.size, "countFiles": self.files
        }

        self.global_index = full_index

        dict_riff = {"metaBlock": self.block_metadata, "metaRun": self.run_metadata,
                     "index": self.global_index}

        with open(os.path.join(drop_dir, self.name+".riff"), "w") as riff:
            json.dump(dict_riff, riff)

        return os.path.join(drop_dir, (self.name+".riff"))


class RecoveryIndex(object):
    """Special utility class for loading and translating Tapestry recovery
    index files and presenting them back to the script in a universal way. Made
    for both the old Recovery Pickle design as well as the NewRIFF format.
    """

    def __init__(self, index_file):
        """Initialize the Recovery Index by handing it a working file.

        :param index_file: a reader object containing the file in question.
        """
        self.pickle_failed = False
        self.json_failed = False

        try:
            self.pickled_data = pickle.load(index_file)
        except pickle.UnpicklingError:  # In this case we must have a newRiff:
            self.pickle_failed = True

        try:
            self.unpacked_json = json.load(index_file)
        except json.JSONDecodeError:
            self.json_failed = True

        if self.pickle_failed and self.json_failed:
            raise RecoveryIndexError("The recovery index is not a valid file type, or corrupt.")
        elif self.pickle_failed:
            self.mode = "json"
        else:
            self.mode = "pkl"

        if self.mode == "json":
            self.run_metadata = self.unpacked_json["metaRun"]
            self.file_index = self.unpacked_json["index"]
            self.blocks = self.unpacked_json["metaRun"]["sumBlock"]
        elif self.mode == "pkl":
            self.blocks, self.rec_paths, self.rec_sections = self.pickled_data
        else:  # We have entered a cursed state...
            raise RecoveryIndexError("The self.mode variable is an unexpected value. Are you hacking?")

    def find(self, file_key):
        """Expectes a FID value as the argument and will return the
        category and sub-path accordingly.

        :param file_key: A string representing a valid file ID.
        """
        if file_key.lower() in ["recovery-pkl", "recovery-riff"]:
            category = "skip"
            sub_path = "skip"
            return category, sub_path
        elif self.mode == "json":
            try:
                category = self.file_index[file_key]["category"]
                sub_path = self.file_index[file_key]["fpath"]
            except KeyError:
                category = b"404"
                sub_path = b"404"
        elif self.mode == "pkl":
            try:
                category = self.rec_sections[file_key]
                sub_path = self.rec_paths[file_key]
            except KeyError:
                category = b"404"
                sub_path = b"404"
        else:
            raise RecoveryIndexError("The self.mode variable is an unexpected value. Are you hacking?")

        return category, sub_path


class SimpleLogger:
    """A minimal class to provide very low-level logging functionality to an
    arbitrary logfile. Tapestry relies on this being instantiated and attached
    to the namespace object shared among functions so that it can be called.
    """

    def __init__(self, landingdir, name, label, date, cfg):
        """Return a SimpleLogger object with the below properties

        :param landingdir: The path (ideally absolute) to the output file.
        :param name: The name of that logfile
        :param label: The machine label for which this backup is being run.
        :param date: The current date, used in the header.
        :param cfg: The path to the config file used, ideally absolute.
        """

        landingAbs = os.path.join(landingdir, name)
        if not os.path.exists(landingdir):
            os.makedirs(landingdir)
        self.loggerfile = open(landingAbs, "w")  # This will REPLACE the existing logfile with the new one so be careful
        header = "This file is a log of all operations that were part of the run of Tapestry %s, " \
                 "using the configuration file at %s, which is for the machine labelled \"%s\"."
        header = (header % (date, cfg, label))
        for each in wrap(header,80):
            self.loggerfile.write(each+'\n')
        cores = os.cpu_count()
        self.loggerfile.write("Cores Available: %s \n" % cores)
        self.loggerfile.write("Operating System: %s \n" % system())
        RAM = os.popen("free -m").readlines()[1].split()[1]
        self.loggerfile.write("RAM Available: %s MB \n" % RAM)
        self.loggerfile.write("================================BEGIN USAGE LOG==================================\n")

    def log(self, foo):
        """# Formats foo nicely and adds it to the log"""
        for each in wrap(foo, 80):
            self.loggerfile.write(each+'\n')

    def save(self):
        """saves the file to disk. Once used you have to re-instance the logger"""
        self.loggerfile.write("\n\n===============================[END OF RECORD]===============================")
        self.loggerfile.flush()
        self.loggerfile.close()