#!/usr/bin/python3
# vi:shiftwidth=4:expandtab
#
# Create a new .treeinfo file from one produced by lorax
#
# Usage: python3 $0 .treeinfo > treeinfo-new
#   Variants are searched as */Packages/ where .treeinfo resides
#
# Prerequiste package: python3-productmd
#

import productmd
import glob
import os
import sys

t = productmd.TreeInfo()
t.load(sys.argv[1])

# Add missing [media] stanza
t.media.discnum=1
t.media.totaldiscs=1

# Clean existing "CentOS" "." variant
for i in t.variants:
    del t.variants[i]

# Get directory of .treeinfo, to search for */Packages/
treedir = os.path.dirname(sys.argv[1])

# Search for */Packages/ and add them as variants
ocwd = os.getcwd()
os.chdir(treedir)
for p in glob.glob("*/Packages"):
    if not os.path.isdir(p):
        continue
    id = os.path.basename(os.path.dirname(p))   # * part
    v = productmd.treeinfo.Variant(t)
    v.id = id
    v.uid = id
    v.name = id
    v.type = "variant"
    v.paths = productmd.treeinfo.VariantPaths(v)
    v.paths.packages = p
    v.paths.repository = id

    t.variants.add(v)

os.chdir(ocwd)

# Output to stdout
t.dump("/dev/fd/1")