import boto3
from typing import List

def get_profiles() -> List:
    return [
        #'gdb-artistapp-dev',
        #'gdb-artistapp-prod',
        #'gdb-apollo-dev',
        #'gdb-apollo-prod',
        'gdb-delphi-dev',
        'gdb-delphi-prod',
        #'gdb-databricks-dev',
        #'gdb-databricks-prod',
        #'gdb-datasci-dev',
        #'gdb-mct-dev',
        #'gdb-mct-prod',
        #'gdb-infra-dev',
        #'gdb-infra-prod',
    ]

def set_retention_settings(client, environment, retention_days):
    prefix = f'/aws/lambda/{environment}-'
    response = client.describe_log_groups(logGroupNamePrefix=prefix)

    for log_group in response['logGroups']:
        log_group_name = log_group['logGroupName']
        retention_policy = {
            'logGroupName': log_group_name,
            'retentionInDays': retention_days
        }
        client.put_retention_policy(**retention_policy)
        print(f"Set retention {retention_days} days for log group: {log_group_name}")

def main():
    profiles = get_profiles()
    sessions = []
    for profile in profiles:
        sessions.append(boto3.session.Session(
            profile_name=profile,
            region_name="us-east-1"
        ))

    clients = []
    for session in sessions:
        clients.append(session.client('logs'))

    environments = {
        'dev': 30,
        'qa-apollo': 60,
        'qa-delphi': 90,
        'stage-apollo': 90,
        'stage-delphi': 180,
        'uat-apollo': 180,
        'prod': 365
    }
    for client in clients:
        for environment, retention_days in environments.items():
            set_retention_settings(client, environment, retention_days)

if __name__ == '__main__':
    main()