#!/bin/python
import boto3
from pprint import pprint, pformat
import json
import sys
import argparse
import logging
import secrets
import psycopg2
from psycopg2 import sql
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT # <-- ADD THIS LINE
from psycopg2.extensions import ISOLATION_LEVEL_DEFAULT

def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--accesskey', help='AWS Access key', default='')
    parser.add_argument('--secretkey', help='AWS Secret key', default='')
    parser.add_argument('--profile', help='AWS profile', default='gdb-delphi-dev')
    parser.add_argument('--region', help='AWS Default region', default='us-east-1')
    parser.add_argument('--instance-id', help='ID of RDS instance', default='')
    parser.add_argument('--cluster-id', help='ID of Aurora cluster', default='')
    requiredNamed = parser.add_argument_group('required named arguments')
    requiredNamed.add_argument('--database', help='Name of the database', required=True)

    return parser.parse_args()

def set_db_secret(
  secrets_client, secret_name, username, hostname,
  database, masterarn, active_endpoint=False):
    pwd = secrets.token_urlsafe(32)
    secret = {
        "username": username,
        "password": pwd,
        "engine": "postgres",
        "host": hostname,
        "port": 5432,
        "database": database,
        "masterarn": masterarn,
    }
    if active_endpoint: secret["active_endpoint"] = hostname
    secrets_client.put_secret_value(
        SecretId=secret_name, SecretString=json.dumps(secret))
    return secret

def rds_create_role(connection, rolename, password='', force=False):
    cur = connection.cursor()
    if force:
        cur.execute(
            sql.SQL("DROP ROLE IF EXISTS {rolename}")\
                .format(rolename=sql.Identifier(rolename)))
    if password:
        query = sql.SQL(
        '''CREATE USER {rolename} WITH
        LOGIN
        NOSUPERUSER
        CREATEDB
        NOCREATEROLE
        INHERIT
        NOREPLICATION
        CONNECTION LIMIT -1
        PASSWORD %s;
        ''').format(rolename=sql.Identifier(rolename))
        cur.execute(query, [password])
    else:
        query = sql.SQL(
        '''CREATE USER {rolename} WITH
        LOGIN
        NOSUPERUSER
        CREATEDB
        NOCREATEROLE
        INHERIT
        NOREPLICATION
        CONNECTION LIMIT -1;
        ''').format(rolename=sql.Identifier(rolename))
        cur.execute(query)
    connection.commit()

def rds_grant_role(connection, rolename, grantie):
    cur = connection.cursor()
    cur.execute(
        sql.SQL("GRANT {rolename} TO {grantie};").format(
            rolename=sql.Identifier(rolename),
            grantie=sql.Identifier(grantie)))
    connection.commit()

def rds_create_database(connection, dbname, owner):
    connection.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
    cur = connection.cursor()

    query = sql.SQL(
        '''CREATE DATABASE {dbname}
            WITH
            OWNER = {owner}
            ENCODING = 'UTF8'
            LC_COLLATE = 'en_US.UTF-8'
            LC_CTYPE = 'en_US.UTF-8'
            CONNECTION LIMIT = -1;
        ''').format(
            dbname=sql.Identifier(dbname),
            owner=sql.Identifier(owner)
        )

    cur.execute(query)
    connection.commit()
    cur.execute(
        sql.SQL("GRANT ALL ON DATABASE {dbname} TO {owner};").format(
            dbname=sql.Identifier(dbname),
            owner=sql.Identifier(owner)
        )
    )
    cur.execute(
        sql.SQL("GRANT TEMPORARY, CONNECT ON DATABASE {dbname} TO PUBLIC;").format(
            dbname=sql.Identifier(dbname)
        )
    )
    connection.commit()

def setup_db(secret_client, platform, db_hostname, env, database, cluster_admin):
    database_owner = database + "_owner"
    database_user = database + "_user"
    database_user_iam = database + "_user_iam"
    database_reader = database + "_reader"
    database_reader_iam = database + "_reader_iam"
    # Get root pwd and hostname
    root_secret_path = f"maintenance/root/{platform}/{env}/{cluster_admin}"
    root_secret_arn = secret_client.describe_secret(SecretId=root_secret_path)['ARN']
    root_secret = json.loads(
        secrets_client.get_secret_value(SecretId=root_secret_path)['SecretString']
    )

    con = psycopg2.connect(
        database="postgres",
        user=root_secret['username'],
        password=root_secret['password'],
        host=db_hostname,
        port="5432"
    )
    # Create roles
    # Generate owner secret and create owner secret
    database_owner_secret_path = f"maintenance/common/{platform}/{env}/{database}_owner"
    database_owner_secret = set_db_secret(
        secrets_client, database_owner_secret_path,
        database_owner, db_hostname,
        database, root_secret_arn,
        active_endpoint=False)
    rds_create_role(con, database_owner, database_owner_secret['password'])
    rds_grant_role(con, database_owner, cluster_admin)

    # Generate user secret and and create db role
    database_user_secret_path = f"{platform}/{env}/{database}/pg/user"
    database_user_secret = set_db_secret(
        secrets_client, database_user_secret_path,
        database_user, db_hostname,
        database, root_secret_arn,
        active_endpoint=True)
    rds_create_role(con, database_user, database_user_secret['password'])

    # Generate reader secret and and create db role
    database_reader_secret_path = f"{platform}/{env}/{database}/pg/reader"
    database_reader_secret = set_db_secret(
        secrets_client, database_reader_secret_path,
        database_reader, db_hostname,
        database, root_secret_arn,
        active_endpoint=True)
    rds_create_role(con, database_reader, database_reader_secret['password'])

    # create user_iam, reader_iam roles
    rds_create_role(con, database_user_iam)
    rds_grant_role(con, "rds_iam", database_user_iam)
    rds_grant_role(con, database_user, database_user_iam)
    rds_create_role(con, database_reader_iam)
    rds_grant_role(con, "rds_iam", database_reader_iam)
    rds_grant_role(con, database_reader, database_reader_iam)

    # create database
    rds_create_database(con, database, database_owner)
    # make dump, load it
    con.close()

if __name__ == "__main__":
    logger = logging.getLogger('')
    logger.setLevel(logging.DEBUG)

    args = parse_arguments()

    # Initialize the client with credentials or with IAM role.
    if (args.accesskey and args.secretkey):
        secrets_client = boto3.client(
          'secretsmanager',
          aws_access_key_id = args.accesskey,
          aws_secret_access_key = args.secretkey,
          region_name = args.region
        )
        rds_client = boto3.client(
          'rds',
          aws_access_key_id = args.accesskey,
          aws_secret_access_key = args.secretkey,
          region_name = args.region
        )
    elif args.profile:
        session = boto3.Session(profile_name=args.profile)
        secrets_client = session.client('secretsmanager')
        rds_client = session.client('rds')
    else:
        secrets_client = boto3.client('secretsmanager', region_name=args.region)
        rds_client = boto3.client('rds', region_name=args.region)

    if (not args.cluster_id and not args.instance_id):
        raise ValueError("No RDS instance or Aurora cluster are specified")
    elif (args.cluster_id and args.instance_id):
        raise ValueError("Error - both RDS instance ID or Aurora cluster ID are specified. It's allowed to pass only one at a time")
    elif (args.instance_id):
        # clusters are named as env-platform-clustername, admin user is named as clustername_admin
        env = args.instance_id.split('-', maxsplit=2)[0]
        platform = args.instance_id.split('-', maxsplit=2)[1]
        admin_user = (args.instance_id.split('-', maxsplit=2)[-1] + '_admin').replace('-', '_')
        hostname = rds_client.describe_db_instances(
            DBInstanceIdentifier=args.instance_id
        )['DBInstances'][0]['Endpoint']['Address']
    else:
        # clusters are named as env-platform-clustername, admin user is named as clustername_admin
        env = args.cluster_id.split('-', maxsplit=2)[0]
        platform = args.cluster_id.split('-', maxsplit=2)[1]
        admin_user = (args.cluster_id.split('-', maxsplit=2)[-1] + '_admin').replace('-', '_')
        endpoints = rds_client.describe_db_cluster_endpoints(
            DBClusterIdentifier=args.cluster_id)['DBClusterEndpoints']
        for endpoint in endpoints:
            if endpoint['EndpointType'] == 'WRITER':
                hostname = endpoint['Endpoint']
                break

    setup_db(secrets_client, platform, hostname, env, args.database, admin_user)