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

from __future__ import annotations

import argparse
import asyncio
import re
from collections.abc import AsyncIterator
from typing import TYPE_CHECKING

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

from common.utils import password

if TYPE_CHECKING:
    from types_aiobotocore_rds.client import RDSClient
    from types_aiobotocore_rds.type_defs import DBClusterTypeDef
    from types_aiobotocore_iam.client import IAMClient
    from types_aiobotocore_sts.client import STSClient

_POLICY_RE = re.compile(
    r"^(?:ABAC|RBAC)-(?P<family>.+?)-(?:resource|tag)-access-policy"
)
_USERNAME_RE = re.compile(r"^[a-zA-Z0-9_-]{2,}$")


async def main() -> None:
    (environment, users, user_type, create, drop) = parse_args()

    session = aioboto3.Session()

    # determine application families the current user has access to
    async with session.client("iam") as iam, session.client("sts") as sts:
        is_admin, app_families = await asyncio.gather(
            is_current_user_admin(iam, sts),
            get_current_user_app_families(iam, sts),
        )
    if not is_admin and not app_families:
        raise Exception(
            "Current user does not have any application family access policies attached"
        )

    # determine dbs user can choose from
    async with session.client("rds") as rds:
        if is_admin:
            all_db_clusters = await describe_mysql_clusters_by_tags(rds, environment)
        else:
            all_db_clusters = await describe_mysql_clusters_by_tags(
                rds, environment, app_families
            )
        if not all_db_clusters:
            raise Exception(
                f"No RDS clusters found for environment {environment.value}"
                + (
                    ""
                    if is_admin
                    else f" and application families {', '.join(app_families)}"
                )
            )
        db_clusters = prompt_select_db_instances(all_db_clusters)
        results = await asyncio.gather(
            *[
                _process_cluster(rds, cluster, users, user_type, create, drop)
                for cluster in db_clusters
            ],
            return_exceptions=True,
        )

    # merge results and report errors
    created_users: dict[str, list[tuple[str, str]]] = {}
    for cluster, result in zip(db_clusters, results):
        if isinstance(result, BaseException):
            print(f"[{cluster['DBClusterIdentifier']}] 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, db_password in details:
            table.add_row([username, db_endpoint, db_password])
        print(table)
        print()


async def _process_cluster(
    rds: RDSClient,
    db_info: DBClusterTypeDef,
    users: list[str],
    user_type: UserType | None,
    create: bool,
    drop: bool,
) -> dict[str, list[tuple[str, str]]]:
    # get database metadata
    db_endpoint = db_info["Endpoint"]
    master_username = db_info["MasterUsername"]

    # reset master password
    master_password = password.generate_random_password()
    await reset_master_password(rds, db_info["DBClusterIdentifier"], 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 and user_type is not None:
        new_users = [
            (username, password.generate_random_password()) for username in users
        ]
        await create_users(
            db_endpoint,
            master_username,
            master_password,
            new_users,
            user_type,
        )
        for username, db_password in new_users:
            result.setdefault(username, []).append((db_endpoint, db_password))
    return result


def parse_args() -> tuple[Environment, list[str], UserType | None, bool, bool]:
    parser = argparse.ArgumentParser(
        prog="rds-user-creation", 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(
        "--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()

    for user in args.users:
        if not _USERNAME_RE.match(user):
            parser.error(f"Invalid username {user!r}")
    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.")
    if args.type == UserType.READ_WRITE and args.env == Environment.PROD:
        parser.error("Read-write users cannot be created in production environment.")
    return (args.env, 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]],
    user_type: UserType,
) -> None:
    engine = _create_engine(db_endpoint, master_username, master_password)
    try:
        async with engine.connect() as connection:
            for username, db_password in users:
                await connection.execute(
                    sqlalchemy.text(
                        "CREATE USER :username@'%' IDENTIFIED BY :password;"
                    ),
                    {"username": username, "password": db_password},
                )
                await connection.execute(
                    sqlalchemy.text(
                        "GRANT SELECT, SHOW VIEW, SHOW_ROUTINE ON *.* TO :username@'%';"
                    ),
                    {"username": username},
                )
                if user_type == 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, db_password: str) -> AsyncEngine:
    return create_async_engine(
        sqlalchemy.URL.create(
            drivername="mysql+aiomysql",
            username=username,
            password=db_password,
            host=db_endpoint,
            port=3306,
            database=None,
        )
    )


async def cluster_ready(rds: RDSClient, cluster_id: str) -> bool:
    response = await rds.describe_db_clusters(DBClusterIdentifier=cluster_id)
    cluster_info = response["DBClusters"][0]
    return "PendingModifiedValues" not in cluster_info or not cluster_info[
        "PendingModifiedValues"
    ].get("MasterUserPassword")


async def reset_master_password(
    rds: RDSClient, cluster_id: str, db_password: str
) -> None:
    await rds.modify_db_cluster(
        DBClusterIdentifier=cluster_id,
        MasterUserPassword=db_password,
        ApplyImmediately=True,
    )
    print(f"[{cluster_id}] Master password reset initiated")
    while True:
        if await cluster_ready(rds, cluster_id):
            break
        await asyncio.sleep(3)
    print(f"[{cluster_id}] Master password reset completed")


async def describe_mysql_clusters_by_tags(
    rds: RDSClient, environment: Environment, app_families: set[str] | None = None
) -> list[DBClusterTypeDef]:
    instances: list[DBClusterTypeDef] = []
    paginator = rds.get_paginator("describe_db_clusters")
    async for page in paginator.paginate(
        Filters=[{"Name": "engine", "Values": ["aurora-mysql"]}]
    ):
        for instance in page["DBClusters"]:
            tags = {t["Key"]: t["Value"] for t in instance.get("TagList", [])}
            if (
                app_families is None or tags.get("application_family") in app_families
            ) and tags.get("environment") == environment.value:
                instances.append(instance)
    return instances


async def _get_caller_principal(sts: STSClient) -> tuple[str, str]:
    """Returns (principal_type, principal_name) for the current caller.

    principal_type is 'user' for IAM users or 'role' for assumed roles.
    """
    identity = await sts.get_caller_identity()
    # ARN formats:
    #   IAM user:      arn:aws:iam::123456789012:user/johndoe
    #   Assumed role:  arn:aws:sts::123456789012:assumed-role/MyRole/MySession
    resource = identity["Arn"].split(":", 5)[5]
    resource_type, resource_name = resource.split("/")[0], resource.split("/")[1]
    return ("role" if resource_type == "assumed-role" else "user"), resource_name


async def _iter_principal_policy_names(
    iam: IAMClient, principal_type: str, principal_name: str
) -> AsyncIterator[str]:
    if principal_type == "user":
        user_policy_paginator = iam.get_paginator("list_attached_user_policies")
        async for page in user_policy_paginator.paginate(UserName=principal_name):
            for policy in page["AttachedPolicies"]:
                yield policy["PolicyName"]

        group_paginator = iam.get_paginator("list_groups_for_user")
        async for group_page in group_paginator.paginate(UserName=principal_name):
            for group in group_page["Groups"]:
                policy_paginator = iam.get_paginator("list_attached_group_policies")
                async for page in policy_paginator.paginate(  # type: ignore[assignment]
                    GroupName=group["GroupName"]
                ):
                    for policy in page["AttachedPolicies"]:
                        yield policy["PolicyName"]
    else:
        role_policy_paginator = iam.get_paginator("list_attached_role_policies")
        async for page in role_policy_paginator.paginate(RoleName=principal_name):
            for policy in page["AttachedPolicies"]:
                yield policy["PolicyName"]


async def is_current_user_admin(iam: IAMClient, sts: STSClient) -> bool:
    principal_type, principal_name = await _get_caller_principal(sts)
    async for policy_name in _iter_principal_policy_names(
        iam, principal_type, principal_name
    ):
        if policy_name == "AdministratorAccess":
            return True
    return False


async def get_current_user_app_families(iam: IAMClient, sts: STSClient) -> set[str]:
    principal_type, principal_name = await _get_caller_principal(sts)
    app_families = set()
    async for policy_name in _iter_principal_policy_names(
        iam, principal_type, principal_name
    ):
        app_family = extract_app_family(policy_name)
        if app_family:
            app_families.add(app_family)
    return app_families


def extract_app_family(policy_name: str) -> str | None:
    m = _POLICY_RE.match(policy_name)
    return m.group("family") if m else None


def prompt_select_db_instances(
    instances: list[DBClusterTypeDef],
) -> list[DBClusterTypeDef]:
    sorted_instances = sorted(
        instances,
        key=lambda i: (
            {t["Key"]: t["Value"] for t in i.get("TagList", [])}.get(
                "application_family", ""
            ),
            i["DBClusterIdentifier"],
        ),
    )

    table = prettytable.PrettyTable()
    table.field_names = ["#", "DBClusterIdentifier", "Application Family"]
    for idx, instance in enumerate(sorted_instances, start=1):
        tags = {t["Key"]: t["Value"] for t in instance.get("TagList", [])}
        table.add_row(
            [
                idx,
                instance["DBClusterIdentifier"],
                tags.get("application_family", "---"),
            ]
        )
    print(table)

    while True:
        raw = input(
            "Select databases by number (space-separated) or '*' for all: "
        ).strip()
        if raw == "*":
            return sorted_instances

        selected: list[DBClusterTypeDef] = []
        for token in raw.split():
            try:
                idx = int(token)
            except ValueError:
                print(f"Invalid input: {token!r}. Enter integers or '*'.")
                break
            if idx < 1 or idx > len(sorted_instances):
                print(
                    f"Selection out of range: {idx}. Enter a number between 1 and {len(sorted_instances)}."
                )
                break
            selected.append(sorted_instances[idx - 1])
        else:
            return selected