#!/usr/bin/env python3
import os
import re
import shutil

def main():
    """
    For each config file in ~/.ssh/conf, find lines matching:
      IdentityFile ~/.ssh/conf/<folderA>/keys/<folderB>/id_ed25519[.pub]?
    and replace them with:
      IdentityFile ~/.ssh/conf/<folderB>/id_ed25519[.pub]?

    Then physically move the key file from:
      ~/.ssh/conf/<folderA>/keys/<folderB>/
    to:
      ~/.ssh/conf/<folderB>/
    """

    conf_root = os.path.expanduser("~/.ssh/conf")
    # Regex capturing the path:
    #   group(1)  -> "IdentityFile ~/.ssh/conf/"
    #   group(2)  -> folderA
    #   group(3)  -> folderB
    #   group(4)  -> "id_ed25519" or "id_ed25519.pub"
    #   group(5)  -> "id_ed25519" (without .pub)
    #   group(6)  -> ".pub" or empty
    pattern = re.compile(r'(IdentityFile\s+~/.ssh/conf/)([^/]+)/keys/([^/]+)/((id_ed25519)(\.pub)?)')

    # Walk the ~/.ssh/conf directory tree
    for root, dirs, files in os.walk(conf_root):
        for f in files:
            # Only process "config" or "*.conf" files
            if f == "config" or f.endswith(".conf"):
                file_path = os.path.join(root, f)

                # Read the file
                with open(file_path, 'r') as fp:
                    lines = fp.readlines()

                new_lines = []
                changed = False

                # Check each line for the pattern
                for line in lines:
                    match = pattern.search(line)
                    if match:
                        # We found a line with the old IdentityFile path
                        new_line = pattern.sub(r'\1\3/\4', line)
                        changed = True

                        # Example of what we captured:
                        folderA = match.group(2)  # e.g. "foxdale"
                        folderB = match.group(3)  # e.g. "atlas"
                        keyfile = match.group(4)  # e.g. "id_ed25519" or "id_ed25519.pub"

                        # Build old/new file paths
                        old_file_path = os.path.join(conf_root, folderA, "keys", folderB, keyfile)
                        new_file_path = os.path.join(conf_root, folderB, keyfile)

                        print(f"Moving {old_file_path} => {new_file_path}")

                        # Ensure the destination folder exists
                        os.makedirs(os.path.dirname(new_file_path), exist_ok=True)

                        # Move the file if it exists
                        if os.path.isfile(old_file_path):
                            try:
                                shutil.move(old_file_path, new_file_path)
                            except Exception as e:
                                print(f"Error moving file: {e}")
                        else:
                            print(f"[WARNING] Key file not found: {old_file_path}")

                        line = new_line

                    new_lines.append(line)

                # If we changed anything, rewrite the config file
                if changed:
                    with open(file_path, 'w') as fp:
                        fp.writelines(new_lines)
                    print(f"Updated IdentityFile paths in {file_path}")

    print("Done fixing IdentityFile paths and moving key files.")

if __name__ == "__main__":
    main()