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

from __future__ import print_function

import os
import sys
from subprocess import check_output, call
import time
import argparse


def str_to_tuple(str):
    """return tuple from all numbers in a string"""
    tmp = ""
    for ch in str:
        tmp += ch if ch.isdigit() else ' '
    return tuple([int(x) for x in tmp.split()])

class NoKernelInstalledError(Exception):
    pass

class ActiveNotInstalledError(Exception):
    pass

class RemoveOldKernels(object):
    def __init__(self, keep=0, verbose=0):
        self._keep = keep
        self._verbose = verbose

    def run(self):
        active_kernel_version = str_to_tuple(check_output(['uname', '-r']))
        if self._verbose > 0:
            print('active_kernel', active_kernel_version)
        pkg_nums, num_pkg_list_map = self.gen_list_map(
            ['dpkg', '-l', 'linux-*'], active_kernel_version, int(self._keep))
        if not pkg_nums:
            print('nothing to uninstall')
            return 1
        cmd = ['sudo', 'apt-get', 'remove', '--purge']
        cmd.extend(self.gather_packages(pkg_nums, num_pkg_list_map))
        if self._verbose > -1:
            print(cmd)
        # do not use check_output or the confirmation prompt disappears
        return call(cmd)

    @staticmethod
    def gather_packages(pkg_nums, num_pkg_list_map):
        ret_val = []
        for pkg_num in pkg_nums:
            ret_val.extend(num_pkg_list_map[pkg_num])
        return ret_val

    @staticmethod
    def gen_list_map(cmd, active_kernel_version, extra_keep=0, verbose=0):
        if isinstance(cmd, list):
            data = check_output(cmd)
        else:
            data = cmd
        num_pkg_list_map = {}
        for line in data.splitlines():
            if not line.startswith('ii'):
                continue
            pkg = line.split()[1]
            if not pkg.startswith('linux-'):
                continue
            num = str_to_tuple(pkg)
            if not num:  # empty tuple
                continue
            pkg_list = num_pkg_list_map.setdefault(num, [])
            pkg_list.append(pkg)
        # remove those installed packages without numbers
        pkg_nums = sorted(num_pkg_list_map)
        if verbose > 0:
            print('all pkg nums', pkg_nums, extra_keep)
        # if the last kernel is not the active one, keep regardless of --keep
        if not pkg_nums:
            raise NoKernelInstalledError()
        if pkg_nums[-1] != active_kernel_version:
            if verbose > 0:
                print('will not delete latest kernel {}, it is the not active'
                      ' kernel'.format(pkg_nums[-1]))
            extra_keep -= 1
            pkg_nums = pkg_nums[:-1]
        try:
            pkg_nums.remove(active_kernel_version)
        except ValueError:
            raise ActiveNotInstalledError
        if extra_keep > 0:  # shorten the list
            print('---------------------- shortening')
            pkg_nums = pkg_nums[:-(extra_keep)]
        if verbose > 0:
            print('pl, nplm', num_pkg_list_map, pkg_nums, extra_keep)
        return pkg_nums, num_pkg_list_map


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-v', '--verbose', action='count', default=0,
                        help='generate more output')
    parser.add_argument(
        '-k', '--keep', default=0,
        help="""how many extra kernels to keep apart from active one.
        by default only current and the latest kernel are kept
        (might be the same)"""
    )
    args = parser.parse_args()
    p = RemoveOldKernels(keep=args.keep, verbose=args.verbose)
    sys.exit(p.run())
