import os

base_dir = r"c:\Users\floris\Documents\GitHub\athena-auth"

def fix_file(path):
    with open(path, 'r', encoding='utf-8') as f:
        content = f.read()

    # The original script prepended "use super::*;\n\n"
    # If the second or third line is another "use super::*;", it means it's a duplicate.
    if content.startswith("use super::*;\n\nuse super::*;"):
        content = content[len("use super::*;\n\n"):]
        with open(path, 'w', encoding='utf-8') as f:
            f.write(content)
        print(f"Fixed double-import in {path}")
    elif content.startswith("use super::*;\n\n") and content.count("use super::*;") > 1:
        # It has another use super::*; somewhere, probably soon after.
        # Let's just remove the first 15 bytes.
        content = content[len("use super::*;\n\n"):]
        with open(path, 'w', encoding='utf-8') as f:
            f.write(content)
        print(f"Fixed duplicate import in {path}")
    elif content.startswith("use super::*;\n"):
        # We can also check if rustc is still complaining about unused import
        # we can just allow it for these generated test files.
        pass

def allow_unused_imports(path):
    # To be extremely safe, we can add #![allow(unused_imports)] to the top of all __tests__/*.rs files
    with open(path, 'r', encoding='utf-8') as f:
        content = f.read()
    if not content.startswith("#![allow(unused_imports)]"):
        content = "#![allow(unused_imports)]\n" + content
        with open(path, 'w', encoding='utf-8') as f:
            f.write(content)
        print(f"Added #![allow(unused_imports)] to {path}")

for root, dirs, files in os.walk(base_dir):
    if "__tests__" in root:
        for file in files:
            if file.endswith(".rs"):
                path = os.path.join(root, file)
                fix_file(path)
                allow_unused_imports(path)
