"""
https://www.notion.so/RDS-Cluster-and-User-Creation-Self-Service-55af0b6e65a44d21a8d0702059cbdaab#12897177520f800d9f79e1c46e3ae885
"""

import argparse
import asyncio
import secrets
import string
from pprint import pprint
from typing import TYPE_CHECKING

import aioboto3
import sqlalchemy
import prettytable
from sqlalchemy.ext.asyncio import create_async_engine, AsyncEngine
from src.constants import DATABASE_OPTIONS, Environment, UserType, DatabaseCluster

if TYPE_CHECKING:
    from types_aiobotocore_rds.client import RDSClient
    from types_aiobotocore_rds.type_defs import DBClusterTypeDef


async def main() -> None:
    # parse args
    (clusters, users, user_type, create, drop) = parse_args(DATABASE_OPTIONS)

    session = aioboto3.Session()
    async with session.client("rds") as rds:
        results = await asyncio.gather(
            *[
                _process_cluster(rds, cluster, users, user_type, create, drop)
                for cluster in clusters
            ],
            return_exceptions=True,
        )

    # merge results and report errors
    created_users: dict[str, list[tuple[str, str]]] = {}
    for cluster, result in zip(clusters, results):
        if isinstance(result, BaseException):
            print(f"[{cluster.cluster_id}] Error: {result}")
        else:
            for username, entries in result.items():
                created_users.setdefault(username, []).extend(entries)

    # print summary by user
    for username, details in created_users.items():
        table = prettytable.PrettyTable()
        table.field_names = ["Username", "Database Endpoint", "Password"]
        for db_endpoint, password in details:
            table.add_row([username, db_endpoint, password])
        print(table)
        print()


async def _process_cluster(
    rds: RDSClient,
    cluster: DatabaseCluster,
    users: list[str],
    user_type: UserType,
    create: bool,
    drop: bool,
) -> dict[str, list[tuple[str, str]]]:
    # get database metadata
    db_info = await describe_db_cluster(rds, cluster.cluster_id)
    db_endpoint = db_info["Endpoint"]
    master_username = db_info["MasterUsername"]

    # reset master password
    master_password = random_password()
    await reset_master_password(rds, cluster.cluster_id, master_password)

    # perform operations
    result: dict[str, list[tuple[str, str]]] = {}
    if drop:
        await drop_users(db_endpoint, master_username, master_password, users)
    if create:
        new_users = [(username, random_password()) for username in users]
        await create_users(
            db_endpoint,
            master_username,
            master_password,
            new_users,
            user_type,
        )
        for username, password in new_users:
            result.setdefault(username, []).append((db_endpoint, password))
    return result


def parse_args(
    options: dict[str, dict[Environment, DatabaseCluster]],
) -> tuple[list[DatabaseCluster], list[str], UserType, bool, bool]:
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        "--users",
        metavar="USERNAME",
        type=str,
        nargs="+",
        required=True,
        help="One or more usernames to add db users for.",
    )
    parser.add_argument(
        "--dbs",
        type=str,
        choices=options.keys(),
        nargs="+",
        required=True,
        help="Databases to manage.",
    )
    parser.add_argument(
        "--env",
        type=Environment,
        choices=list(Environment),
        required=True,
        help="Database environment.",
    )
    parser.add_argument(
        "--type",
        type=UserType,
        choices=list(UserType),
        help="Type of database user.",
    )
    parser.add_argument("-c", action="store_true", help="Create users.")
    parser.add_argument("-d", action="store_true", help="Drop users.")
    args = parser.parse_args()

    clusters: list[DatabaseCluster] = []
    for db in args.dbs:
        cluster = options[db].get(args.env)
        if not cluster:
            parser.error(f"Database {db} not available in environment {args.env}.")
        if args.type and args.type not in cluster.allowed_user_types:
            parser.error(
                f"User type {args.type} not allowed for database {db} in environment {args.env}."
            )
        clusters.append(cluster)
    if args.c and not args.type:
        parser.error("User type must be specified when creating users.")
    if not args.c and not args.d:
        parser.error("At least one of -c or -d must be specified.")
    return (clusters, args.users, args.type, args.c, args.d)


async def create_users(
    db_endpoint: str,
    master_username: str,
    master_password: str,
    users: list[tuple[str, str]],
    UserType,
) -> None:
    engine = _create_engine(db_endpoint, master_username, master_password)
    try:
        async with engine.connect() as connection:
            for username, password in users:
                await connection.execute(
                    sqlalchemy.text(
                        "CREATE USER :username@'%' IDENTIFIED BY :password;"
                    ),
                    {"username": username, "password": password},
                )
                await connection.execute(
                    sqlalchemy.text(
                        "GRANT SELECT, SHOW VIEW, SHOW_ROUTINE ON *.* TO :username@'%';"
                    ),
                    {"username": username},
                )
                if UserType.READ_WRITE:
                    await connection.execute(
                        sqlalchemy.text(
                            "GRANT INSERT, UPDATE, DELETE, TRIGGER ON *.* TO :username@'%';"
                        ),
                        {"username": username},
                    )
            await connection.commit()
    finally:
        await engine.dispose()


async def drop_users(
    db_endpoint: str, master_username: str, master_password: str, users: list[str]
) -> None:
    engine = _create_engine(db_endpoint, master_username, master_password)
    try:
        async with engine.connect() as connection:
            for username in users:
                await connection.execute(
                    sqlalchemy.text("DROP USER IF EXISTS :username@'%';"),
                    {"username": username},
                )
            await connection.commit()
    finally:
        await engine.dispose()


def _create_engine(db_endpoint: str, username: str, password: str) -> AsyncEngine:
    return create_async_engine(
        sqlalchemy.URL.create(
            drivername="mysql+aiomysql",
            username=username,
            password=password,
            host=db_endpoint,
            port=3306,
            database=None,
        )
    )


async def reset_master_password(rds: RDSClient, cluster_id: str, password: str) -> None:
    response = await rds.modify_db_cluster(
        DBClusterIdentifier=cluster_id,
        MasterUserPassword=password,
        ApplyImmediately=True,
    )
    if response["ResponseMetadata"]["HTTPStatusCode"] != 200:
        pprint(response)
        raise Exception(f"Failed to reset master password for DB cluster {cluster_id}")
    print(f"[{cluster_id}] Master password reset initiated")
    while True:
        cluster_info = await describe_db_cluster(rds, cluster_id)
        if "PendingModifiedValues" not in cluster_info or not cluster_info[
            "PendingModifiedValues"
        ].get("MasterUserPassword"):
            break
        await asyncio.sleep(3)
    print(f"[{cluster_id}] Master password reset completed!")


async def describe_db_cluster(rds: RDSClient, cluster_id: str) -> DBClusterTypeDef:
    response = await rds.describe_db_clusters(DBClusterIdentifier=cluster_id)
    if response["ResponseMetadata"]["HTTPStatusCode"] != 200:
        pprint(response)
        raise Exception(f"Failed to describe DB cluster {cluster_id}")
    return response["DBClusters"][0]


def random_password() -> str:
    """Generate a complex random password."""
    password = ""
    while len(password) < 24:
        upper = secrets.choice(string.ascii_uppercase)
        lower = secrets.choice(string.ascii_lowercase)
        num = secrets.choice(string.digits)
        symbol = secrets.choice("#%^()")
        chars = list(upper + lower + num + symbol)
        secrets.SystemRandom().shuffle(chars)
        password += "".join(chars)
    return password