def fix_imports(file_path):
    """This function fixes the imports in the files generated by ANTLR.
    Some files have imports that check for older python versions which
    causes issues with the linter. This function removes the if/else check
    and only keeps the import statement that works for our python version.
    """
    with open(file_path) as file:
        lines = file.readlines()
    updated_lines = []
    i = 0
    while i < len(lines):
        line = lines[i]
        if "if sys.version_info[1] > 5:" in line or 'if "." in __name__:' in line:
            updated_lines.append(lines[i + 1].strip())
            i += 4
        else:
            updated_lines.append(line)
            i += 1
    with open(file_path, "w") as file:
        file.writelines(updated_lines)


def add_lines_to_start(file_path, lines_to_add):
    """This function adds lines to the start of a file. This is used to
    add `flake8: noqa` and `type: ignore` to the files generated by ANTLR
    to ignore linter errors caused by ANTLR.
    """
    with open(file_path) as file:
        original_content = file.readlines()

    updated_content = lines_to_add + original_content

    with open(file_path, "w") as file:
        file.writelines(updated_content)


files = [
    "generated/AssetSelectionLexer.py",
    "generated/AssetSelectionListener.py",
    "generated/AssetSelectionParser.py",
    "generated/AssetSelectionVisitor.py",
]

for file in files:
    fix_imports(file)
    add_lines_to_start(
        file,
        [
            "# flake8: noqa\n",
            "# type: ignore\n",
        ],
    )
