import subprocess
import yaml
import os
import argparse

# Set up argument parsing
parser = argparse.ArgumentParser(
    description="Checkout modules from target branch that are not in source branch."
)
parser.add_argument("yaml_file", help="Path to the YAML file")
parser.add_argument("source_branch", help="The source branch")
parser.add_argument("target_branch", help="The target branch")
parser.add_argument("root_repo", help="Path to the root repository")

# Parse the arguments
args = parser.parse_args()

yaml_file = args.yaml_file
source_branch = args.source_branch
target_branch = args.target_branch
root_repo = args.root_repo

# Change to the repository directory
os.chdir(root_repo)
# Read YAML file
with open(yaml_file, "r") as file:
    data = yaml.safe_load(file)

# Extract module lists for source and target branches
modules_source = data["branches"].get(source_branch, {}).get("modules", [])
modules_target = data["branches"].get(target_branch, {}).get("modules", [])

# Ensure the latest changes are fetched
subprocess.run(["git", "fetch", "origin"], check=True)

# Checkout source branch first
print(f"Checking out source branch: {source_branch}")
subprocess.run(["git", "checkout", source_branch], check=True)

# Checkout modules in target_branch that are not in source_branch
for module in modules_target:
    if module not in modules_source:
        print(f"Checking out module: {module}")
        subprocess.run(["git", "checkout", target_branch, "--", module], check=True)