#!/usr/bin/python3 -u

import argparse
import json
import os
import shutil
import sys

import createrepo_c as cr

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from cosalib import cmdlib
from cosalib.builds import Builds
from cosalib.meta import GenericBuildMeta


def main():
    args = parse_args()

    workdir = os.path.abspath(os.getcwd())

    builds = Builds()
    builddir = builds.get_build_dir(args.build)
    buildmeta = GenericBuildMeta(workdir=workdir, build=args.build)

    if 'extensions' in buildmeta and not args.force:
        print(f"Extensions already exist: {buildmeta['extensions']['path']}")
        print("Use --force to force a rebuild")
        return

    init_config = 'src/config.json'
    if os.path.exists(init_config):
        with open(init_config, encoding='utf-8') as f:
            init_cfg = json.loads(f.read())
            variant = init_cfg["coreos-assembler.config-variant"]
            treefile_src = f"src/config/manifest-{variant}.yaml"
            extensions_src = f"src/config/extensions-{variant}.yaml"
    else:
        treefile_src = 'src/config/manifest.yaml'
        extensions_src = 'src/config/extensions.yaml'

    if not os.path.exists(extensions_src):
        raise Exception(f"Missing {extensions_src}")

    commit = buildmeta['ostree-commit']
    cmdlib.import_ostree_commit(workdir, builddir, buildmeta)

    tmpworkdir = prepare_tmpworkdir()
    changed = run_rpmostree(tmpworkdir, commit, treefile_src, extensions_src)
    if not changed:
        # For now, rpm-ostree will always detect a change because we don't seed
        # state from the previous build, so we won't hit this. Need to rework
        # how change detection is wired in `cmd-build` to do this properly.
        return

    outputdir = f"{tmpworkdir}/output"
    with open(f'{outputdir}/.rpm-ostree-state-chksum', encoding='utf-8') as f:
        rpm_ostree_state_chksum = f.read()

    pkglist = create_yumrepo(outputdir)
    extensions_tarball = create_tarball(buildmeta, outputdir, tmpworkdir)
    extensions_tarball_base = os.path.basename(extensions_tarball)

    buildmeta['extensions'] = {
        "path": extensions_tarball_base,
        "sha256": cmdlib.sha256sum_file(extensions_tarball),
        "rpm-ostree-state": rpm_ostree_state_chksum,
        "manifest": pkglist,
    }

    cmdlib.rm_allow_noent(f'{builddir}/{extensions_tarball_base}')
    shutil.move(extensions_tarball, builddir)
    buildmeta.write(artifact_name='extensions')

    shutil.rmtree(tmpworkdir)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--build", help="Build ID", default='latest')
    parser.add_argument("--force", help="Force rebuild", action='store_true')
    return parser.parse_args()


def prepare_tmpworkdir():
    tmpworkdir = 'tmp/extensions'
    if os.path.exists(tmpworkdir):
        shutil.rmtree(tmpworkdir)
    os.mkdir(tmpworkdir)
    configdir = 'src/config'
    for f in os.listdir(configdir):
        if os.path.isfile(f"{configdir}/{f}") and f.endswith('.repo'):
            shutil.copyfile(f"{configdir}/{f}", f"{tmpworkdir}/{f}")
    yumreposdir = 'src/yumrepos'
    if os.path.exists(yumreposdir):
        for f in os.listdir(yumreposdir):
            if os.path.isfile(f"{yumreposdir}/{f}") and f.endswith('.repo'):
                shutil.copyfile(f"{yumreposdir}/{f}", f"{tmpworkdir}/{f}")
    return tmpworkdir


def run_rpmostree(workdir, commit, treefile, extensions):
    cmdlib.cmdlib_sh(f'''
        cat > "{workdir}/manifest-override.yaml" <<EOF
include: ../../{treefile}
EOF
        changed_stamp={workdir}/changed
        runcompose_extensions {workdir}/output {workdir}/manifest-override.yaml {extensions} \
            --base-rev {commit}''')
    return os.path.exists(f'{workdir}/changed')


def create_yumrepo(repodir):
    cmdlib.runcmd(['createrepo_c', repodir])
    # we could also have rpm-ostree output the pkglist for us, but meh... we
    # need to run createrepo_c anyway and it's nice that we're using it as the
    # source of truth, since that's what rpm-ostree clients will also use
    repomd = cr.Repomd(os.path.join(repodir, "repodata/repomd.xml"))
    pkglist = {}

    def cb(pkg):
        epoch = ''
        if pkg.epoch and int(pkg.epoch) > 0:
            epoch = f'{pkg.epoch}:'
        pkglist[pkg.name] = f'{epoch}{pkg.version}-{pkg.release}.{pkg.arch}'

    for record in repomd.records:
        if record.type == 'primary':
            primary_xml = os.path.join(repodir, record.location_href)
            cr.xml_parse_primary(primary_xml, do_files=False, pkgcb=cb)
            break

    if len(pkglist) == 0:
        raise Exception("No RPMs found in output dir")
    return pkglist


def create_tarball(buildmeta, srcdir, destdir):
    destdir = os.path.abspath(destdir)
    basearch = buildmeta['coreos-assembler.basearch']
    tarfile = f'{destdir}/{buildmeta["name"]}-{buildmeta["buildid"]}-extensions.{basearch}.tar'
    cmdlib.runcmd(['tar', '-cf', tarfile, '.'], cwd=srcdir)
    return tarfile


if __name__ == '__main__':
    sys.exit(main())
