#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright (C) 2015 Canonical
#
# Authors:
#  Didier Roche
#
# This program is free software; you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation; version 3.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA

import apt
import argparse
from contextlib import suppress
import logging
import os
import re
import sys

logger = logging.getLogger(__name__)


def extract_package_list(seed_file):
    '''Extract root package list (which can pull new deps) from the seed'''
    # we only treat one subtitution per line (the seed doesn't handle more)
    declared_pkgs = []
    substitution_re = re.compile("\${(.*)}")
    substitutions = {}
    with open(seed_file, 'r') as f:
        for line in f:
            line = line.strip()
            if not line.startswith("*"):
                continue
            line = line[1:].strip()
            # new subtitution definition
            with suppress(ValueError):
                key, values = line.split(':')
                substitutions[key] = [value.strip() for value in values.split()]
                continue
            pkg_name = line
            # we only treat one subtitution per line (the seed doesn't handle more)
            res = substitution_re.search(pkg_name)
            with suppress(AttributeError):
                key = res.group(1)
                try:
                    for value in substitutions[key]:
                        declared_pkgs.append(pkg_name.replace("${" + key + "}", value))
                except KeyError:
                    logger.error("{} hasn't been declared before {}. Exiting".format(key, pkg_name))
                    sys.exit(1)
                continue
            # last: package name with no substitution
            declared_pkgs.append(pkg_name)

    return declared_pkgs


def filter_packages(package_list, apt_cache):
    '''Filter input package list by excluding already installed package and those in universe'''
    filtered_list = []
    for pkg_name in package_list:
        if pkg_name not in apt_cache:
            logger.warning("{} isn't available in the archive".format(pkg_name))
            continue
        pkg = apt_cache[pkg_name]
        # discare installed packages
        if pkg.is_installed:
            logger.debug("{} is already installed".format(pkg_name))
            continue
        # discare if package only available in universe
        in_universe = True
        for origin in pkg.candidate.origins:
            if origin.origin == 'Ubuntu' and origin.component == 'main':
                in_universe = False
        if in_universe:
            logger.info("{} is in universe, discare".format(pkg_name))
            continue
        filtered_list.append(pkg_name)
    return filtered_list


def compute_install_size(package_list, apt_cache):
    '''Return additional install size from the package_list, including deps'''
    for pkg_name in package_list:
        pkg = apt_cache[pkg_name]
        pkg.mark_install()
    return apt_cache.required_download


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Count the compressed diff size between current installation and " +
                                                 "what the seed file requires")
    parser.add_argument('file', help="Path to the file in a seed format")
    parser.add_argument("-d", "--debug", action="store_true", help="Add debug infos")
    args = parser.parse_args()

    log_level = logging.INFO
    if args.debug:
        log_level = logging.DEBUG
    logging.basicConfig(level=log_level, format="%(levelname)s: %(message)s")

    seed_file = args.file
    if not os.path.isfile(seed_file):
        print("{} doesn't exit".format(seed_file))
        sys.exit(1)
    all_pkgs = extract_package_list(seed_file)

    cache = apt.Cache()
    filtered_pkgs = filter_packages(all_pkgs, cache)

    print("Packages to install:")
    print(", ".join(filtered_pkgs))

    print("Estimated additional compressed size on live: {:.2f} Mb".format(compute_install_size(filtered_pkgs, cache) /
                                                                           (1024 * 1024)))
    cache.clear()
