#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Reconcile and cleanse where necessary an OMERO data directory of orphaned data.
"""
from __future__ import print_function

#
#  Copyright (c) 2009-2016 University of Dundee. All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions
#  are met:
#  1. Redistributions of source code must retain the above copyright
#     notice, this list of conditions and the following disclaimer.
#  2. Redistributions in binary form must reproduce the above copyright
#     notice, this list of conditions and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
#
#  THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
#  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
#  ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
#  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
#  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
#  OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
#  HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
#  LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
#  OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
#  SUCH DAMAGE.

from builtins import range
from builtins import object
import omero.clients
import omero
import sys
import os
import getpass
import Ice

from Glacier2 import PermissionDeniedException
from getopt import getopt, GetoptError
from omero.util import get_user
from math import ceil
from stat import ST_SIZE


# The directories underneath an OMERO data directory to search for "dangling"
# files and reconcile with the database. Directory name key and corresponding
# OMERO data type value.
SEARCH_DIRECTORIES = {
    'Pixels': 'Pixels',
    'Files': 'OriginalFile',
    'Thumbnails': 'Thumbnail'
}


def usage(error):
    """
    Prints usage so that we don't have to. :)
    """
    cmd = sys.argv[0]
    print("""%s
Usage: %s [--dry-run] [-u username | -k] <omero.data.dir>
Cleanses files in the OMERO data directory that have no reference in the
OMERO database. NOTE: As this script is designed to be run via cron or in
a scheduled manner it produces NO output unless a dry run is performed.

Options:
  -u          Administrator username to log in to OMERO with
  -k          Session key to log in to OMERO with
  --dry-run   Just prints out what would have been done

Examples:
  %s --dry-run -u root /OMERO

Report bugs at https://forum.image.sc/""" % \
        (error, cmd, cmd))
    sys.exit(2)


class Cleanser(object):

    """
    Keeps file cleansing state and performs OMERO database reconciliation of
    files within an OMERO binary repository.
    """

    # Number of objects to defer before we actually make a query
    QUERY_THRESHOLD = 25

    # Strings identifying pyramid files
    PYRAMID_FILE = "_pyramid"
    PYRAMID_LOCK = ".pyr_lock"
    PYRAMID_TEMP = ".tmp"

    def __init__(self, query_service, object_type):
        self.query_service = query_service
        self.object_type = object_type
        self.cleansed = list()
        self.bytes_cleansed = 0
        self.deferred_paths = list()
        self.dry_run = False

    def cleanse(self, root):
        """
        Begins a cleansing operation from a given OMERO binary repository
        root directory. /OMERO/Files or /OMERO/Pixels for instance.
        """
        for file in os.listdir(root):
            path = os.path.join(root, file)
            if os.path.isdir(path):
                self.cleanse(path)
            else:
                self.query_or_defer(path)

    def query_or_defer(self, path):
        """
        Adds a given path to the list of deferred paths. If the number of
        deferred paths has reached the QUERY_THRESHOLD (to reduce database
        hits) a reconciliation check will happen against OMERO.
        """
        self.deferred_paths.append(path)
        if len(self.deferred_paths) == self.QUERY_THRESHOLD:
            self.do_cleanse()

    def do_cleanse(self):
        """
        Actually performs the reconciliation check against OMERO and
        removes relevant files.
        """
        if len(self.deferred_paths) == 0:
            return
        split = os.path.split
        object_ids = []
        for path in self.deferred_paths:
            file_name = split(path)[1]
            try:
                object_id = omero.rtypes.rlong(int(file_name))
            except ValueError:
                try:
                    file_name.index(self.PYRAMID_FILE)
                    id_part = file_name.split("_")[0]
                    if file_name.endswith(self.PYRAMID_FILE):
                        object_id = omero.rtypes.rlong(int(id_part))
                    elif (file_name.endswith(self.PYRAMID_LOCK)
                            or file_name.endswith(self.PYRAMID_TEMP)):
                        object_id = omero.rtypes.rlong(
                            int(id_part.lstrip('.')))
                    else:
                        object_id = omero.rtypes.rlong(-1)
                except ValueError:
                    object_id = omero.rtypes.rlong(-1)
            object_ids.append(object_id)

        parameters = omero.sys.Parameters()
        parameters.map = {'ids': omero.rtypes.rlist(object_ids)}
        rows = self.query_service.projection(
            "select o.id from %s as o where o.id in (:ids)" % self.object_type,
            parameters, {"omero.group": "-1"})
        existing_ids = [cols[0].val for cols in rows]

        for i, object_id in enumerate(object_ids):
            path = self.deferred_paths[i]
            if object_id.val not in existing_ids:
                if object_id.val == -1:
                    if self.dry_run:
                        print("   \_ %s (ignored/keep)" % path)
                else:
                    size = os.stat(path)[ST_SIZE]
                    self.cleansed.append(path)
                    self.bytes_cleansed = size
                    if self.dry_run:
                        print("   \_ %s (remove)" % path)
                    else:
                        try:
                            os.unlink(path)
                        except OSError as e:
                            print(e)
            elif self.dry_run:
                print("   \_ %s (keep)" % path)
        self.deferred_paths = list()

    def finalize(self):
        """
        Takes the final set of deferred paths and performs a reconciliation
        check against OMERO for them. This method's purpose is basically to
        catch the final set of paths in the deferred path list and/or perform
        any cleanup.
        """
        self.do_cleanse()

    def __str__(self):
        return "Cleansing context: %d files (%d bytes)" % \
            (len(self.cleansed), self.bytes_cleansed)


def initial_check(config_service, admin_service=None):
    if admin_service is None:
        raise Exception("No admin service provided!")

    ctx = admin_service.getEventContext()
    if not ctx.isAdmin:
        raise Exception('SecurityViolation: Admins only!')

    #
    # Compare server versions. See ticket #3123
    #
    if config_service is None:
        print("No config service provided! "
              "Waiting 10 seconds to allow cancellation")
        from threading import Event
        Event().wait(10)

    server_version = config_service.getVersion()
    server_tuple = tuple([int(x) for x in server_version.split(".")])
    if server_tuple < (4, 2, 1):
        print("Server version is too old! (%s) Aborting..." % server_version)
        sys.exit(3)


def cleanse(data_dir, client, dry_run=False):
    client.getImplicitContext().put(omero.constants.GROUP, '-1')

    admin_service = client.sf.getAdminService()
    query_service = client.sf.getQueryService()
    config_service = client.sf.getConfigService()

    initial_check(config_service, admin_service)

    try:
        cleanser = ""
        for directory in SEARCH_DIRECTORIES:
            full_path = os.path.join(data_dir, directory)
            if not os.path.exists(full_path):
                print("%s does not exist. Skipping..." % full_path)
                continue
            if dry_run:
                print("Reconciling OMERO data directory...\n %s" % full_path)
            object_type = SEARCH_DIRECTORIES[directory]
            cleanser = Cleanser(query_service, object_type)
            cleanser.dry_run = dry_run
            cleanser.cleanse(full_path)
            cleanser.finalize()
    finally:
        if dry_run:
            print(cleanser)

    # delete empty directories from the managed repositories
    proxy, description = client.getManagedRepository(description=True)
    if proxy:
        root = description.path.val + description.name.val
        print("Removing empty directories from...\n %s" % root)
        delete_empty_dirs(proxy, root, client, dry_run)


def delete_empty_dirs(repo, root, client, dry_run):
    # empty subdirectories are to be appended to this list
    to_delete = []

    # find the empty subdirectories
    is_empty_dir(repo, '/', False, to_delete)

    if dry_run:
        for directory in to_delete:
            print("   \_ %s%s (remove)" % (root, directory))
    elif to_delete:
        # probably less than a screenful
        batch_size = 20

        for from_index in range(0, len(to_delete), batch_size):
            batch_to_delete = to_delete[from_index:from_index + batch_size]
            for directory in batch_to_delete:
                print("Removing %s%s" % (root, directory))
            handle = repo.deletePaths(batch_to_delete, True, False)
            try:
                client.waitOnCmd(handle, closehandle=True)
            except omero.CmdError as ce:
                if isinstance(ce.err, omero.cmd.GraphException):
                    raise Exception("failed delete: " + ce.err.message)
                else:
                    raise Exception("failed: " + ce.err.name)


def is_empty_dir(repo, directory, may_delete_dir, to_delete):
    empty_subdirs = []
    is_empty = True

    for entry in repo.listFiles(directory):
        subdirectory = directory + entry.name.val + '/'
        may_delete_subdir = entry.details.permissions.canDelete()
        if entry.mimetype is not None and \
           entry.mimetype.val == 'Directory' and \
           is_empty_dir(repo, subdirectory, may_delete_subdir, empty_subdirs):
            if may_delete_subdir:
                # note empty subdirectories that can be deleted
                empty_subdirs.append(subdirectory)
        else:
            is_empty = False

    if not (may_delete_dir and is_empty):
        # cannot delete this directory, so note empty subdirectories
        to_delete.extend(empty_subdirs)

    return is_empty


def fixpyramids(data_dir, query_service,
                dry_run=False, config_service=None, admin_service=None):
    initial_check(config_service, admin_service)

    # look for any pyramid files with length 0
    # if there is no matching .*.tmp or .*.pyr_lock file, then
    # the pyramid file will be removed

    pixels_dir = os.path.join(data_dir, "Pixels")
    for root, dirs, files in os.walk(pixels_dir):
        for f in files:
            pixels_file = os.path.join(root, f)
            length = os.path.getsize(pixels_file)
            if length == 0 and f.endswith("_pyramid"):
                delete_pyramid = True
                for lockfile in os.listdir(pixels_dir):
                    if lockfile.startswith("." + f) and \
                        (lockfile.endswith(".tmp") or
                            lockfile.endswith(".pyr_lock")):
                        delete_pyramid = False
                        break

                if delete_pyramid:
                    if dry_run:
                        print("Would remove %s" % f)
                    else:
                        print("Removing %s" % f)
                        os.remove(pixels_file)


def removepyramids(client, little_endian=None, dry_run=False,
                   imported_after=None, wait=25, limit=500):
    client.getImplicitContext().put(omero.constants.GROUP, '-1')
    admin_service = client.sf.getAdminService()
    config_service = client.sf.getConfigService()

    initial_check(config_service, admin_service)
    value = int(limit)
    # If the default limit is changed, please update the help
    # in admin.py
    if value > 500 or value <= 0:
        value = int(500)
    print("No more than %s pyramids will be removed" % value)
    # look for any pyramid files with the specified endianness
    # the pyramid file will be removed
    request = omero.cmd.FindPyramids()
    request.littleEndian = little_endian
    request.importedAfter = imported_after
    request.limit = value

    cb = None
    ms = 500
    loops = ceil(wait * 1000.0 / ms)

    try:
        cb = client.submit(request, loops=loops, ms=ms,
                           failonerror=True,
                           failontimeout=True)
        rsp = cb.getResponse()
    except omero.CmdError as ce:
        print("Failed to load pyramids: %s" % ce.err.name)
        return
    finally:
        if cb:
            cb.close(True)

    if len(rsp.pyramidFiles) == 0:
        print("No pyramids to remove")
        return

    for j in range(len(rsp.pyramidFiles)):
        image_id = rsp.pyramidFiles[j]
        if dry_run:
            print("Would remove pyramid for image %s" % image_id)
        else:
            req = omero.cmd.ManageImageBinaries()
            req.imageId = image_id
            req.deletePyramid = True
            req.deleteThumbnails = True
            try:
                cb = client.submit(req, loops=loops, ms=ms,
                                   failonerror=True,
                                   failontimeout=True)
                print("Pyramid removed for image %s" % image_id)
            except omero.CmdError as ce:
                print("Failed to remove for image %s: %s" % (
                    image_id, ce.err.name))
            finally:
                if cb:
                    try:
                        cb.close(True)
                    except Ice.NotRegisteredException:
                        print("Error closing callback for %s" % image_id)


def main():
    """
    Default main() that performs OMERO data directory cleansing.
    """
    try:
        options, args = getopt(sys.argv[1:], "u:k:", ["dry-run"])
    except GetoptError as xxx_todo_changeme:
        (msg, opt) = xxx_todo_changeme.args
        usage(msg)

    try:
        data_dir, = args
    except:
        usage('Expecting single OMERO data directory!')

    username = get_user("root")
    session_key = None
    dry_run = False
    for option, argument in options:
        if option == "-u":
            username = argument
        if option == "-k":
            session_key = argument
        if option == "--dry-run":
            dry_run = True

    if session_key is None:
        print("Username: %s" % username)
        try:
            password = getpass.getpass()
        except KeyboardInterrupt:
            sys.exit(2)

    try:
        client = omero.client('localhost')
        client.setAgent("OMERO.cleanse")
        if session_key is None:
            client.createSession(username, password)
        else:
            client.createSession(session_key)
    except PermissionDeniedException:
        print("%s: Permission denied" % sys.argv[0])
        print("Sorry.")
        sys.exit(1)

    try:
        cleanse(data_dir, client, dry_run)
    finally:
        if session_key is None:
            client.closeSession()


if __name__ == '__main__':
    main()
