#!/usr/bin/python3

import ruamel.yaml
import pathlib
import argparse
import sys

yaml = ruamel.yaml.YAML()

def get_chart_version(chart_name: str) -> str:
    with open(f"./{chart_name}-chart/Chart.yaml") as f:
        chart = yaml.load(f)
        return chart["version"]

def get_charts(chart):
    if not chart["chart"].startswith("%%CHART_REPO%%"):
        # Not a locally managed chart
        return {}

    chart_name = chart["chart"][len("%%CHART_REPO%%/%%CHART_PREFIX%%"):]
    charts = { chart_name: chart["version"] }
    for child_chart in chart.get("dependencyCharts", []) + chart.get("addonCharts", []):
        charts.update(get_charts(child_chart))
    return charts

def get_charts_list():
    with open("./release-manifest-image/release_manifest.yaml") as f:
        manifest = yaml.load(f)
    charts = {}
    for chart in manifest["spec"]["components"]["workloads"]["helm"]:
        charts.update(get_charts(chart))
    return charts

def check_charts(fix: bool) -> bool:
    success = True
    charts = get_charts_list()
    to_fix = {}
    for chart in charts:
        expected_version = get_chart_version(chart)
        if expected_version != charts[chart]:
            success = False
            to_fix[f'%%CHART_REPO%%/%%CHART_PREFIX%%{chart}'] = expected_version
            print(f"{chart}: Expected: {expected_version}, Got: {charts[chart]}")
    if fix and not success:
        fix_charts(to_fix)
        return True
    return success

def fix_charts(to_fix):
    manifest_path = pathlib.Path("./release-manifest-image/release_manifest.yaml")
    manifest = yaml.load(manifest_path)
    yaml.indent(mapping=2, sequence=4, offset=2)
    yaml.width = 4096
    for chart_index, chart in enumerate(manifest["spec"]["components"]["workloads"]["helm"]):
        changed = False
        if chart["chart"] in to_fix.keys():
            changed = True
            chart["version"] = to_fix[chart["chart"]]
        for subchart_index, subchart in enumerate(chart.get("addonCharts", [])):
            if subchart["chart"] in to_fix.keys():
                changed = True
                subchart["version"] = to_fix[subchart["chart"]]
                chart["addonCharts"][subchart_index] = subchart
        for subchart_index, subchart in enumerate(chart.get("dependencyCharts", [])):
            if subchart["chart"] in to_fix.keys():
                changed = True
                subchart["version"] = to_fix[subchart["chart"]]
                chart["dependencyCharts"][subchart_index] = subchart
        if changed:
            manifest["spec"]["components"]["workloads"]["helm"][chart_index] = chart
    yaml.dump(manifest, manifest_path)

def main():
    print("Checking charts versions in release manifest")
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--check', action='store_true')
    args = parser.parse_args()
    if not check_charts(not args.check):
        sys.exit(1)
    else:
        print("All local charts in release manifest are using the right version")

if __name__ == "__main__":
    main()