#!usr/bin/env python
# -*- coding: utf-8 -*-
import argparse, os, fileinput


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        prog="GANDLF_UpdateVersion",
        formatter_class=argparse.RawTextHelpFormatter,
        description="Update versions when creating a new release of GaNDLF, also useful when updating the version for development.\n\n",
    )
    parser.add_argument(
        "-ov",
        "--old_version",
        metavar="",
        type=str,
        required=True,
        help="The old version number",
    )
    parser.add_argument(
        "-nv",
        "--new_version",
        metavar="",
        type=str,
        required=True,
        help="The new version number",
    )

    args = parser.parse_args()

    def in_place_string_replace(
        filename: str, old_string: str, new_string: str
    ) -> None:
        """
        Replace a string in a file in place.

        Args:
            filename (str): The file to replace the string in
            old_string (str): The string to replace
            new_string (str): The string to replace with
        """
        if os.path.exists(filename):
            with fileinput.FileInput(filename, inplace=True) as file:
                for line in file:
                    print(line.replace(old_string, new_string), end="")

    cwd = os.getcwd()
    in_place_string_replace(
        os.path.join(cwd, "GANDLF/version.py"),
        args.old_version,
        args.new_version,
    )

    # find all yaml files in samples and testing directories
    folders_to_iterate = [
        os.path.join(cwd, "samples"),
        os.path.join(cwd, "testing"),
    ]

    files_where_version_is_stored = [
        os.path.join(cwd, "mlcube/model_mlcube/workspace/config.yml"),
        os.path.join(cwd, "tutorials/classification_medmnist_notebook/config.yaml"),
    ]

    for folder in folders_to_iterate:
        if os.path.isdir(folder):
            files_in_dir = os.listdir(folder)
            for file in files_in_dir:
                if file.endswith(".yaml") or file.endswith(".yml"):
                    files_where_version_is_stored.append(os.path.join(folder, file))

    args.old_version = args.old_version.replace("-dev", "")
    args.new_version = args.new_version.replace("-dev", "")

    # update the version.py file
    for filename in files_where_version_is_stored:
        in_place_string_replace(filename, args.old_version, args.new_version)

    print("Version updated successfully in `version.py` and all configuration files!")
