import argparse
import boto3
import os
import json

parser = argparse.ArgumentParser(description="Update RES module package URI")
parser.add_argument(
    "--environment-name", help="Name of the RES environment", required=True, type=str
)
parser.add_argument(
    "--module",
    help="Name of the module to update",
    required=True,
    type=str,
    choices=["cluster-manager", "virtual-desktop-controller", "virtual-desktop"],
)
parser.add_argument("--zip-file", help="Path to the zip file", required=True, type=str)
parser.add_argument("--s3-bucket", help="S3 bucket name", required=True, type=str)

args = parser.parse_args()

# Map module names to DDB keys
if args.module == "virtual-desktop-controller":
    module_key = "vdc.controller"
elif args.module == "virtual-desktop":
    module_key = "vdi-app"
else:
    module_key = args.module

ddb_key = f"{module_key}.app_package_uri"

session = boto3.session.Session()
s3_client = session.client("s3")
dynamodb_client = session.client("dynamodb")
iam_client = session.client("iam")

# 1. Upload zip file to S3
zip_filename = os.path.basename(args.zip_file)
s3_key = f"patches/{zip_filename}"
s3_uri = f"s3://{args.s3_bucket}/{s3_key}"

print(f"Uploading {args.zip_file} to {s3_uri}")
s3_client.upload_file(args.zip_file, args.s3_bucket, s3_key)
print(f"Successfully uploaded to {s3_uri}")

# 2. Update DynamoDB record
print(f"Updating {ddb_key} to {s3_uri}")
dynamodb_client.update_item(
    TableName=f"{args.environment_name}.cluster-settings",
    Key={"key": {"S": ddb_key}},
    UpdateExpression="SET #val = :val",
    ExpressionAttributeNames={"#val": "value"},
    ExpressionAttributeValues={":val": {"S": s3_uri}}
)
print(f"Successfully updated {ddb_key} to {s3_uri}")

# 3. Update IAM role with S3 GetObject permission
if args.module == "virtual-desktop-controller":
    role_name = f"{args.environment_name}-vdc-controller-role"
elif args.module == "virtual-desktop":
    role_name = f"{args.environment_name}-vdc-host-scoped-down-role"
else:
    role_name = f"{args.environment_name}-{args.module}-role"
policy_name = f"{args.environment_name}-s3-patch-access"
s3_resource_arn = f"arn:aws:s3:::{args.s3_bucket}/{s3_key}"

# Get existing policy if it exists
try:
    existing_policy = iam_client.get_role_policy(
        RoleName=role_name,
        PolicyName=policy_name
    )
    policy_document = existing_policy['PolicyDocument']
    
    # Add new resource to existing statement
    resources = policy_document['Statement'][0].get('Resource', [])
    if isinstance(resources, str):
        resources = [resources]
    if s3_resource_arn not in resources:
        resources.append(s3_resource_arn)
    policy_document['Statement'][0]['Resource'] = resources
    
except iam_client.exceptions.NoSuchEntityException:
    # Policy doesn't exist, create new one
    policy_document = {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Action": "s3:GetObject",
                "Resource": [s3_resource_arn]
            }
        ]
    }

print(f"Updating IAM role {role_name} with S3 access to {s3_resource_arn}")
iam_client.put_role_policy(
    RoleName=role_name,
    PolicyName=policy_name,
    PolicyDocument=json.dumps(policy_document)
)

if args.module == "virtual-desktop":
    print("Please launch a new VDI to use the updated package.")
else:
    print("Please terminate the existing infra instance and wait for a new one to be launched automatically.")