#!/usr/bin/python3.6

import os
import createrepo_c as cr
import shutil
from argparse import ArgumentParser
from pyfaf.storage import getDatabase
from pyfaf.queries import get_packages_by_osrelease

faf_names = {'rhel': "Red Hat Enterprise Linux",
             'fedora': 'Fedora',
             'centos': 'CentOS'}


def get_pkglist(db, opsys, release, arch):
    if opsys in faf_names.keys():
        opsys = faf_names[opsys]
    q = get_packages_by_osrelease(db, opsys, release, arch)
    for pkg in q:
        if pkg.has_lob("package"):
            print("Adding package: %s-%s-%s" % (pkg.name, pkg.build.version, pkg.build.release))
            yield os.path.abspath(pkg.get_lob_path("package"))


def generate_repo(db, outputdir, opsys, release, arch):

    repodata_path = os.path.join(outputdir, "repodata")

    if os.path.exists(repodata_path):
        shutil.rmtree(repodata_path)

    os.makedirs(repodata_path)

    # Prepare metadata files
    repomd_path = os.path.join(repodata_path, "repomd.xml")
    pri_xml_path = os.path.join(repodata_path, "primary.xml.gz")
    fil_xml_path = os.path.join(repodata_path, "filelists.xml.gz")
    oth_xml_path = os.path.join(repodata_path, "other.xml.gz")
    pri_db_path = os.path.join(repodata_path, "primary.sqlite")
    fil_db_path = os.path.join(repodata_path, "filelists.sqlite")
    oth_db_path = os.path.join(repodata_path, "other.sqlite")

    pri_xml = cr.PrimaryXmlFile(pri_xml_path)
    fil_xml = cr.FilelistsXmlFile(fil_xml_path)
    oth_xml = cr.OtherXmlFile(oth_xml_path)
    pri_db = cr.PrimarySqlite(pri_db_path)
    fil_db = cr.FilelistsSqlite(fil_db_path)
    oth_db = cr.OtherSqlite(oth_db_path)

    # Prepare list of packages to process
    pkg_list = list(get_pkglist(db, opsys, release, arch))

    pkg_list_len = len(pkg_list)
    pri_xml.set_num_of_pkgs(pkg_list_len)
    fil_xml.set_num_of_pkgs(pkg_list_len)
    oth_xml.set_num_of_pkgs(pkg_list_len)

    # Process all packages
    for filename in pkg_list:
        pkg = cr.package_from_rpm(filename)
        pkg.location_href = os.path.relpath(filename, outputdir)
        pri_xml.add_pkg(pkg)
        fil_xml.add_pkg(pkg)
        oth_xml.add_pkg(pkg)
        pri_db.add_pkg(pkg)
        fil_db.add_pkg(pkg)
        oth_db.add_pkg(pkg)

    pri_xml.close()
    fil_xml.close()
    oth_xml.close()

    # Prepare repomd.xml
    repomd = cr.Repomd()

    # Add records into the repomd.xml
    repomdrecords = (("primary",      pri_xml_path, pri_db),
                     ("filelists",    fil_xml_path, fil_db),
                     ("other",        oth_xml_path, oth_db),
                     ("primary_db",   pri_db_path,  None),
                     ("filelists_db", fil_db_path,  None),
                     ("other_db", oth_db_path, None))

    for name, path, db_to_update in repomdrecords:
        # Compress sqlite files with bzip2
        if path.endswith('.sqlite'):
            new_path = '%s.bz2' % path
            cr.compress_file(path, new_path, cr.BZ2)
            os.unlink(path)
            path = new_path

        record = cr.RepomdRecord(name, path)
        record.fill(cr.SHA256)
        record.rename_file()
        if (db_to_update):
            db_to_update.dbinfo_update(record.checksum)
            db_to_update.close()
        repomd.set_record(record)

    # Write repomd.xml
    open(repomd_path, "w").write(repomd.xml_dump())


if __name__ == "__main__":
    parser = ArgumentParser(description="Generates dnf repo from FAF databse")
    parser.add_argument("OPSYS")
    parser.add_argument("RELEASE")
    parser.add_argument("ARCHITECTURE")
    parser.add_argument("--outputdir", default=os.getcwd())
    args = parser.parse_args()

    print("Creating repo in '%s'" % (args.outputdir))

    generate_repo(getDatabase(), args.outputdir, args.OPSYS, args.RELEASE, args.ARCHITECTURE)
