#!/usr/bin/python3
#
# Copyright (C) 2015 Matthias Klumpp <mak@debian.org>
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 3.0 of the License, or (at your option) any later version.
#
# 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
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this program.

import os
import sys
import yaml
import apt_pkg
import gzip
import tarfile
import glob
import shutil
import time
from jinja2 import Environment, FileSystemLoader
from kyotocabinet import *
from optparse import OptionParser
from dep11.multiprocessing import *
import multiprocessing as mp

from dep11.extractor import MetadataExtractor
from dep11.component import DEP11Component, DEP11YamlDumper, get_dep11_header
from dep11.iconfinder import ContentsListIconFinder
from dep11.utils import read_packages_dict_from_file

def debugln(msg):
    if os.environ.get("DEBUG"):
        print(msg)

def safe_move_file(old_fname, new_fname):
    if not os.path.isfile(old_fname):
        return
    if os.path.isfile(new_fname):
        os.remove(new_fname)
    os.rename(old_fname, new_fname)

def get_pkg_id(name, version, arch):
    return "%s_%s_%s" % (name, version, arch)

class MetadataPool:
    '''
    A pool of component metadata
    '''

    def __init__(self, db, hints_db):
        '''
        Initialize the metadata pool.
        '''
        self._mcpts = dict()
        self._db = db
        self._db_hints = hints_db

    def append_cptdata(self, pkgname, version, arch, cptlist):
        '''
        Makes a list of all the DEP11Component objects in a arch pool
        '''
        cpts = self._mcpts.get(arch)
        if not cpts:
            self._mcpts[arch] = list()
            cpts = self._mcpts[arch]
        for c in cptlist:
            cpts.append({'version': version, 'cpt': c, 'pkg': pkgname})
        if not cptlist:
            cpts.append({'version': version, 'cpt': None, 'pkg': pkgname})

    def save(self):
        """
        Saves metadata in db (serialized to YAML)
        """

        self._db.begin_transaction()
        self._db_hints.begin_transaction()

        for arch, items in self._mcpts.items():
            for item in items:
                cpt = item.get('cpt')
                version = item['version']
                idname = get_pkg_id(item['pkg'], version, arch)

                # if the package has no components, mark it as always-ignored
                if not cpt:
                    self._db_hints.set(idname, "ignore")
                    continue

                # get the metadata in YAML format
                metadata = cpt.to_yaml_doc()
                hints_yml = cpt.get_hints_yaml()
                if not hints_yml:
                    hints_yml = ""

                if not cpt.has_ignore_reason():
                    self._db.set(idname, metadata)
                self._db_hints.set(idname, hints_yml)

        self._db.end_transaction()
        self._db_hints.end_transaction()

def extract_metadata(mde, sn, pkgname, package_fname, version, arch):
    cpts = mde.process(pkgname, package_fname)

    data = dict()
    data['arch'] = arch
    data['cpts'] = list(cpts)
    data['version'] = version
    data['pkgname'] = pkgname
    data['message'] = "Processed package: %s (%s/%s)" % (pkgname, sn, arch)

    return (PROC_STATUS_SUCCESS, data)

class DEP11Generator:
    def __init__(self):
        pass

    def initialize(self, dep11_dir):
        conf_fname = os.path.join(dep11_dir, "dep11-config.yml")
        if not os.path.isfile(conf_fname):
            print("Could not find configuration! Make sure 'dep11-config.yml' exists!")
            return False

        f = open(conf_fname, 'r')
        conf = yaml.safe_load(f.read())
        f.close()

        if not conf:
            print("Configuration is empty!")
            return False

        if not conf.get("ArchiveRoot"):
            print("You need to specify an archive root path.")
            return False

        if not conf.get("Suites"):
            print("Config is missing information about suites!")
            return False

        if not conf.get("DataUrl"):
            print("You need to specify an URL where additional data (like screenshots) can be downloaded.")
            return False

        self._dep11_url = conf.get("DataUrl")
        self._icon_sizes = conf.get("IconSizes")
        if not self._icon_sizes:
            self._icon_sizes = ["128x128", "64x64"]

        self._archive_root = conf.get("ArchiveRoot")

        self._cache_dir = os.path.join(dep11_dir, "cache")
        if conf.get("CacheDir"):
            self._cache_dir = conf.get("CacheDir")

        self._export_dir = os.path.join(dep11_dir, "export")
        if conf.get("ExportDir"):
            self._export_dir = conf.get("ExportDir")

        if not os.path.exists(self._cache_dir):
            os.makedirs(self._cache_dir)
        if not os.path.exists(self._export_dir):
            os.makedirs(self._export_dir)

        self._suites_data = conf['Suites']

        self._db = DB()
        if not self._db.open(os.path.join(self._cache_dir, "data_cache.kch"), DB.OREADER | DB.OWRITER | DB.OCREATE):
            print("Could not open cache: %s" % (str(self._db.error())), file=sys.stderr)
            return False

        self._db_hints = DB()
        if not self._db_hints.open(os.path.join(self._cache_dir, "hints_cache.kch"), DB.OREADER | DB.OWRITER | DB.OCREATE):
            print("Could not open hints cache: %s" % (str(self._db_hints.error())), file=sys.stderr)
            return False

        os.chdir(dep11_dir)
        return True

    def _get_asset_dir(self):
        asset_dir = os.path.join(self._export_dir, "assets")
        if not os.path.exists(asset_dir):
            os.makedirs(asset_dir)
        return asset_dir

    def _get_packages_for(self, suite, component, arch):
        return read_packages_dict_from_file(self._archive_root, suite, component, arch).values()

    def make_icon_tar(self, suitename, component, pkglist):
        '''
         Generate icons-%(size).tar.gz
        '''
        dep11_assetdir = self._get_asset_dir()
        names_seen = set()
        tar_location = os.path.join(self._export_dir, "dep11", suitename, component)

        size_tars = dict()

        for pkg in pkglist:
            idname = get_pkg_id(pkg['name'], pkg['version'], pkg['arch'])

            for size in self._icon_sizes:
                icon_location_glob = os.path.join (dep11_assetdir, component, idname, "icons", size, "*.*")

                tar = None
                if size not in size_tars:
                    icon_tar_fname = os.path.join(tar_location, "icons-%s.tar.gz" % (size))
                    size_tars[size] = tarfile.open(icon_tar_fname+".new", "w:gz")
                tar = size_tars[size]

                for filename in glob.glob(icon_location_glob):
                    icon_name = os.path.basename(filename)
                    if icon_name in names_seen:
                        continue
                    tar.add(filename, arcname=icon_name)
                    names_seen.add(icon_name)

        for tar in size_tars.values():
            tar.close()
            # FIXME Ugly....
            safe_move_file(tar.name, tar.name.replace(".new", ""))

    def process_suite(self, suite_name):
        '''
        Extract new metadata for a given suite.
        '''

        suite = self._suites_data.get(suite_name)
        if not suite:
            print("Suite '%s' not found!" % (suite_name))
            return False

        dep11_assetdir = self._get_asset_dir()

        # We need 'forkserver' as startup method to prevent deadlocks on join()
        # Something in the extractor is doing weird things, makes joining impossible
        # when using simple fork as startup method.
        mp.set_start_method('forkserver')

        for component in suite['components']:
            all_cpt_pkgs = list()
            for arch in suite['architectures']:
                pkglist = self._get_packages_for(suite_name, component, arch)

                pool = ExtractorProcessPool(maxtasksperchild=16)
                dpool = MetadataPool(self._db, self._db_hints)

                def parse_results(message):
                    (code, msg) = message
                    if code != PROC_STATUS_SUCCESS:
                        print("######\nERROR:\n#\n%s" % (str(msg)))
                        pool.terminate()
                        pool.join()
                        sys.exit(5)
                    debugln(msg['message'])
                    dpool.append_cptdata(msg['pkgname'], msg['version'], msg['arch'], msg['cpts'])

                iconf = ContentsListIconFinder(suite_name, component, arch, self._archive_root)
                mde = MetadataExtractor(suite_name, component,
                                dep11_assetdir,
                                self._dep11_url,
                                self._icon_sizes,
                                iconf)

                for pkg in pkglist:
                    package_fname = os.path.join (self._archive_root, pkg['filename'])
                    idname = get_pkg_id(pkg['name'], pkg['version'], pkg['arch'])

                    # check if we scanned the package already
                    if self._db.get_str(idname) or self._db_hints.get_str(idname):
                        continue

                    if not os.path.exists(package_fname):
                        print('Package not found: %s' % (package_fname))
                        continue
                    pool.apply_async(extract_metadata,
                                (mde, suite_name, pkg['name'], package_fname, pkg['version'], pkg['arch']), callback=parse_results)
                pool.close()
                pool.join()

                # save new metadata to the database
                dpool.save()

                hints_dir = os.path.join(self._export_dir, "hints", suite_name, component)
                if not os.path.exists(hints_dir):
                    os.makedirs(hints_dir)
                dep11_dir = os.path.join(self._export_dir, "dep11", suite_name, component)
                if not os.path.exists(dep11_dir):
                    os.makedirs(dep11_dir)

                # now write data to disk
                hints_fname = os.path.join(hints_dir, "DEP11Hints_%s.yml.gz" % (arch))
                data_fname = os.path.join(dep11_dir, "Components-%s.yml.gz" % (arch))

                hints_f = gzip.open(hints_fname+".new", 'wb')
                data_f = gzip.open(data_fname+".new", 'wb')

                dep11_header = get_dep11_header(suite_name, component)
                data_f.write(bytes(dep11_header, 'utf-8'))

                for pkg in pkglist:
                    idname = get_pkg_id(pkg['name'], pkg['version'], pkg['arch'])
                    data = self._db.get_str(idname)
                    if data:
                        data_f.write(bytes(data, 'utf-8'))
                    hint = self._db_hints.get_str(idname)
                    if hint and hint != "ignore":
                        hints_f.write(bytes(hint, 'utf-8'))

                data_f.close()
                safe_move_file(data_fname+".new", data_fname)

                hints_f.close()
                safe_move_file(hints_fname+".new", hints_fname)

                all_cpt_pkgs.extend(pkglist)

            # create icon tarball
            self.make_icon_tar(suite_name, component, all_cpt_pkgs)

            print("Completed metadata extraction for suite %s/%s" % (suite_name, component))

    def expire_cache(self):
        dep11_assetdir = self._get_asset_dir()

        pkgids = set()
        for suite_name in self._suites_data:
            suite = self._suites_data[suite_name]
            for component in suite['components']:
                for arch in suite['architectures']:
                    pkglist = self._get_packages_for(suite_name, component, arch)
                    for pkg in pkglist:
                        idname = get_pkg_id(pkg['name'], pkg['version'], pkg['arch'])
                        pkgids.add(idname)

        # clean caches and database
        for pkid in self._db:
            if pkid in pkgids:
                continue
            self._db.remove(pkid)
            dirs = glob.glob(os.path.join(dep11_assetdir, "*", pkid))
            if dirs:
                shutil.rmtree(dirs[0])

        for pkid in self._db_hints:
            if pkid in pkgids:
                continue
            self._db_hints.remove(pkid)
            dirs = glob.glob(os.path.join(dep11_assetdir, "*", pkid))
            if dirs:
                shutil.rmtree(dirs[0])

    def render_template(self, name, out_dir, out_name = None, *args, **kwargs):
        if not out_name:
            out_path = os.path.join(out_dir, name)
        else:
            out_path = os.path.join(out_dir, out_name)
        # create subdirectories if necessary
        out_dir = os.path.dirname(os.path.realpath(out_path))
        if not os.path.exists(out_dir):
            os.makedirs(out_dir)

        template_dir = os.path.dirname(os.path.realpath(__file__))
        template_dir = os.path.realpath(os.path.join(template_dir, "..", "templates", "default"))
        j2_env = Environment(loader=FileSystemLoader(template_dir))

        template = j2_env.get_template(name)
        content = template.render(*args, **kwargs)
        debugln(out_path)
        with open(out_path, 'w') as f:
            f.write(content)

    def update_html(self):
        dep11_hintsdir = os.path.join(self._export_dir, "hints")
        if not os.path.exists(dep11_hintsdir):
            return

        # FIXME: Properly account for architecture, and remove old HTML files
        for suite_name in self._suites_data:
            suite = self._suites_data[suite_name]
            for component in suite['components']:
                for arch in suite['architectures']:
                    package_summaries = list()
                    export_dir = os.path.join(self._export_dir, "hints_html", suite_name, component)
                    h_fname = os.path.join(dep11_hintsdir, suite_name, component, "DEP11Hints_%s.yml.gz" % (arch))
                    if not os.path.isfile(h_fname):
                        continue
                    f = gzip.open(h_fname, 'r')
                    hints_data = yaml.safe_load_all(f.read())
                    f.close()
                    if not hints_data:
                        continue

                    for hint in hints_data:
                        pkg_name = hint['Package']
                        errors = hint['Hints'].get('errors', list())
                        warnings = hint['Hints'].get('warnings', list())
                        infos = hint['Hints'].get('infos', list())
                        for i in range(0, len(infos)):
                            infos[i] = infos[i].replace("http://freedesktop.org/software/appstream/docs/chap-Quickstart.html",
                                            "<a href=\"http://freedesktop.org/software/appstream/docs/chap-Quickstart.html\">the XML quickstart guide</a>")

                        package_summaries.append({'pkgname': pkg_name, 'errors_count': len(errors), 'warnings_count': len(warnings), 'infos_count': len(infos)})

                        self.render_template("issues_page.html", export_dir, "%s.html" % (pkg_name),
                                package_name=pkg_name, errors=errors, warnings=warnings, infos=infos, suite=suite_name, time=time.strftime("%c"))

                    # Now render our index page
                    self.render_template("issues_index.html", export_dir, "index.html",
                                package_summaries=package_summaries, suite=suite_name, time=time.strftime("%c"))

def main():
    parser = OptionParser()

    (options, args) = parser.parse_args()
    if len(args) == 0:
        print("You need to specify a command!")
        sys.exit(1)

    command = args[0]

    if command == "process":
        if len(args) != 3:
            print("You need to specify a DEP-11 data dir and a suite.")
            sys.exit(1)
        gen = DEP11Generator()
        ret = gen.initialize(args[1])
        if not ret:
            print("Initialization failed, can not continue.")
            sys.exit(2)

        gen.process_suite(args[2])
    elif command == "cleanup":
        if len(args) != 2:
            print("You need to specify a DEP-11 data dir.")
            sys.exit(1)
        gen = DEP11Generator()
        ret = gen.initialize(args[1])
        if not ret:
            print("Initialization failed, can not continue.")
            sys.exit(2)

        gen.expire_cache()
    elif command == "update_html":
        if len(args) != 2:
            print("You need to specify a DEP-11 data dir.")
            sys.exit(1)
        gen = DEP11Generator()
        ret = gen.initialize(args[1])
        if not ret:
            print("Initialization failed, can not continue.")
            sys.exit(2)

        gen.update_html()
    else:
        print("Run with --help for a list of available command-line options!")

if __name__ == "__main__":
    apt_pkg.init()
    main()
