#!/usr/bin/python3.6

from argparse import ArgumentParser
from os.path import relpath
from pathlib import Path
import shutil
from typing import Generator

import createrepo_c as cr
from pyfaf.queries import get_packages_by_osrelease
from pyfaf.storage import Database, getDatabase

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


def get_pkglist(db: Database, opsys: str, release: str, arch: str) -> Generator[Path, None, None]:
    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 Path(pkg.get_lob_path("package")).resolve()


def generate_repo(db: Database, outputdir: Path, opsys: str, release: str, arch: str) -> None:
    repodata_path = outputdir / "repodata"

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

    repodata_path.mkdir(parents=True)

    # Prepare metadata files
    repomd_path = repodata_path / "repomd.xml"
    pri_xml_path = repodata_path / "primary.xml.gz"
    fil_xml_path = repodata_path / "filelists.xml.gz"
    oth_xml_path = repodata_path / "other.xml.gz"
    pri_db_path = repodata_path / "primary.sqlite"
    fil_db_path = repodata_path / "filelists.sqlite"
    oth_db_path = 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 = 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
    #pylint: disable=bad-whitespace
    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.suffix == '.sqlite':
            new_path = '%s.bz2' % path
            cr.compress_file(path, new_path, cr.BZ2)
            path.unlink()
            path = Path(new_path)

        record = cr.RepomdRecord(name, str(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())


def main() -> None:
    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=Path.cwd())
    args = parser.parse_args()

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

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

if __name__ == "__main__":
    main()
