"""PostgreSQL logic.

Mirrors the public interface of src.logic.mysql (copy_users and run_scripts)
so the handler can dispatch on engine without special-casing.

Unlike MySQL, PostgreSQL roles are cluster-wide (not per user@host) and are
shared across all databases in the cluster. Roles and their attributes are read
from pg_roles, a view the RDS master user (rds_superuser) is permitted to read.

Password hashes are NOT migrated: on RDS the master user is rds_superuser, not
a true superuser, and cannot read pg_authid or pg_shadow, so role passwords are
simply not obtainable. Instead, known dev/QA passwords are supplied to
copy_users as a {role_name: password} map (read from Secrets Manager by
src.logic.secrets) and applied to matching roles via the PASSWORD clause on
both CREATE and ALTER. Setting the password on ALTER intentionally overwrites
the prod password carried over in the restored snapshot, which is the point of
sanitisation. A role with no entry in the map is left as-is: a brand-new role
is created without a password (and flagged), and an existing role keeps
whatever password it already has. Roles that need a password but have no secret
entry can still get one from the per-database users.sql sanitisation script
that runs after copy_users.

Both login roles and group/NOLOGIN roles are copied: group roles are commonly
the parents referenced by pg_auth_members memberships, so omitting them would
make membership grants fail. Roles are upserted (CREATE if missing, otherwise
ALTER) rather than dropped and recreated, because in a cluster restored from
the source snapshot a role almost always owns objects and DROP ROLE would fail.

The SUPERUSER, REPLICATION and BYPASSRLS attributes are NOT copied. Setting any
of them (even to the negative form, NOSUPERUSER/NOREPLICATION/NOBYPASSRLS, and
even on a non-superuser target role) requires a *true* superuser. On RDS the
master user is only rds_superuser, so emitting these clauses raises
"must be superuser to alter superusers" / "permission denied". Managed dev/QA
roles never need these attributes, so they are simply omitted: a new role gets
the PostgreSQL defaults (NO*) and an existing role keeps whatever it carried
over from the restored snapshot. This mirrors the password-hash limitation
above — both are cases where local Docker postgres (a real superuser) masks an
RDS-only restriction.

Sanitisation scripts run against the *application* database, not the cluster.
Unlike MySQL, PostgreSQL has no USE statement and a connection can only
reference the database it is connected to (cross-database names like
db.schema.table are not supported), so each script must declare which database
to connect to. Every .sql script must therefore begin with a directive:

    -- @database: <db_name>

run_scripts parses this header and opens a connection to the named database for
that file (a folder may target more than one database). A script with no
directive raises MissingDatabaseDirective. Role copy (copy_users) is unaffected
and stays on the cluster-wide "postgres" database.
"""
import glob
import re

import psycopg
import sqlparse
from lambdacommon.common_config import logger
from psycopg import sql
from psycopg.rows import dict_row

# Every sanitisation script must declare its target database with a leading
# `-- @database: <db_name>` directive (the PostgreSQL analogue of MySQL's USE;
# see the module docstring). By convention it is the first line, but it may be
# preceded by other comment lines.
_DATABASE_DIRECTIVE = re.compile(
    r'^\s*--\s*@database:\s*(\S+)\s*$', re.MULTILINE)

# Roles managed by RDS/PostgreSQL itself that must never be dropped or
# recreated. rdsadmin is the underlying superuser and the rds_* roles are
# predefined by RDS; pg_* roles are predefined by PostgreSQL.
RESERVED_ROLES = (
    'postgres_readonly_group',
    'rds_ad',
    'rds_extension',
    'rds_iam',
    'rds_password',
    'rds_replication',
    'rds_superuser',
    'rdsadmin',
    'readonly_group',
    'rdsrepladmin',
)


def run_scripts(db_credentials, scripts_dir):
    """Run a set of SQL scripts from a directory.

    Each script must begin with a `-- @database: <db_name>` directive naming
    the application database it targets; the connection is opened against that
    database for that script. PostgreSQL has no USE statement and cannot
    reference a database other than the one connected to, so the database is
    chosen per file rather than once for the directory (a folder may target
    more than one database). A script without the directive raises
    MissingDatabaseDirective (see the module docstring).

    Args:
        db_credentials: connection credentials (host/username/password) for the
            cluster; the target database is taken from each script's directive
        scripts_dir: the directory containing the scripts

    Returns:
        The number of scripts run
    """
    scripts = sorted(glob.glob(f'{scripts_dir}/*.sql'))
    logger.info(f'Found {len(scripts)} scripts in {scripts_dir}')
    for script in scripts:
        with open(script) as f:
            script_body = f.read()
        dbname = _parse_target_database(script_body, script)
        logger.info(f'Running SQL script {script} against database {dbname}')
        with _get_connection(db_credentials, dbname=dbname) as connection:
            with connection.cursor() as cursor:
                _run_script(cursor, script_body)
    return len(scripts)


def _parse_target_database(script_body, script_path):
    """Return the database named by a script's `-- @database:` directive.

    Raises MissingDatabaseDirective if the script has no directive.
    """
    match = _DATABASE_DIRECTIVE.search(script_body)
    if not match:
        raise MissingDatabaseDirective(
            f'{script_path} has no "-- @database: <db_name>" directive; every '
            'PostgreSQL sanitisation script must declare its target database.')
    return match.group(1)


def _run_script(cursor, script_body):
    """Run a single SQL script body against a database."""
    commands = sqlparse.split(script_body)
    for command in commands:
        # Logged at DEBUG: sanitisation scripts may embed sensitive
        # literals (emails, tokens) we do not want in CloudWatch/Sentry.
        logger.debug(f'Running command {command}')
        cursor.execute(command)


def copy_users(source_db_credentials, target_db_credentials,
               role_passwords=None):
    """Copy PostgreSQL roles from one database to another.

    role_passwords is an optional {role_name: password} map of known dev/QA
    passwords to apply to matching roles (see the module docstring). Roles
    absent from the map are left without a password change.
    """
    roles = _get_role_details(source_db_credentials)
    logger.info(
        f"Copying {len(roles)} roles from {source_db_credentials['host']} "
        f"to {target_db_credentials['host']}")
    return _restore_roles(
        target_db_credentials, roles, role_passwords or {})


def _get_role_details(db_credentials):
    """Retrieve roles, their attributes, password hashes and memberships.

    Returns all non-reserved roles in the cluster (both login and group roles)
    except for the role used to connect.
    """
    connection = _get_connection(db_credentials)
    with connection:
        with connection.cursor(row_factory=dict_row) as cursor:
            roles = _get_roles(cursor, db_credentials['username'])
        with connection.cursor() as cursor:
            for role in roles:
                role['member_of'] = _get_memberships(cursor, role['rolname'])
    return roles


def _restore_roles(creds, roles, role_passwords):
    """Restore roles and their memberships to a cluster.

    Roles are upserted rather than dropped first: DROP ROLE fails when a role
    owns objects or holds privileges, which is the norm in a cluster restored
    from the source snapshot. Roles are created/updated in one pass and then
    granted their memberships in a second pass so every parent role referenced
    in pg_auth_members already exists. role_passwords is a
    {role_name: password} map; a role's password is set when it has an entry.
    """
    connection = _get_connection(creds)
    with connection:
        with connection.cursor() as cursor:
            existing_roles = _get_existing_role_names(cursor)

            logger.info(f"Upserting {len(roles)} roles in {creds['host']}")
            for role in roles:
                _upsert_role(
                    cursor, role, role['rolname'] in existing_roles,
                    role_passwords.get(role['rolname']))
            for role in roles:
                _grant_memberships(cursor, role)
            logger.info(f"Successfully upserted roles in {creds['host']}")
            return len(roles)


def _get_connection(db_credentials, dbname='postgres'):
    """Return a connection from the given credentials to dbname.

    Defaults to the cluster-wide "postgres" database used for role copy; script
    running overrides it with the database from each script's directive.
    """
    return psycopg.connect(
        host=db_credentials['host'],
        user=db_credentials['username'],
        password=db_credentials['password'],
        dbname=dbname,
        autocommit=True,
    )


def _get_roles(cursor, role_to_exclude):
    """Retrieve roles from a PostgreSQL cluster.

    Returns all roles (both login and group/NOLOGIN roles) with their
    attributes, except for RDS/PostgreSQL reserved roles and the role specified
    by the role_to_exclude parameter. Group roles are included because they are
    commonly the parents referenced by pg_auth_members memberships. Roles are
    read from pg_roles; password hashes are not read because the RDS master
    user cannot access them (see the module docstring). Only the attributes the
    RDS master user can set are selected; rolsuper, rolreplication and
    rolbypassrls are deliberately omitted (see the module docstring).
    """
    excluded = list(RESERVED_ROLES) + [role_to_exclude]
    cursor.execute(
        'select rolname, rolinherit, rolcreaterole, rolcreatedb, '
        '       rolcanlogin, rolconnlimit '
        '  from pg_roles '
        ' where rolname not like %s '
        '   and rolname != all(%s)',
        ('pg\\_%', excluded))
    return cursor.fetchall()


def _get_existing_role_names(cursor):
    """Return the set of role names that already exist in the cluster."""
    cursor.execute('select rolname from pg_roles')
    return {row[0] for row in cursor.fetchall()}


def _get_memberships(cursor, rolname):
    """Return the names of roles that the given role is a member of."""
    cursor.execute(
        'select parent.rolname '
        '  from pg_auth_members member '
        '  join pg_roles parent on parent.oid = member.roleid '
        '  join pg_roles child on child.oid = member.member '
        ' where child.rolname = %s',
        (rolname,))
    return [row[0] for row in cursor.fetchall()]


def _build_role_options(role, password=None):
    """Build the WITH option clause shared by CREATE ROLE and ALTER ROLE.

    Only the attributes the RDS master user (rds_superuser) is permitted to set
    are emitted. SUPERUSER, REPLICATION and BYPASSRLS are superuser-only — even
    their negative forms fail on RDS — so they are intentionally not copied
    (see the module docstring).

    When password is not None a PASSWORD clause is appended. The password is
    embedded as a SQL literal because PASSWORD cannot be a bound parameter in
    CREATE/ALTER ROLE; sql.Literal escapes it safely.
    """
    options = [
        sql.SQL('LOGIN') if role['rolcanlogin'] else sql.SQL('NOLOGIN'),
        sql.SQL('INHERIT') if role['rolinherit'] else sql.SQL('NOINHERIT'),
        sql.SQL('CREATEROLE') if role['rolcreaterole']
        else sql.SQL('NOCREATEROLE'),
        sql.SQL('CREATEDB') if role['rolcreatedb'] else sql.SQL('NOCREATEDB'),
        sql.SQL('CONNECTION LIMIT {}').format(
            sql.Literal(role['rolconnlimit'])),
    ]
    # Source password hashes are not readable on RDS (see the module
    # docstring). A PASSWORD clause is only emitted when a known dev/QA
    # password was supplied; otherwise the role is upserted without one.
    if password is not None:
        options.append(sql.SQL('PASSWORD {}').format(sql.Literal(password)))
    return sql.SQL(' ').join(options)


def _upsert_role(cursor, role, exists, password=None):
    """Create the role, or update it in place if it already exists.

    Existing roles are altered rather than dropped because DROP ROLE fails when
    a role owns objects or holds privileges. When a password is supplied it is
    set on both the CREATE and ALTER paths, intentionally overwriting any
    password carried over from the source snapshot on ALTER. When no password
    is supplied an ALTER leaves the existing password intact and a new role is
    left passwordless until a users.sql script grants it one. The
    composed statement is never logged because it can contain the password.
    """
    verb = 'alter' if exists else 'create'
    logger.debug(f"{verb.capitalize()} role {role['rolname']}")
    statement = sql.SQL('{verb} role {name} with {options}').format(
        verb=sql.SQL(verb),
        name=sql.Identifier(role['rolname']),
        options=_build_role_options(role, password))
    cursor.execute(statement)

    if not exists and role['rolcanlogin'] and password is None:
        logger.warning(
            f"Created login role {role['rolname']} without a password; it "
            'cannot authenticate until a users.sql script sets one.')


def _grant_memberships(cursor, role):
    """Grant the role membership in each of its parent roles."""
    for parent in role['member_of']:
        logger.debug(f"Granting {role['rolname']} membership in {parent}")
        cursor.execute(
            sql.SQL('grant {parent} to {child}').format(
                parent=sql.Identifier(parent),
                child=sql.Identifier(role['rolname'])))


class MissingDatabaseDirective(Exception):
    """Raised when a sanitisation script has no `-- @database:` directive."""

    pass