import boto3
import argparse
import logging


def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument("--usernames", help="list of comma-separated user names")
    parser.add_argument("--roles", help="list of comma-separated role names")
    parser.add_argument(
        "--shared-databases",
        default="youtube_api_db",
        help="list of comma-separated shared database names: currently supported youtube_api_db and facebook_main_db")
    return parser.parse_args()

if __name__ == "__main__":
    logger = logging.getLogger("")
    handler = logging.StreamHandler()
    formatter = logging.Formatter('%(message)s')
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logger.setLevel(logging.INFO)

    args = parse_arguments()

    lf_client = boto3.client("lakeformation")
    glue_client = boto3.client("glue")



    users = [ f"arn:aws:iam::437795906767:user/{user}" for user in args.usernames.split(",")] if args.usernames is not None else []
    roles = [ f"arn:aws:iam::437795906767:role/{role}" for role in args.roles.split(",")] if args.roles is not None else []
    principals = users + roles
    shared_databases = args.shared_databases.split(",")
    # Map of shared databases to resource links
    shared_databases_map = {
        'youtube_api_db' : 'sme_youtube_api_db',
        'facebook_main_db': 'facebook_main_db',
    }
    catalog_id = '250735107403'
    rl_catalog_id = '437795906767'

    for principal in principals:
        for database in shared_databases:
            response = glue_client.get_tables(
                CatalogId=catalog_id,
                DatabaseName=database,
            )

            tables = ([r['Name'] for r in response['TableList']])
            logger.info(f"Grant DESCRIBE on database {catalog_id}/{database} to {principal}")
            response = lf_client.grant_permissions(
                Principal={
                    'DataLakePrincipalIdentifier': principal
                },
                Resource={
                    'Database': {
                        'CatalogId': catalog_id,
                        'Name': database
                    },
                },
                Permissions=[
                    'DESCRIBE'
                ],
            )

            for table in tables:
                logger.info(f"Grant DESCRIBE, SELECT on table {catalog_id}/{database}/{table} to {principal}")
                response = lf_client.grant_permissions(
                    Principal={
                        'DataLakePrincipalIdentifier': principal
                    },
                    Resource={
                        'Table': {
                            'CatalogId': catalog_id,
                            'DatabaseName': database,
                            'Name': table,
                        },
                    },
                    Permissions=[
                        'SELECT',
                        'DESCRIBE',
                    ],
                )

            if database in shared_databases_map:
                logger.info(f"Grant DESCRIBE on resource link {rl_catalog_id}/{shared_databases_map[database]} to {principal}")
                response = lf_client.grant_permissions(
                    Principal={
                        'DataLakePrincipalIdentifier': principal
                    },
                    Resource={
                        'Database': {
                            'CatalogId': rl_catalog_id,
                            'Name': shared_databases_map[database]
                        },
                    },
                    Permissions=[
                        'DESCRIBE'
                    ],
                )