#!/usr/bin/python
#
# Copyright 2015-2021 University of Southern California
# Distributed under the Apache License, Version 2.0. See LICENSE for more info.
#
import os
import re
import sys
import datetime
import shutil
import argparse
import base64
import logging
import traceback
import requests
from hatrac.core import config

# do a run-time import of deriva-py so that we don't need to declare it as a formal dependency
try:
    from deriva.core import hatrac_store, format_exception, format_credential, get_credential, get_transfer_summary, \
        urlsplit, urlunsplit, make_dirs, BaseCLI, DEFAULT_CHUNK_SIZE, DEFAULT_SESSION_CONFIG
    from deriva.core.utils import hash_utils
except ModuleNotFoundError:
    sys.stderr.write("Cannot import deriva.core python module. Ensure that deriva-py is installed.\n\n")
    sys.exit(2)


DESC = "DERIVA HATRAC Migration Tool"
INFO = "For more information see: https://github.com/informatics-isi-edu/hatrac"
VERSION = 0.2

logger = logging.getLogger(__name__)


class HatracMigrateCLI (BaseCLI):
    """Deriva Hatrac Migration Tool CLI.
    """
    cache_dir_name = "hatrac-migration-cache"

    def __init__(self, description, epilog, version):
        super(HatracMigrateCLI, self).__init__(description, epilog, version)
        self.store = None
        self.backend = None
        self.credentials = None
        self.directory = None
        self.cache_dir = "./%s" % self.cache_dir_name
        self.headers = {'Connection': 'keep-alive', 'Accept-Encoding': 'identity'}
        self.chunk_size = DEFAULT_CHUNK_SIZE
        self.chunk_size_multiplier = 1
        self.progress_percentage_checkpoint = 10

        self.parser.add_argument("-l", "--link-redirects", action="store_true",
                                 help="Set ALL object version redirect fields to: "
                                 "{<host> + <path-prefix> + <existing object name>}")
        self.parser.add_argument("-t", "--transfer-linked-redirects", action="store_true",
                                 help="Transfer files that have linked redirects to this instance, "
                                 "deleting link redirect on successful transfer completion.")
        self.parser.add_argument("-p", "--path-prefix", metavar="</path>",
                                 help="Base resource path prefix")
        self.parser.add_argument("-c", "--cache-dir", metavar="<path>",
                                 help="Cache directory to use for transfers. Defaults to '%s'." % self.cache_dir)
        self.parser.add_argument("-a", "--allow-missing-hashes", action="store_true", default=False,
                                 help="If hashes for existing objects cannot be found in both the source and "
                                      "destination, allow them to be generated based on locally downloaded content "
                                      "from the source.")
        self.parser.add_argument("-w", "--whitelist-path-pattern", metavar="regex>",
                                 help="Regular expression to apply to object names for inclusion in "
                                      "linking/transferring.")
        self.parser.add_argument("-e", "--enumerate-linked-objects", action="store_true",
                                 help="Enumerate the set of linked objects that have not yet been migrated.")
        self.parser.add_argument("-d", "--dry-run", action="store_true", default=False,
                                 help="Dry run will not modify the database when linking and will not any transfer any "
                                      "previously linked files. It will simply log the list of files that would be "
                                      "linked, based on the input parameters.")
        self.parser.add_argument("-x", "--chunk-size-multiplier", metavar="<1-10>", type=int,
                                 choices=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
                                 help="Chunk size multiplier to use for chunked transfers. "
                                      "Defaults to %dx of default chunk size %d bytes." %
                                      (self.chunk_size_multiplier, self.chunk_size),
                                 default=self.chunk_size_multiplier)
        self.parser.add_argument("-i", "--transfer-progress-info-percentage", metavar="<1-100>", type=int,
                                 help="Transfer statistics will be logged at this completion percentage interval. "
                                      "Defaults to %s percent." % self.progress_percentage_checkpoint,
                                 default=self.progress_percentage_checkpoint)

    @staticmethod
    def make_session_config():
        session_config = DEFAULT_SESSION_CONFIG.copy()
        session_config.update({
            "retry_status_forcelist": [502, 503, 504],
        })
        return session_config

    def enumerate_objects(self, linked=False, whitelist_regex=None):
        objects = dict()
        try:
            root = self.directory.name_resolve("/")
            names = self.directory.namespace_enumerate_names(root)
            for name in names:
                if name.is_object():
                    versions = self.directory.object_enumerate_versions(name)
                    for version in versions:
                        if linked:
                            url = version.aux.get("url")
                            if url:
                                objects[url] = version
                        else:
                            objects["%s:%s" % (version.name, version.version)] = version

        except Exception as e:
            raise RuntimeError(
                "Exception while enumerating object versions from the database: %s" % format_exception(e))
        total = len(objects)
        logger.info("Found %d%s object versions." % (total, " linked" if linked else ""))

        # 2. Prune the object set based on the whitelist regex, if any
        pruned_objects = dict()
        if whitelist_regex:
            logger.info("Applying whitelist regex '%s' to prune object set." % whitelist_regex)
            regex = re.compile(whitelist_regex)
            for key, obj in objects.items():
                path = "%s:%s" % (obj.name, obj.version)
                if regex.match(path):
                    if logger.isEnabledFor(logging.DEBUG):
                        logger.debug("Including object path [%s] in whitelisted object set." % path)
                    pruned_objects[key] = obj
            total = len(pruned_objects)
            logger.info("Filtered object set now contains %d entries." % total)
        else:
            pruned_objects = objects

        return pruned_objects

    def link_objects(self, base_url, whitelist_regex=None, dry_run=False):
        # 1. Enumerate the objects to link
        logger.info("Enumerating existing object versions. %sPlease wait..." %
                    ("Dry-run is enabled. " if dry_run else ""))
        objects = self.enumerate_objects(linked=False, whitelist_regex=whitelist_regex)

        # 3. Link the objects to the source URLs
        succeeded = 0
        for path, obj in objects.items():
            try:
                msg = "Linking object [%s:%s] to [%s]" % (obj.name, obj.version, base_url + obj.asurl())
                if dry_run:
                    logger.info(msg)
                else:
                    if logger.isEnabledFor(logging.DEBUG):
                        logger.debug(msg)
                    self.directory.version_aux_url_update(obj, base_url)
                    succeeded += 1
            except Exception as e:
                logger.warning(
                    "Exception while attempting to link object [%s:%s] to [%s]: %s" %
                    (obj.name, obj.version, base_url + obj.asurl(), format_exception(e)))
        if not dry_run:
            logger.info("Successfully linked %d of %d object versions." % (succeeded, len(objects)))

    def transfer_linked_objects(self, allow_missing_hashes=False, whitelist_regex=None):
        # Get the current working set of linked objects.
        logger.info("Enumerating linked object versions. Please wait...")
        objects = self.enumerate_objects(linked=True, whitelist_regex=whitelist_regex)

        logger.info("Attempting to transfer file(s)...")
        for url, obj in objects.items():
            if not self.store:
                host = urlsplit(url)
                if not self.credentials:
                    self.credentials = get_credential(host.hostname)
                self.store = hatrac_store.HatracStore(
                    host.scheme, host.hostname, self.credentials, session_config=self.make_session_config())

            # 1. GET the specified file from the remote hatrac system into the working cache.
            success = False
            cached = False
            input_file = None
            filename = os.path.basename(url)
            resource_path = urlsplit(url).path
            dirname = hash_utils.compute_hashes(resource_path.encode())['md5'][0]
            output_dir = os.path.abspath(os.path.join(self.cache_dir, dirname))
            file_path = os.path.abspath(os.path.join(output_dir, filename))
            try:
                if logger.isEnabledFor(logging.DEBUG):
                    logger.debug("Performing HEAD request for metadata on resource: %s" % resource_path)
                resp = self.store.head(resource_path, self.headers)
                file_size = int(resp.headers.get('Content-Length', 0))
                file_md5 = resp.headers.get('Content-MD5') or ""
                if not file_md5:
                    logger.warning("Remote resource %s is missing 'Content-MD5' header." % resource_path)
                file_sha256 = resp.headers.get('Content-SHA256') or ""
                if not file_sha256:
                    logger.warning("Remote resource %s is missing 'Content-SHA256' header." % resource_path)
                if os.path.isfile(file_path):
                    logger.info("Computing hashes for file %s found in local cache: [%s]." %
                                (filename, file_path))
                    cached_file_size = os.path.getsize(file_path)
                    cached_hashes = hash_utils.compute_file_hashes(file_path, ["md5", "sha256"])
                    cached_md5 = cached_hashes["md5"][1]
                    cached_sha256 = cached_hashes["sha256"][1]
                    if (cached_file_size == file_size) and (cached_md5 == file_md5 or cached_sha256 == file_sha256):
                        cached = True
                    else:
                        logger.warning("Cached file [%s] exists but does not pass validation check. "
                                       "Local file size: [%d]. Remote file size: [%d]. "
                                       "Local file MD5: [%s]. Remote file MD5: [%s]. "
                                       "Local file SHA256: [%s]. Remote file SHA256: [%s]."
                                       % (file_path, cached_file_size, file_size, cached_md5, file_md5,
                                          cached_sha256, file_sha256))
                if not cached:
                    logger.info("Transferring [%s] with size %s bytes to cache location: %s" %
                                (url, file_size, file_path))
                    make_dirs(output_dir, mode=0o755)
                    chunk_progress_threshold = int((file_size / self.chunk_size) *
                                                   min((self.progress_percentage_checkpoint * .01), 1))
                    start = datetime.datetime.now()

                    def cache_transfer_callback(**kwargs):
                        try:
                            elapsed_time = datetime.datetime.now() - start
                            total_bytes = kwargs.get("total_bytes", None)
                            current_chunk = kwargs.get("current_chunk", None)
                            if not (total_bytes and current_chunk and chunk_progress_threshold):
                                return True
                            if total_bytes == file_size:
                                return True
                            if (current_chunk % chunk_progress_threshold) == 0:
                                logger.info("Transferred %s%% of file [%s] with %s" %
                                            (round((total_bytes / file_size) * 100),
                                             file_path,
                                             get_transfer_summary(total_bytes, elapsed_time)))
                                return True
                        except:
                            traceback.print_exc()
                            raise
                        finally:
                            return True

                    self.store.get_obj(resource_path,
                                       headers={'Connection': 'keep-alive'},
                                       destfilename=file_path,
                                       chunk_size=self.chunk_size,
                                       callback=cache_transfer_callback)
                    success = True

            except requests.HTTPError as e:
                logger.error("File [%s] transfer failed: %s" % (file_path, e))
                continue
            except Exception as e:
                logger.error(format_exception(e))
                continue
            finally:
                if not success:
                    if os.path.isdir(output_dir):
                        shutil.rmtree(output_dir)
                else:
                    success = False

            # 2. Sanity check the hashes against the known object version hashes from the database.
            md5 = base64.b64encode(obj.metadata.get("content-md5", b"")).decode()
            if (md5 and file_md5) and (md5 != file_md5):
                logger.warning(
                    "Downloaded file MD5 [%s] does not equal known object version MD5 [%s]. Skipping." %
                    (file_md5, md5))
                continue

            sha256 = base64.b64encode(obj.metadata.get("content-sha256", b"")).decode()
            if (sha256 and file_sha256) and (sha256 != file_sha256):
                logger.warning(
                    "Downloaded file SHA256 [%s] does not equal known object version SHA256 [%s]. Skipping." %
                    (file_sha256, sha256))
                continue

            # If hashes don't exist in both remote and local, create them if "allow_missing_hashes" is enabled.
            # No matter what, skip on size mismatch.
            if not (md5 and file_md5) and not (sha256 and file_sha256):
                if not allow_missing_hashes:
                    logger.warning("Missing both local and remote MD5 and SHA256 hashes. Skipping.")
                    continue
                if os.path.getsize(file_path) != file_size:
                    logger.warning("Missing both local and remote MD5 and SHA256 hashes, and local file size does not "
                                   "equal remote file size. Skipping.")
                    continue
                logger.warning("Missing both local and remote MD5 and SHA256 hashes. "
                               "Calculating new hash values based on local content.")
                hashes = hash_utils.compute_file_hashes(file_path, ["md5", "sha256"])
                obj.metadata.update({"content-md5": base64.b64decode(hashes["md5"][1]),
                                     "content-sha256": base64.b64decode(hashes["sha256"][1])})

            # 3. Transfer the input file to the storage backend.
            try:
                if self.backend == "filesystem":
                    version = obj.version
                    dest_file_path = "%s/%s" % (self.directory.storage._dirname_relname(obj.name, version))
                    make_dirs(dest_file_path, mode=0o755)
                    logger.info("Moving file [%s] from local cache to [%s]." % (file_path, dest_file_path))
                    shutil.move(file_path, dest_file_path)
                else:
                    input_file = open(file_path, "rb")

                    # 3.1 Non-chunked transfer if under self.chunk_size
                    if file_size < self.chunk_size:
                        logger.info("Transferring file [%s] from local cache to backend storage type: [%s]." %
                                    (file_path, self.backend))
                        version = self.directory.storage.create_from_file(obj.name,
                                                                          input_file,
                                                                          file_size,
                                                                          obj.metadata)
                    # 3.2 Do chunked transfer.
                    else:
                        logger.info("Transferring file [%s] from local cache to backend storage type [%s] using "
                                    "chunked upload." % (file_path, self.backend))

                        # 3.2.1 Create upload job
                        tbytes = 0
                        position = 0
                        chunk_aux = list()
                        remainder = file_size % self.chunk_size
                        nchunks = file_size // self.chunk_size
                        tchunks = nchunks if remainder == 0 else nchunks + 1
                        start = datetime.datetime.now()
                        upload = self.directory.storage.create_upload(obj.name, file_size, obj.metadata)

                        # 3.2.2 Upload chunks
                        while tbytes < file_size:
                            nbytes = self.chunk_size if position < nchunks else remainder
                            metadata = obj.metadata.copy()
                            c_start = datetime.datetime.now()
                            try:
                                aux = self.directory.storage.upload_chunk_from_file(obj.name,
                                                                                    upload,
                                                                                    position,
                                                                                    self.chunk_size,
                                                                                    input_file,
                                                                                    nbytes,
                                                                                    metadata)
                            except:
                                self.directory.storage.cancel_upload(obj.name, upload)
                                raise

                            chunk_aux.append({"position": position, "aux": aux})
                            position += 1
                            tbytes += nbytes
                            c_elapsed = datetime.datetime.now() - c_start
                            c_summary = get_transfer_summary(nbytes, c_elapsed)
                            if logger.isEnabledFor(logging.DEBUG):
                                logger.debug("Object [%s] successfully uploaded chunk: %d of %d. "
                                             "Total bytes processed: %d of %d. %s" %
                                             (obj.name, position, tchunks, tbytes, file_size, c_summary))

                        # 3.2.3 Finalize upload job
                        version = self.directory.storage.finalize_upload(obj.name, upload, chunk_aux, obj.metadata)
                        elapsed = datetime.datetime.now() - start
                        summary = get_transfer_summary(tbytes, elapsed)
                        logger.info("Chunked upload of file [%s] successful. %s" % (file_path, summary))

                # 3.3 Post-transfer database updates
                # We only care about storing the returned aux version in the non-filesystem backend case
                version = version if self.backend != "filesystem" else None
                self.directory.version_aux_version_update(obj, version)
                # On a successful transfer, we will automatically delete the aux redirect url
                self.directory.version_aux_url_delete(obj)
                success = True
            except Exception as e:
                logger.error("Error during transfer of [%s] to backend storage: %s" %
                             (file_path, format_exception(e)))
                continue
            finally:
                if input_file:
                    input_file.close()
                if success:
                    if os.path.isdir(output_dir):
                        shutil.rmtree(output_dir)

    def main(self):
        args = self.parse_cli()
        self.chunk_size *= args.chunk_size_multiplier
        if args.cache_dir:
            self.cache_dir = os.path.join(args.cache_dir, self.cache_dir_name)

        if args.transfer_progress_info_percentage:
            self.progress_percentage_checkpoint = args.transfer_progress_info_percentage

        # credential initialization
        if args.host:
            token = args.token if args.token else None
            oauth2_token = args.oauth2_token if args.oauth2_token else None
            credential_file = args.credential_file if args.credential_file else None
            if token or oauth2_token:
                self.credentials = format_credential(token=token, oauth2_token=oauth2_token)
            elif credential_file:
                self.credentials = get_credential(args.host, credential_file)

        try:
            # This will generally fail unless run locally in ~/hatrac on the target hatrac system for the migration
            from hatrac import directory
            self.directory = directory
            self.backend = config.get('storage_backend')

            # Enumerate the set of linked objects.
            if args.enumerate_linked_objects:
                logger.info("Enumerating linked object versions. Please wait...")
                self.enumerate_objects(linked=True, whitelist_regex=args.whitelist_path_pattern)

            # The "link_redirects" logic will set alternate "redirect" URLs for every object version in the database
            # based on the following path components:
            # <host> + <path-prefix> + <existing object name from the database name table>.
            # The server will then automatically redirect requests for these objects to these URLs.
            elif args.link_redirects:
                if not args.host:
                    raise argparse.ArgumentError(
                        None, 'The "--host" argument is required when the "--link-redirects" argument is specified.')
                base_url = (args.host if args.host.startswith("http") else "https://" + args.host)
                self.directory.prefix = \
                    args.path_prefix if args.path_prefix else config.get("service_prefix", "/hatrac")
                self.link_objects(base_url, whitelist_regex=args.whitelist_path_pattern, dry_run=args.dry_run)

            # The "transfer_linked_redirects" logic will attempt to transfer files referenced by any pre-existing
            # redirect links to the "local" hatrac instance. This could involve local file caching when additional IO
            # is performed to backend storage off system, e.g. to S3.
            elif args.transfer_linked_redirects:
                if args.dry_run:
                    logger.info("Skipping transfer of linked objects for dry-run.")
                else:
                    self.transfer_linked_objects(allow_missing_hashes=args.allow_missing_hashes,
                                                 whitelist_regex=args.whitelist_path_pattern)
            else:
                self.parser.print_usage()
            return 0
        except argparse.ArgumentError as e:
            logger.error(format_exception(e))
        except RuntimeError as e:
            logger.error(format_exception(e))
            if args.debug:
                traceback.print_exc()
        except KeyboardInterrupt:
            logger.warning("Process interrupted by user.")
        except:
            logger.error("Unhandled exception!")
            traceback.print_exc()
        return 1


def main():
    return HatracMigrateCLI(DESC, INFO, VERSION).main()


if __name__ == '__main__':
    sys.exit(main())
