import argparse
import json
import os
from git import Repo # Requires: pip install GitPython

def update_keys_in_git_repo(
    repo_path,
    file_to_update,
    key_file_path,
    commit_message="Automated IAM key rotation update"
):
    """
    Updates a specified file in a local Git repository with new AWS IAM credentials,
    then commits and pushes the changes to the remote repository.

    Args:
        repo_path (str): The local path to the Git repository.
        file_to_update (str): The path to the file within the repository to update (e.g., 'config/aws_credentials.json').
        key_file_path (str): The path to the JSON file containing the new AWS IAM credentials.
        commit_message (str): The commit message for the Git commit.
    """
    print(f"Starting Git repository update for new IAM keys in '{repo_path}'...")

    # 1. Load new access keys from the provided file
    print("\n>>> Step 1: Loading new access keys...")
    try:
        with open(key_file_path, 'r') as f:
            new_keys = json.load(f)
        new_access_key_id = new_keys['AccessKeyId']
        new_secret_access_key = new_keys['SecretAccessKey']
        print(f"   Successfully loaded new keys for user: {new_keys['UserName']}")
    except FileNotFoundError:
        print(f"Error: Key file not found at '{key_file_path}'. Exiting.")
        return
    except json.JSONDecodeError:
        print(f"Error: Invalid JSON format in key file '{key_file_path}'. Exiting.")
        return
    except KeyError as e:
        print(f"Error: Missing expected key in JSON file: {e}. Exiting.")
        return

    # 2. Open the Git repository
    print("\n>>> Step 2: Opening Git repository...")
    try:
        repo = Repo(repo_path)
        if repo.is_dirty():
            print("Warning: Repository is dirty. Please commit or stash changes before running this script.")
            # return # Uncomment to exit if repo is dirty
        print(f"   Successfully opened repository at '{repo_path}'.")
    except Exception as e:
        print(f"Error opening Git repository: {e}. Ensure Git is installed and the path is correct.")
        return

    # 3. Update the specified file within the repository
    full_file_path = os.path.join(repo_path, file_to_update)
    print(f"\n>>> Step 3: Updating file '{full_file_path}' with new keys...")
    try:
        # Example: Update a JSON config file. Adapt this logic for other file formats.
        if os.path.exists(full_file_path):
            with open(full_file_path, 'r+') as f:
                config_data = json.load(f)
                config_data['aws_access_key_id'] = new_access_key_id
                config_data['aws_secret_access_key'] = new_secret_access_key
                f.seek(0)
                json.dump(config_data, f, indent=4)
                f.truncate()
            print(f"   Successfully updated '{file_to_update}'.")
        else:
            print(f"Error: File to update '{full_file_path}' not found in repository. Exiting.")
            return
    except json.JSONDecodeError:
        print(f"Error: Invalid JSON format in file to update '{full_file_path}'. Exiting.")
        return
    except Exception as e:
        print(f"Error updating file '{file_to_update}': {e}")
        return

    # 4. Commit and push changes
    print("\n>>> Step 4: Committing and pushing changes...")
    try:
        repo.index.add([file_to_update]) # Stage the changed file
        repo.index.commit(commit_message)
        print(f"   Committed changes with message: '{commit_message}'.")

        origin = repo.remotes.origin
        origin.push()
        print("   Successfully pushed changes to remote repository.")
    except Exception as e:
        print(f"Error committing or pushing changes: {e}. Ensure your Git credentials are configured.")
        return

    print("\nGit repository update completed successfully.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Update a file in a Git repository with new AWS IAM credentials and push changes.")
    parser.add_argument("--repo-path", required=True, help="The local path to the Git repository.")
    parser.add_argument("--file-to-update", required=True, help="The path to the file within the repository to update (e.g., 'config/aws_credentials.json').")
    parser.add_argument("--key-file-path", required=True, help="The path to the JSON file containing the new AWS IAM credentials.")
    parser.add_argument("--commit-message", default="Automated IAM key rotation update", help="The commit message for the Git commit.")

    args = parser.parse_args()

    update_keys_in_git_repo(
        repo_path=args.repo_path,
        file_to_update=args.file_to_update,
        key_file_path=args.key_file_path,
        commit_message=args.commit_message
    )