#!/opt/dosis2/bin/python3-venv
import json
import os
import re
import shutil
import subprocess
import syslog as sl
from collections import namedtuple
from subprocess import check_call
from syslog import syslog

File = namedtuple('File', ['size', 'date', 'name'])


class Rotate(object):
    def __init__(self, settings_file, **kwargs):
        self._settings_file = settings_file
        self._debug = kwargs.get('debug', False)

        with open(self._settings_file, 'r') as f:
            self._settings = json.load(f)

        self._allow_dir_access = self._settings.get('allow_dir_access', [])
        self._disk_reserve = self._settings.get('disk_reserve', 10240)
        _directories = self._settings.get('directories', {})

        self._directories = sorted(
            [(c.get('order', 100000000), d, c) for d, c in list(_directories.items())]
        )

    def syscall(self, cmd):
        if not self._debug:
            try:
                check_call(cmd)
                return 0
            except subprocess.CalledProcessError as e:
                return e.returncode
        return 0

    def remove_file(self, filename):
        if not self._debug:
            try:
                if os.path.isfile(filename):
                    os.remove(filename)
                elif os.path.isdir(filename):
                    shutil.rmtree(filename)
                elif os.path.islink(filename):
                    os.unlink(filename)
                return True
            except Exception as e:
                self.log(e)
                return False
        return True

    def log(self, msg):
        syslog(sl.LOG_INFO, 'Rotate.py - %s' % msg)
        print(msg)

    def log_err(self, msg):
        syslog(sl.LOG_ERR, 'Rotate.py - %s' % msg)
        print(msg)

    def disk_free_space(self):
        _statvfs = os.statvfs('/')
        return (_statvfs.f_frsize * _statvfs.f_bavail) / (1024.0 ** 2)

    def directory_size(self, directory):
        size = 0
        for path, dnames, fnames in os.walk(directory):
            for f in fnames:
                size += os.path.getsize(os.path.join(path, f))
            for d in dnames:
                size += os.path.getsize(os.path.join(path, d))
        return size / (1024.0 ** 2)

    def run(self):
        self.log('starting')
        self.log('free disk space before: %sM' % self.disk_free_space())
        for order, directory, config in self._directories:
            self.cleanup(directory, **config)

        _freespace = self.disk_free_space()
        self.log('free disk space after: %sM' % _freespace)

        if _freespace >= self._disk_reserve:
            self.log('complete - success')
        else:
            self.log_err('complete - failed to meet disk reserve!')

    def cleanup(self, directory, **config):
        self.log('cleanup: %s' % directory)

        _size_type = config.get('size_type', '')
        _size = config.get('size')
        _logrotate_files = config.get('logrotate_files', [])

        _dirsize = self.directory_size(directory)

        if _logrotate_files:
            for f in _logrotate_files:
                try:
                    self.log('  logrotate %s' % f)
                    self.cleanup_logrotate(f)
                except Exception as e:
                    self.log_err('Error during logrotate for %s: %s' % (f, e))

        try:
            if _size_type.lower() == 'chop':
                if _size is not None:
                    self.log('  clean type=%s size=%s dirsize=%sM' % (_size_type, _size, _dirsize))
                    self.cleanup_files_chop(
                        directory, _dirsize, 
                        config.get('cleanup_type', 'files'),
                        config.get('cleanup', {}),
                        _size,
                        log_indent='    '
                    )
                    _cleaned = _dirsize - self.directory_size(directory)
                    self.log('  cleaned up %sM' % _cleaned)
            elif _size_type.lower() == 'limit':
                if _size is not None:
                    self.log('  clean type=%s size=%s dirsize=%sM' % (_size_type, _size, _dirsize))
                    self.cleanup_files_limit(
                        directory, _dirsize,
                        config.get('cleanup_type', 'files'),
                        config.get('cleanup', {}),
                        _size,
                        log_indent='    '
                    )
                    _cleaned = _dirsize - self.directory_size(directory)
                    self.log('  cleaned up %sM' % _cleaned)
            else:
                self.log('  clean dirsize=%sM' % _dirsize)
                self.cleanup_files_nosize(
                    directory, _dirsize,
                    config.get('cleanup_type', 'files'),
                    config.get('cleanup', {}),
                    log_indent='    '
                )

                _cleaned = _dirsize - self.directory_size(directory)
                self.log('  cleaned up %sM' % _cleaned)
        except Exception as e:
            self.log_err('Error during file cleanup: %s' % e)

    def cleanup_logrotate(self, logrotate_file):
        self.syscall(['/usr/sbin/logrotate', logrotate_file])

    def get_matching_files(self, start_dir, regex_list, ftype='file'):
        for r in regex_list:
            for path, dnames, fnames in os.walk(start_dir):
                _rpath = os.path.relpath(path, start_dir)
                if _rpath == '.':
                    _rpath = ''

                if ftype.lower() == 'file':
                    for f in fnames:
                        if re.match(r, os.path.join(_rpath, f)):
                            _file = os.path.join(path, f)
                            yield File(os.path.getsize(_file) / (1024.0 ** 2), os.path.getmtime(_file), _file)
                elif ftype.lower() == 'dir':
                    for d in dnames:
                        if re.match(r, os.path.join(_rpath, d)):
                            _dir = os.path.join(path, d)
                            yield File(self.directory_size(_dir), os.path.getmtime(_dir), _dir)

    def cleanup_files_chop(self, directory, dirsize, ftype, files, chop_size, log_indent=''):
        _free_space = self.disk_free_space()
        if _free_space < self._disk_reserve:
            if ftype.lower() == 'files':
                _files_to_clean = self.get_cleanup_files_chop(
                    directory, dirsize, chop_size, 'file',
                    remove_all=files.get('remove_all', []),
                    sort_oldest_first=files.get('sort_oldest_first', [])
                )
                self.cleanup_raw(_files_to_clean, log_indent=log_indent)
            elif ftype.lower() == 'directories':
                _files_to_clean = self.get_cleanup_files_chop(
                    directory, dirsize, chop_size, 'dir',
                    remove_all=files.get('remove_all', []),
                    sort_oldest_first=files.get('sort_oldest_first', [])
                )
                self.cleanup_raw(_files_to_clean, log_indent=log_indent)

    def cleanup_files_limit(self, directory, dirsize, ftype, files, limit_size, log_indent=''):
        if dirsize > limit_size:
            if ftype.lower() == 'files':
                _files_to_clean = self.get_cleanup_files_limit(
                    directory, dirsize, limit_size, 'file',
                    keep_latest_one=files.get('keep_latest_one', []),
                    remove_all=files.get('remove_all', []),
                    sort_oldest_first=files.get('sort_oldest_first', [])
                )
                self.cleanup_raw(_files_to_clean, log_indent=log_indent)
            elif ftype.lower() == 'directories':
                _files_to_clean = self.get_cleanup_files_limit(
                    directory, dirsize, limit_size, 'dir',
                    keep_latest_one=files.get('keep_latest_one', []),
                    remove_all=files.get('remove_all', []),
                    sort_oldest_first=files.get('sort_oldest_first', [])
                )
                self.cleanup_raw(_files_to_clean, log_indent=log_indent)

    def cleanup_files_nosize(self, directory, dirsize, ftype, files, log_indent=''):
        _dirsize = self.directory_size(directory)
        self.log(log_indent + 'current directory size: %sM' % _dirsize)
        if ftype.lower() == 'files':
            _files_to_clean = self.get_cleanup_files_limit(
                directory, _dirsize, 0, 'file',
                keep_latest_one=files.get('keep_latest_one', []),
                remove_all=files.get('remove_all', [])
            )
            self.cleanup_raw(_files_to_clean, log_indent=log_indent)
        elif ftype.lower() == 'directories':
            _files_to_clean = self.get_cleanup_files_limit(
                directory, _dirsize, 0, 'dir',
                keep_latest_one=files.get('keep_latest_one', []),
                remove_all=files.get('remove_all', [])
            )
            self.cleanup_raw(_files_to_clean, log_indent=log_indent)

    def cleanup_raw(self, files_to_clean, log_indent=''):
        for f in files_to_clean:
            _realpath = os.path.realpath(f)
            for lm in self._allow_dir_access:
                if _realpath.startswith(lm):
                    break
            else:
                self.log(
                    log_indent + 'preventing removal of %s; must be a subdir of %s' % (
                        _realpath, self._allow_dir_access
                    )
                )
                continue

            self.log(log_indent + 'removing %s' % _realpath)
            if not self.remove_file(_realpath):
                self.log(log_indent + 'failed to remove %s' % _realpath)

    def get_cleanup_files_limit(self, directory, dirsize, limit_size, ftype, **cleanup):
        _files_to_clean = []
        _keep_latest_one = cleanup.get('keep_latest_one', [])
        _remove_all = cleanup.get('remove_all', [])
        _sort_oldest_first = cleanup.get('sort_oldest_first', [])

        _dirsize = dirsize

        for f in self.get_matching_files(directory, _remove_all, ftype=ftype):
            _files_to_clean.append(f.name)
            _dirsize -= f.size

        _klo = sorted(self.get_matching_files(directory, _keep_latest_one, ftype=ftype), key=lambda x: x.date)
        for f in _klo[:len(_klo) - 1]:
            _files_to_clean.append(f.name)
            _dirsize -= f.size

        if _dirsize > limit_size:
            for f in sorted(self.get_matching_files(directory, _sort_oldest_first, ftype=ftype), key=lambda x: x.date):
                _files_to_clean.append(f.name)
                _dirsize -= f.size
                if _dirsize < limit_size:
                    break

        return _files_to_clean

    def get_cleanup_files_chop(self, directory, dirsize, chop_size, ftype, **cleanup):
        _files_to_clean = []
        _remove_all = cleanup.get('remove_all', [])
        _sort_oldest_first = cleanup.get('sort_oldest_first', [])

        _chop_size = chop_size

        for f in self.get_matching_files(directory, _remove_all, ftype=ftype):
            _files_to_clean.append(f.name)
            _chop_size -= f.size

        if _chop_size <= 0:
            return _files_to_clean

        for f in sorted(self.get_matching_files(directory, _sort_oldest_first, ftype=ftype), key=lambda x: x.date):
            _files_to_clean.append(f.name)
            _chop_size -= f.size
            print(_chop_size)
            if _chop_size <= 0:
                break

        return _files_to_clean


if __name__ == '__main__':
    r = Rotate('/opt/dosis2/etc/dosis-system-cleanup-rules.json', debug=False)
    r.run()
