from __future__ import annotations

import re
from typing import Protocol


class MigrationConnection(Protocol):
    def exec_driver_sql(self, statement: str) -> object: ...

    def commit(self) -> None: ...


def migration_role_statement(role: str | None) -> str | None:
    """Return a safely quoted SET ROLE statement for an admitted owner role."""
    if not role:
        return None
    if not re.fullmatch(r"[a-z_][a-z0-9_]{0,62}", role):
        raise ValueError("HUB_CORE_MIGRATION_ROLE is not a safe PostgreSQL role name")
    return f'SET ROLE "{role}"'


def migration_schema_statement(schema: str | None) -> str | None:
    """Return an explicit, safely quoted migration search path."""
    if not schema:
        return None
    if not re.fullmatch(r"[a-z_][a-z0-9_]{0,62}", schema):
        raise ValueError("HUB_CORE_MIGRATION_SCHEMA is not a safe PostgreSQL schema name")
    return f'SET search_path TO "{schema}", public'


def configure_migration_session(
    connection: MigrationConnection,
    *,
    role: str | None,
    schema: str | None,
) -> None:
    """Apply and commit session settings before Alembic opens its transaction."""
    statements = tuple(
        statement
        for statement in (
            migration_role_statement(role),
            migration_schema_statement(schema),
        )
        if statement is not None
    )
    for statement in statements:
        connection.exec_driver_sql(statement)
    if statements:
        # SQLAlchemy 2 autobegins on exec_driver_sql().  Without this commit,
        # Alembic joins the setup transaction and its DDL is rolled back when
        # the connection closes.
        connection.commit()