#!/usr/bin/env python
from __future__ import print_function
import sys
import argparse
from datetime import datetime, timedelta
from dateutil.parser import parse as parse_datetime
from gitlab import Gitlab
from gitlab.exceptions import *
import gitlab.config

__version__ = '1.0.1'

try:
    # python3 only
    from datetime import timezone
    def utcnow():
       return datetime.now(timezone.utc)    
except ImportError:
    # python2 fallback.
    try:
        import pytz
    except ImportError:
        print("gitlab-artifact-cleanup requires either Python 3 or the pytz package")
        sys.exit(2)
    def utcnow():
       return datetime.utcnow().replace(tzinfo=pytz.utc)
    

class GitlabArtifactCleanup(object):
    def __init__(self, dry_run=False, min_age=None):
        self.dry_run = dry_run
        self.min_age = timedelta() if min_age is None else min_age

        self.total_size = 0
        self.total_count = 0


    def cleanup_project(self, proj):
        print(proj.name_with_namespace)

        subtotal_size = 0
        subtotal_count = 0
        now = utcnow()

        for build in proj.builds.list(all=True):
            try:
                af = build.artifacts_file
            except AttributeError:
                # Skip builds without artifacts
                continue

            print('  Build {} ({}): '.format(build.id, format_datasize(af['size'])), end='')

            # Skip builds run for tagged commits
            if build.tag:
                print('Skipping, tagged "{}"'.format(build.ref))
                continue

            # Skip builds that are too recent
            ts = parse_datetime(build.created_at)
            age = now - ts
            if age < self.min_age:
                print('Skipping, too recent')
                continue

            print('{action} {filename}'.format(
                action = 'Would delete' if self.dry_run else 'Deleting',
                filename = af['filename'],
                ))

            if not self.dry_run:
                build.erase()

            subtotal_size += af['size']
            self.total_size += af['size']
            subtotal_count += 1
            self.total_count += 1


        if subtotal_count > 0:
            print('  {action} {size} in {count} artifacts'.format(
                action = 'Would delete' if self.dry_run else 'Deleted',
                size = format_datasize(subtotal_size),
                count = subtotal_count))


def format_datasize(b):
    prefixes = ['', 'Ki', 'Mi', 'Gi', 'Ti']

    for pref in prefixes:
        if (b < 1024) or (pref == prefixes[-1]):
            return '{:.02f} {}B'.format(b, pref) if pref else '{} B'.format(b)
        b /= 1024.0

def parse_timedelta(s):
    parts = s.split()
    n = int(parts[0])

    if len(parts) == 1:
        return timedelta(seconds=n)

    if len(parts) == 2:
        unit = parts[1].lower()

        # Allow singular units to be given
        if unit[-1] != 's':
            unit += 's'

        # Support units not supported by timedelta constructor
        if unit == 'years':
            unit = 'days'
            n *= 365
        elif unit == 'months':
            unit = 'days'
            n *= 30

        # Leverage timedelta kwargs for handling most units
        return timedelta(**{unit: n})

    raise ValueError

def check_old_gitlab():
    """Detect an old version of python-gitlab"""
    from gitlab import ProjectBuild
    try:
        ProjectBuild.erase
    except AttributeError:
        print("\nYou have an old version of python-gitlab that is missing the ability to erase builds.")
        print("See https://github.com/gpocentek/python-gitlab/pull/159.")
        print("Monkey-patching a solution into python-gitlab...\n")

        # Monkey-patch the functionality into python-gitlab

        class GitlabBuildEraseError(GitlabRetryError):
            pass

        def erase(self, **kwargs):
            """Erase the build (remove build artifacts and trace)."""
            url = '/projects/%s/builds/%s/erase' % (self.project_id, self.id)
            r = self.gitlab._raw_post(url)
            raise_error_from_response(r, GitlabBuildEraseError, 201)

        ProjectBuild.erase = erase


def parse_args():
    ap = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter)
    ap.add_argument('-g', '--gitlab',
            help='GitLab server defined in config file .python-gitlab')
    ap.add_argument('--all-projects', action='store_true',
            help='Cleanup artifacts for all accessible projects')
    ap.add_argument('-m', '--min-age', type=parse_timedelta, default=timedelta(),
            help='Minimum age for artifacts to be deleted; Examples:\n'
                 '  -m 60    (60 seconds)\n'
                 '  -m "8 hours"\n'
                 '  -m "30 days"\n'
                 '  -m "3 weeks"\n'
                 '  -m "6 months"\n'
                 '  -m "2 years"\n')
    ap.add_argument('-n', '--dry-run', action='store_true',
            help="Don't actually delete anything")
    ap.add_argument('-p', '--project', dest='projects', action='append',
            help='Project for which to cleanup artifacts -- can be given multiple times')
    ap.add_argument('-V', '--version', action='version', version='%(prog)s ' + __version__)

    args = ap.parse_args()

    if not (args.all_projects or args.projects):
        ap.error('-p or --all-projects must be specified')
    if (args.all_projects and args.projects):
        ap.error('-p and --all-projects are mutually exclusive')

    return args

def main():
    args = parse_args()
    check_old_gitlab()

    try:
        gl = Gitlab.from_config(gitlab_id=args.gitlab)
    except gitlab.config.ConfigError as cfgerr:
        print("Error loading python-gitlab config:", cfgerr)
        print("See http://python-gitlab.readthedocs.io/en/stable/cli.html#configuration")
        sys.exit(1)

    print('Would delete' if args.dry_run else 'Deleting', 'non-tagged artifacts', end='')
    if args.min_age:
        print(' older than', args.min_age, end='')
    print('\n')

    cleanup = GitlabArtifactCleanup(
            dry_run = args.dry_run,
            min_age = args.min_age,
            )

    if args.all_projects:
        for proj in gl.projects.all(all=True):
            cleanup.cleanup_project(proj)
    else:
        for pname in args.projects:
            try:
                proj = gl.projects.get(pname)
            except GitlabGetError as e:
                print("Error getting project", pname, e)
                continue
            cleanup.cleanup_project(proj)

    print('\n{action} {size} in {count} artifacts total'.format(
        action = 'Would delete' if args.dry_run else 'Deleted',
        size = format_datasize(cleanup.total_size),
        count = cleanup.total_count))

if __name__ == '__main__':
    main()
