import os
import re

files_with_tests = [
    "src/lib.rs",
    "crates/core/src/middleware/rate_limit.rs",
    "crates/core/src/middleware/csrf.rs",
    "crates/core/src/middleware/cors.rs",
    "crates/core/src/middleware/body_limit.rs",
    "crates/core/src/hooks.rs",
    "crates/core/src/openapi.rs",
    "crates/core/src/email.rs",
    "crates/athena-auth-api/src/plugins/session_management.rs",
    "crates/athena-auth-api/src/plugins/password_management.rs",
    "crates/api/src/plugins/user_management/mod.rs",
    "crates/api/src/plugins/passkey/mod.rs",
    "crates/api/src/plugins/session_management.rs",
    "crates/api/src/plugins/password_management/mod.rs",
    "crates/api/src/plugins/admin/mod.rs",
    "crates/api/src/plugins/organization/rbac.rs",
    "crates/api/src/plugins/oauth/encryption.rs",
    "crates/api/src/plugins/mod.rs",
    "crates/api/src/plugins/email_verification/mod.rs",
    "crates/api/src/plugins/email_password.rs",
    "crates/api/src/plugins/api_key/mod.rs"
]

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

def main():
    for f in files_with_tests:
        file_path = os.path.join(base_dir, f)
        if not os.path.exists(file_path):
            continue
        
        with open(file_path, 'r', encoding='utf-8') as file:
            content = file.read()
            
        out_content = content
        matches = list(re.finditer(r'#\[cfg\(test\)\]\s*(?:#\[[^\]]*\]\s*)*mod\s+(\w+)\s*\{', out_content))
        if not matches:
            continue
            
        for match in reversed(matches):
            start_idx = match.start()
            mod_name = match.group(1)
            
            brace_start = match.end() - 1
            assert out_content[brace_start] == '{'
            
            brace_count = 1
            end_idx = brace_start + 1
            while brace_count > 0 and end_idx < len(out_content):
                if out_content[end_idx] == '{':
                    brace_count += 1
                elif out_content[end_idx] == '}':
                    brace_count -= 1
                end_idx += 1
                
            if brace_count == 0:
                test_block = out_content[brace_start+1:end_idx-1].strip()
                
                dir_name = os.path.dirname(file_path)
                base_name = os.path.basename(file_path)
                tests_dir = os.path.join(dir_name, "__tests__")
                os.makedirs(tests_dir, exist_ok=True)
                
                # if multiple tests in a file, we could have collisions if test_file_name = base_name, 
                # but usually there's only one. For `lib.rs` maybe multiple?
                test_file_name = base_name
                # if there is more than 1 test module in the same file, use the mod_name instead.
                if len(matches) > 1:
                    test_file_name = f"{base_name.replace('.rs', '')}_{mod_name}.rs"
                    
                test_file_path = os.path.join(tests_dir, test_file_name)
                
                with open(test_file_path, 'w', encoding='utf-8') as tf:
                    tf.write("use super::*;\n\n" + test_block + "\n")
                    
                test_mod_decl = f'#[cfg(test)]\n#[path = "__tests__/{test_file_name}"]\nmod {mod_name};\n'
                out_content = out_content[:start_idx] + test_mod_decl + out_content[end_idx:]

        with open(file_path, 'w', encoding='utf-8') as file:
            file.write(out_content)

if __name__ == '__main__':
    main()
