"""Shared, fail-closed TLS configuration for database clients."""

from __future__ import annotations

import ssl
from urllib.parse import parse_qsl, urlencode, urlsplit

_TLS_QUERY_KEYS = {
    "ssl",
    "ssl_ca",
    "ssl_verify_cert",
    "ssl_verify_identity",
}
_FALSE_VALUES = {"0", "false", "no", "off"}
_TRUE_VALUES = {"1", "true", "yes", "on", "required", "verify_identity"}
# SQLAlchemy forwards unrecognised URL query values to aiomysql.connect(),
# which then boots with an opaque TypeError. In particular, asyncmy accepted
# read/write timeout keywords that aiomysql 0.3.2 does not expose; silently
# dropping an operator's timeout would be worse because it creates unbounded
# I/O. Keep an explicit driver-signature allowlist at the shared configuration
# boundary. ``connect_timeout`` remains supported by aiomysql.
_AIOMYSQL_QUERY_KEYS = {
    "auth_plugin",
    "autocommit",
    "charset",
    "client_flag",
    "connect_timeout",
    "echo",
    "init_command",
    "local_infile",
    "program_name",
    "read_default_file",
    "read_default_group",
    "server_public_key",
    "sql_mode",
    "unix_socket",
    "use_unicode",
}


def _query_options(url: str) -> dict[str, str]:
    return {key.lower(): value for key, value in parse_qsl(urlsplit(url).query)}


def database_tls_required(url: str) -> bool:
    """Return whether a URL requires an authenticated TLS connection."""
    options = _query_options(url)
    ssl_value = options.get("ssl", "").strip().lower()
    hostname = urlsplit(url).hostname or ""
    is_tidb_cloud = hostname == "tidbcloud.com" or hostname.endswith(".tidbcloud.com")
    return is_tidb_cloud or ssl_value in _TRUE_VALUES


def normalize_database_url(url: str) -> str:
    """Select the safe async driver and remove TLS control parameters.

    Production may still hold a legacy ``mysql+asyncmy`` URL in Parameter
    Store. Translate it before SQLAlchemy imports a driver so the unpatched
    asyncmy dict-key SQL injection cannot remain reachable during rollout.
    """
    if url.startswith("mysql+asyncmy://"):
        url = url.replace("mysql+asyncmy://", "mysql+aiomysql://", 1)
    if url.startswith("mysql+aiomysql://"):
        driver_keys = [
            key.lower()
            for key, _value in parse_qsl(urlsplit(url).query, keep_blank_values=True)
            if key.lower() not in _TLS_QUERY_KEYS
        ]
        duplicates = sorted({key for key in driver_keys if driver_keys.count(key) > 1})
        if duplicates:
            raise ValueError(
                "duplicate aiomysql database URL option(s) are ambiguous: "
                + ", ".join(duplicates)
            )
        unsupported = sorted(set(driver_keys) - _AIOMYSQL_QUERY_KEYS)
        if unsupported:
            names = ", ".join(unsupported)
            raise ValueError(
                f"aiomysql does not support database URL option(s): {names}. "
                "Remove them and enforce operation timeouts at the application "
                "boundary; connect_timeout and SQLAlchemy pool_timeout remain supported."
            )
    base, sep, query = url.partition("?")
    if not sep:
        return url
    kept = [
        (key, value)
        for key, value in parse_qsl(query, keep_blank_values=True)
        if key.lower() not in _TLS_QUERY_KEYS
    ]
    return base + ("?" + urlencode(kept) if kept else "")


def build_database_connect_args(url: str) -> dict[str, ssl.SSLContext]:
    """Build a certificate- and hostname-verifying SSL context.

    Weak verification switches are rejected instead of being silently
    accepted.  This keeps the API, workers, Alembic, and one-off scripts on
    the same authenticated-TLS contract.
    """
    options = _query_options(url)
    hostname = urlsplit(url).hostname or ""
    is_tidb_cloud = hostname == "tidbcloud.com" or hostname.endswith(".tidbcloud.com")
    if (
        is_tidb_cloud
        and options.get("ssl", "").strip().lower() in _FALSE_VALUES
    ):
        raise ValueError("TLS cannot be disabled for a TiDB Cloud database")
    for key in ("ssl_verify_cert", "ssl_verify_identity"):
        if (
            options.get(key, "").strip().lower() in _FALSE_VALUES
            and not is_tidb_cloud
        ):
            raise ValueError(f"insecure database TLS option is forbidden: {key}=false")

    # TiDB Cloud connection strings issued before authenticated TLS became
    # mandatory may still carry ``ssl_verify_cert=false``.  The hostname is a
    # stronger trust signal than that legacy client hint: TiDB Cloud always
    # requires TLS, and ``normalize_database_url`` removes the weak switch
    # before the driver sees it.  We then install our own CERT_REQUIRED,
    # hostname-checking context below.  This is a secure in-place upgrade, not
    # a compatibility downgrade; non-TiDB hosts remain fail-closed above.

    if not database_tls_required(url):
        return {}

    cafile = options.get("ssl_ca") or None
    context = ssl.create_default_context(cafile=cafile)
    # Explicit assertions make accidental weakening visible in tests and
    # protect this security boundary from future refactors.
    context.check_hostname = True
    context.verify_mode = ssl.CERT_REQUIRED
    return {"ssl": context}


def database_engine_options(url: str) -> tuple[str, dict[str, ssl.SSLContext]]:
    """Return the normalized URL and shared SQLAlchemy connect arguments."""
    return normalize_database_url(url), build_database_connect_args(url)