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

from __future__ import division

import argparse
import os
import sys
import subprocess
from pwd import getpwuid
from os import listdir
from os.path import join, isfile, isdir, dirname
import shutil
import requests
import re

from cp3_llbb.SAMADhi.SAMADhi import Dataset, Sample, File, DbStore

def get_options():
    """
    Parse and return the arguments provided by the user.
    """
    username = getpwuid(os.stat(os.getcwd()).st_uid).pw_name
    parser = argparse.ArgumentParser(description='Provide a list of things to be deleted in /storage/data/cms/store/user/')
    parser.add_argument('--crabUsername', action='store', dest='crabUsername', default=username, type=str,
        help='crab / storage username')
    parser.add_argument('--ingridUsername', action='store', dest='ingridUsername', default=username, type=str,
        help='ingrid username')
    parser.add_argument('--debug', action='store_true', help='debug mode: running on only one production and do not delete anything', dest='debug')
    parser.add_argument('-d', '--dry-run', action='store_true', help='dry mode: simulate what would happend but do not actually perform any action', dest='dry')
    parser.add_argument('-s', '--evaluateSize', action='store_true', help='evaluate disk size that would be freed', dest='evaluateSize')
    parser.add_argument('-o', '--output', action='store', help='store list of directories to delete in this file, and do not delete anything', dest='output')
    options = parser.parse_args()
    return options

def main(crabUsername, ingridUsername, DEBUG = False, evaluateSize = False, dryRun=False, output=None):
    if DEBUG:
        print "RUNNING IN DEBUG MODE"
        print "Nothing will be deleted\n"

    if dryRun:
        print("Only simulating what is happening, nothing will be deleted.")

    dbstore = DbStore()

    print "##### Get the list of potential DB samples of interest"
    list_allDBsamples = []
    results = dbstore.find(Sample)
    for r in results:
        if crabUsername in r.path:
            list_allDBsamples.append([r.name, r.sample_id])
    print ""

    print "##### Get the list of existing productions"
    # before anything else: get the list of tags to not touch
    whitelist = requests.get('https://raw.githubusercontent.com/cp3-llbb/GridIn/master/data/SAMADhi_doNOTdelete_whitelist.json').json()
    if DEBUG:
        print "production whitelist= ", whitelist
    list_all_productions = []
    for i, s in enumerate(list_allDBsamples):
        s_name, s_id = s

        isProdAlreadyListed = False
        isSampleProtected = False
        for FWtag, Anatag in list_all_productions:
            if FWtag in str(s_name) and Anatag in str(s_name):
#                print "This prod is already in the list, FWtag= ", FWtag, "Anatag= ", Anatag
                isProdAlreadyListed = True
                break
        if isProdAlreadyListed:
            continue
        tags = str(s_name)
        # Get the tags: 
        # First of all: check if the sample is protected or not
        for ana in whitelist:
            part = str(ana)
            for protectedtag in whitelist[ana]:
                t = str(protectedtag).split('_%s_' % part)
                if t[0] in tags and t[1] in tags:
                    if DEBUG:
                        print '\tSkipping whitelisted sample %s' % s_name
                    isSampleProtected = True
            if not isSampleProtected:
                tags = tags.replace(part, '') # remove HHAnalyzer and the like from the name of the sample
        if isSampleProtected:
            continue
        # now extract the fw and analyzer tags
        # for analyzer, this is always the last part of the sample name so we don't have to worry about naming conventions there (fortunately)
        tags = tags.split('_')
        Anatag = tags[-1]
        tags = tags[:-1]
        # for FW the following regex should work ((v\d+.\d+.\d+\+\d+X?)(-\d+-g[0-9a-f]{7,40})?)|([0-9a-f]{7,40})
        # it matches either:
        #   - a framework tag (possibly with a final X): v1.2.0+7415
        #   - possibly followed by a number of commits and a 'g' plus 7 to 40 characters git hash: v1.2.0+7415-79-ga5b16ff
        #   - or alternatively a 7 to 40 characters git hash: f2f0a44
        tags = [x for x in tags if re.match('((v\d+.\d+.\d+\+\d+X?)(-\d+-g[0-9a-f]{7,40})?)|([0-9a-f]{7,40})', x)]
        if DEBUG:
            print tags, Anatag
        if len(tags) != 1:
            print "ERROR, there are spurious things in the sample name, please figure out what is happening:"
            print "FWtags= ", tags
            return 1
        FWtag = tags[0]
        list_all_productions.append([FWtag, Anatag])

    globalTotalSize = 0

    # If needed, open output file
    f = None
    if output:
        f = open(output, 'w')

    for i, p in enumerate(list_all_productions):
        if DEBUG and i > 0:
            break
        FWtag, Anatag = p

        extrastring = ''
        if not evaluateSize:
            extrastring = ' (evaluation of the disk size is OFF by default)'
        print("\n##### Now looking at prod FWtag: {}, Anatag: {} and list the associated folders{}".format(FWtag, Anatag, extrastring))
        totalSize = 0
        cannotManageToDeleteThisProd = False

        sample_ids = set()

        for s_name, s_id in list_allDBsamples:
            if FWtag in str(s_name) and Anatag in str(s_name):
                result = dbstore.find(Sample, Sample.name == s_name)
                s = result.one()

                if evaluateSize:
                    totalSize += int(subprocess.check_output(["du", '-s', str(s.path)]).split()[0].decode('utf-8'))

                sample_ids.add(s.sample_id)

                if s.source_sample is not None:
                    print "WARNING, the sample", s.name, "depend on another sample, aborting now"
                    cannotManageToDeleteThisProd = True
                    break

                if s.results.count() > 0:
                    print "WARNING, the sample", s.name, "has derived results, aborting now"
                    cannotManageToDeleteThisProd = True
                    break

                if s.derived_samples.count() > 0:
                    # Find all derived samples in the database
                    r = dbstore.find(Sample, Sample.source_sample == s.sample_id)
                    for derived_sample in r:
                        sample_ids.add(derived_sample.sample_id)

        if cannotManageToDeleteThisProd:
            continue

        totalSamples = len(sample_ids)

        globalTotalSize += totalSize

        print("\tNumber of samples: {}, total size: {}".format(totalSamples, sizeof_fmt(totalSize) if evaluateSize else 'N/A'))
        print("")

        message = "\tDo you REALLY want to DELETE this prod from SAMADhi"
        if not f:
            message += " and from disk?"
        else:
            message += "?"

        if dryRun or confirm(prompt=message, resp=False):
            for sample_id in sorted(sample_ids):
                result = dbstore.find(Sample, Sample.sample_id == sample_id)
                s = result.one()

                if s.source_sample:
                    # This sample is a merged sample, no need to remove any files, only the database entry

                    if dryRun:
                        print("\tDeleting merged sample {}.".format(s.name))
                    else:
                        dbstore.remove(s)

                else:

                    if f:
                        f.write(str(s.path).rsplit('/0000', 1)[0] + '\n')

                    if dryRun:
                        print("\tDeleting sample {}.\n\t\tPath on storage: {}".format(s.name, s.path))
                        if DEBUG:
                            print 'rm -r %s' % s.path
                            print 'rm -r %s' % str(s.path).rsplit('/0000', 1)[0]
                    else:
                        try:
                            if not f:
                                shutil.rmtree(s.path)
                                shutil.rmtree(str(s.path).rsplit('/0000', 1)[0])
                        except OSError:
                            print "Seems we have a buggy path: %s" % s.path
                            print "deleting the DB entry then moving on..."
                        dbstore.remove(s)

            dbstore.commit()

    if dryRun:
        print("")
        print("Potential freeable space: {}".format(sizeof_fmt(globalTotalSize) if evaluateSize else 'N/A'))

    if f:
        f.close()
    

# FIXME: deal with extensions and merged samples

def sizeof_fmt(num, suffix='B'):
# Taken from http://stackoverflow.com/questions/1094841/reusable-library-to-get-human-readable-version-of-file-size
    for unit in ['','k','M','G','T','P','E','Z']:
        if abs(num) < 1024.0:
            return "%3.1f%s%s" % (num, unit, suffix)
        num /= 1024.0
    return "%.1f%s%s" % (num, 'Y', suffix)

def confirm(prompt=None, resp=False):
# shamelessly taken from SAMADhi/scripts/userPrompt.py
    """prompts for yes or no response from the user. Returns True for yes and
    False for no. 'resp' should be set to the default value assumed by the caller when
    user simply types ENTER.
    >>> confirm(prompt='Create Directory?', resp=True)
    Create Directory? [y]|n: 
    True
    >>> confirm(prompt='Create Directory?', resp=False)
    Create Directory? [n]|y: 
    False
    >>> confirm(prompt='Create Directory?', resp=False)
    Create Directory? [n]|y: y
    True
    """
    if prompt is None:
        prompt = 'Confirm'
    if resp:
        prompt = '%s [%s]|%s: ' % (prompt, 'y', 'n')
    else:
        prompt = '%s [%s]|%s: ' % (prompt, 'n', 'y')
    while True:
        ans = raw_input(prompt)
        if not ans:
            return resp
        if ans not in ['y', 'Y', 'n', 'N']:
            print 'please enter y or n.'
            continue
        if ans == 'y' or ans == 'Y':
            return True
        if ans == 'n' or ans == 'N':
            return False

if __name__ == '__main__':
    options = get_options()
    main(options.crabUsername, options.ingridUsername, DEBUG = options.debug, evaluateSize = options.evaluateSize, dryRun=options.dry, output=options.output) 
