import os
from azure.identity import DefaultAzureCredential
from azure.mgmt.compute import ComputeManagementClient

# Authenticate
credential = DefaultAzureCredential()
subscription_id = os.environ["AZURE_SUBSCRIPTION_ID"]
compute_client = ComputeManagementClient(credential, subscription_id)

def stop_vms_by_tag(tag_name, tag_value):
    print(f"Stopping VMs with tag {tag_name}={tag_value}...")
    vms = compute_client.virtual_machines.list_all()
    
    for vm in vms:
        if vm.tags and vm.tags.get(tag_name) == tag_value:
            print(f"Stopping VM: {vm.name} in {vm.location}")
            # Deallocate stops the VM and releases the compute resources (stops billing)
            async_vm_stop = compute_client.virtual_machines.begin_deallocate(
                resource_group_name=vm.id.split("/")[4],
                vm_name=vm.name
            )
            async_vm_stop.wait()
            print(f"VM {vm.name} stopped.")

if __name__ == "__main__":
    # Example: Stop all VMs tagged with Environment=Dev
    stop_vms_by_tag("Environment", "Dev")
