#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Location: ./mcpgateway/utils/db_isready.py
Copyright 2026
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti

db_isready - Wait until the configured database is ready
==========================================================
This helper blocks until the given database (defined by an **SQLAlchemy** URL)
successfully answers a trivial round-trip - ``SELECT 1`` - and then returns.
It is useful as a container **readiness/health probe** or imported from Python
code to delay start-up of services that depend on the DB.

Exit codes when executed as a script
-----------------------------------
* ``0`` - database ready.
* ``1`` - all attempts exhausted / timed-out.
* ``2`` - :pypi:`SQLAlchemy` is **not** installed.
* ``3`` - invalid parameter combination (``max_tries``/``interval``/``timeout``).

Features
--------
* Accepts **any** SQLAlchemy URL supported by the installed version.
* **Exponential backoff with jitter** - prevents thundering herd on reconnect:
  - Retry delays: 2s → 4s → 8s → 16s → 30s (capped) → 30s...
  - Random jitter of ±25% prevents synchronized reconnection storms
  - Default: 30 retries ≈ 5 minutes total wait before giving up
* Timing knobs (tries, interval, connect-timeout) configurable through
  *environment variables* **or** *CLI flags* - see below.
* Works **synchronously** (blocking) or **asynchronously** - simply
  ``await wait_for_db_ready()``.
* Credentials appearing in log lines are automatically **redacted**.
* Depends only on ``sqlalchemy`` (already required by *mcpgateway*).

Environment variables
---------------------
The script falls back to :pydata:`mcpgateway.config.settings`, but the values
below can be overridden via environment variables *or* the corresponding
command-line options.

+----------------------------+----------------------------------------------+-----------+
| Name                       | Description                                  | Default   |
+============================+==============================================+===========+
| ``DATABASE_URL``           | SQLAlchemy connection URL                    | ``sqlite:///./mcp.db`` |
| ``DB_WAIT_MAX_TRIES``      | Maximum attempts before giving up            | ``30``    |
| ``DB_WAIT_INTERVAL``       | Delay between attempts *(seconds)*           | ``2``     |
| ``DB_CONNECT_TIMEOUT``     | Per-attempt connect timeout *(seconds)*      | ``2``     |
| ``DB_MAX_BACKOFF_SECONDS`` | Max backoff cap *(seconds, jitter added)*    | ``30``    |
| ``LOG_LEVEL``              | Log verbosity when not set via ``--log-level`` | ``INFO`` |
+----------------------------+----------------------------------------------+-----------+

Usage examples
--------------
Shell ::

    python3 db_isready.py
    python3 db_isready.py --database-url "sqlite:///./mcp.db" --max-tries 2 --interval 1 --timeout 1

Python ::

    from mcpgateway.utils.db_isready import wait_for_db_ready

    # Synchronous/blocking
    wait_for_db_ready(sync=True)

    # Asynchronous
    import asyncio
    asyncio.run(wait_for_db_ready())

Doctest examples
----------------
>>> from mcpgateway.utils.db_isready import wait_for_db_ready
>>> import logging
>>> class DummyLogger:
...     def __init__(self): self.infos = []
...     def info(self, msg): self.infos.append(msg)
...     def debug(self, msg): pass
...     def error(self, msg): pass
...     @property
...     def handlers(self): return [True]
>>> import sys
>>> sys.modules['sqlalchemy'] = type('sqlalchemy', (), {
...     'create_engine': lambda *a, **k: type('E', (), {'connect': lambda self: type('C', (), {'execute': lambda self, q: 1, '__enter__': lambda self: self, '__exit__': lambda self, exc_type, exc_val, exc_tb: None})()})(),
...     'text': lambda q: q,
...     'engine': type('engine', (), {'Engine': object, 'URL': object, 'url': type('url', (), {'make_url': lambda u: type('U', (), {'get_backend_name': lambda self: "sqlite"})()}),}),
...     'exc': type('exc', (), {'OperationalError': Exception})
... })
>>> wait_for_db_ready(database_url='sqlite:///./mcp.db', max_tries=1, interval=1, timeout=1, logger=DummyLogger(), sync=True)
>>> try:
...     wait_for_db_ready(database_url='sqlite:///./mcp.db', max_tries=0, interval=1, timeout=1, logger=DummyLogger(), sync=True)
... except RuntimeError as e:
...     print('error')
error
"""

# Future
from __future__ import annotations

# Standard
# ---------------------------------------------------------------------------
# Standard library imports
# ---------------------------------------------------------------------------
import argparse
import asyncio
import logging
import os
import random
import re
import sys
import time
from typing import Any, Dict, Final, Optional

# ---------------------------------------------------------------------------
# Third-party imports - abort early if SQLAlchemy is missing
# ---------------------------------------------------------------------------
try:
    # Third-Party
    from sqlalchemy import create_engine, text
    from sqlalchemy.engine import Engine, URL
    from sqlalchemy.engine.url import make_url
    from sqlalchemy.exc import OperationalError
except ImportError:  # pragma: no cover - handled at runtime for the CLI
    sys.stderr.write("SQLAlchemy not installed - aborting (pip install sqlalchemy)\n")
    sys.exit(2)

# ---------------------------------------------------------------------------
# Optional project settings (silently ignored if mcpgateway package is absent)
# ---------------------------------------------------------------------------
try:
    # First-Party
    from mcpgateway.config import settings
except Exception:  # pragma: no cover - fallback minimal settings

    class _Settings:
        """Fallback dummy settings when *mcpgateway* is not import-able."""

        database_url: str = "sqlite:///./mcp.db"
        log_level: str = "INFO"

    settings = _Settings()  # type: ignore

# ---------------------------------------------------------------------------
# Environment variable names
# ---------------------------------------------------------------------------
ENV_DB_URL: Final[str] = "DATABASE_URL"
ENV_MAX_TRIES: Final[str] = "DB_WAIT_MAX_TRIES"
ENV_INTERVAL: Final[str] = "DB_WAIT_INTERVAL"
ENV_TIMEOUT: Final[str] = "DB_CONNECT_TIMEOUT"
ENV_MAX_BACKOFF: Final[str] = "DB_MAX_BACKOFF_SECONDS"

# ---------------------------------------------------------------------------
# Defaults - overridable via env-vars or CLI flags
# ---------------------------------------------------------------------------
DEFAULT_DB_URL: Final[str] = os.getenv(ENV_DB_URL, settings.database_url)
DEFAULT_MAX_TRIES: Final[int] = int(os.getenv(ENV_MAX_TRIES, "30"))
DEFAULT_INTERVAL: Final[float] = float(os.getenv(ENV_INTERVAL, "2"))
DEFAULT_TIMEOUT: Final[int] = int(os.getenv(ENV_TIMEOUT, "2"))
DEFAULT_MAX_BACKOFF: Final[float] = float(os.getenv(ENV_MAX_BACKOFF, "30"))
DEFAULT_LOG_LEVEL: Final[str] = os.getenv("LOG_LEVEL", settings.log_level).upper()

# ---------------------------------------------------------------------------
# Helpers - sanitising / formatting util functions
# ---------------------------------------------------------------------------
_CRED_RE: Final[re.Pattern[str]] = re.compile(r"://([^:/?#]+):([^@]+)@")
_PWD_RE: Final[re.Pattern[str]] = re.compile(r"(?i)(password|pwd)=([^\s]+)")


def _sanitize(txt: str) -> str:
    """Hide credentials contained in connection strings or driver errors.

    Args:
        txt: Arbitrary text that may contain a DB DSN or ``password=...``
            parameter.

    Returns:
        Same *txt* but with credentials replaced by ``***``.
    """

    redacted = _CRED_RE.sub(r"://\1:***@", txt)
    return _PWD_RE.sub(r"\1=***", redacted)


def _format_target(url: URL) -> str:
    """Return a concise *host[:port]/db* representation for logging.

    Args:
        url: A parsed :class:`sqlalchemy.engine.url.URL` instance.

    Returns:
        Human-readable connection target string suitable for log messages.
    """

    if url.get_backend_name() == "sqlite":
        return url.database or "<memory>"

    host: str = url.host or "localhost"
    port: str = f":{url.port}" if url.port else ""
    db: str = f"/{url.database}" if url.database else ""
    return f"{host}{port}{db}"


# ---------------------------------------------------------------------------
# Public API - *wait_for_db_ready*
# ---------------------------------------------------------------------------


def wait_for_db_ready(
    *,
    database_url: str = DEFAULT_DB_URL,
    max_tries: int = DEFAULT_MAX_TRIES,
    interval: float = DEFAULT_INTERVAL,
    timeout: int = DEFAULT_TIMEOUT,
    max_backoff: float = DEFAULT_MAX_BACKOFF,
    logger: Optional[logging.Logger] = None,
    sync: bool = False,
) -> None:
    """
    Block until the database replies to ``SELECT 1``.

    The helper can be awaited **asynchronously** *or* called in *blocking*
    mode by passing ``sync=True``.

    Uses **exponential backoff with jitter** to prevent thundering herd when
    multiple workers attempt to reconnect simultaneously. The delay between
    attempts doubles each time (capped at max_backoff), with ±25% random jitter.

    Example retry progression with interval=2s, max_backoff=30s:
        Attempt 1: 2s, Attempt 2: 4s, Attempt 3: 8s, Attempt 4: 16s,
        Attempt 5+: 30s (capped), each ±25% jitter

    Args:
        database_url: SQLAlchemy URL to probe. Falls back to ``$DATABASE_URL``
            or the project default (usually an on-disk SQLite file).
        max_tries: Total number of connection attempts before giving up.
        interval: Base delay *in seconds* between attempts. Actual delay uses
            exponential backoff: ``min(interval * 2^(attempt-1), max_backoff)``, then ±25% jitter.
        timeout: Per-attempt connection timeout in seconds (passed to the DB
            driver when supported).
        max_backoff: Maximum backoff delay in seconds (default 30). Jitter is applied
            after this cap, so actual sleep can be ±25% of this value.
        logger: Optional custom :class:`logging.Logger`. If omitted, a default
            one named ``"db_isready"`` is lazily configured.
        sync: When *True*, run in the **current** thread instead of scheduling
            the probe inside an executor. Setting this flag from inside a
            running event-loop will block that loop!

    Raises:
        RuntimeError: If *invalid* parameters are supplied or the database is
            still unavailable after the configured number of attempts.

    Doctest:
    >>> from mcpgateway.utils.db_isready import wait_for_db_ready
    >>> import logging
    >>> class DummyLogger:
    ...     def __init__(self): self.infos = []
    ...     def info(self, msg): self.infos.append(msg)
    ...     def debug(self, msg): pass
    ...     def error(self, msg): pass
    ...     @property
    ...     def handlers(self): return [True]
    >>> import sys
    >>> sys.modules['sqlalchemy'] = type('sqlalchemy', (), {
    ...     'create_engine': lambda *a, **k: type('E', (), {'connect': lambda self: type('C', (), {'execute': lambda self, q: 1, '__enter__': lambda self: self, '__exit__': lambda self, exc_type, exc_val, exc_tb: None})()})(),
    ...     'text': lambda q: q,
    ...     'engine': type('engine', (), {'Engine': object, 'URL': object, 'url': type('url', (), {'make_url': lambda u: type('U', (), {'get_backend_name': lambda self: "sqlite"})()}),}),
    ...     'exc': type('exc', (), {'OperationalError': Exception})
    ... })
    >>> wait_for_db_ready(database_url='sqlite:///./mcp.db', max_tries=1, interval=1, timeout=1, logger=DummyLogger(), sync=True)
    >>> try:
    ...     wait_for_db_ready(database_url='sqlite:///./mcp.db', max_tries=0, interval=1, timeout=1, logger=DummyLogger(), sync=True)
    ... except RuntimeError as e:
    ...     print('error')
    error
    """

    log = logger or logging.getLogger("db_isready")
    if not log.handlers:  # basicConfig **once** - respects *log.setLevel* later
        logging.basicConfig(
            level=getattr(logging, DEFAULT_LOG_LEVEL, logging.INFO),
            format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
            datefmt="%Y-%m-%dT%H:%M:%S",
        )

    if max_tries < 1 or interval <= 0 or timeout <= 0:
        raise RuntimeError("Invalid max_tries / interval / timeout values")

    url_obj: URL = make_url(database_url)
    backend: str = url_obj.get_backend_name()
    target: str = _format_target(url_obj)

    log.info(f"Probing {backend} at {target} (timeout={timeout}s, interval={interval}s, max_tries={max_tries}, max_backoff={max_backoff}s)")

    connect_args: Dict[str, Any] = {}
    if backend.startswith("postgresql"):
        # Most drivers honour this parameter - harmless for others.
        connect_args["connect_timeout"] = timeout

    if backend == "sqlite":
        # SQLite doesn't support pool overflow/timeout parameters
        engine: Engine = create_engine(
            database_url,
            connect_args=connect_args,
        )
    else:
        # Other databases support full pooling configuration
        engine: Engine = create_engine(
            database_url,
            pool_pre_ping=True,
            pool_size=1,
            max_overflow=0,
            connect_args=connect_args,
        )

    def _probe() -> None:  # noqa: D401 - internal helper
        """Inner synchronous probe running in either the current or a thread.

        Returns:
            None - the function exits successfully once the DB answers.

        Raises:
            RuntimeError: Forwarded after exhausting ``max_tries`` attempts.
        """

        start = time.perf_counter()
        for attempt in range(1, max_tries + 1):
            try:
                with engine.connect() as conn:
                    conn.execute(text("SELECT 1"))
                elapsed = time.perf_counter() - start
                log.info(f"Database ready after {elapsed:.2f}s (attempt {attempt})")
                return
            except OperationalError as exc:
                if attempt < max_tries:  # Don't sleep on the last attempt
                    # Exponential backoff: interval * 2^(attempt-1), capped at max_backoff
                    backoff = min(interval * (2 ** (attempt - 1)), max_backoff)
                    # Add jitter (±25%) to prevent thundering herd
                    jitter = backoff * random.uniform(-0.25, 0.25)  # nosec B311
                    sleep_time = max(0.1, backoff + jitter)  # Ensure minimum 0.1s
                    log.debug(f"Attempt {attempt}/{max_tries} failed ({_sanitize(str(exc))}) - retrying in {sleep_time:.1f}s")
                    time.sleep(sleep_time)
                else:
                    log.debug(f"Attempt {attempt}/{max_tries} failed ({_sanitize(str(exc))})")
        raise RuntimeError(f"Database not ready after {max_tries} attempts")

    if sync:
        try:
            _probe()
        finally:
            # Readiness probe should not keep pooled connections open.
            if hasattr(engine, "dispose"):
                engine.dispose()
    else:
        loop = asyncio.get_event_loop()
        try:
            # Off-load to default executor to avoid blocking the event-loop.
            loop.run_until_complete(loop.run_in_executor(None, _probe))
        finally:
            if hasattr(engine, "dispose"):
                engine.dispose()


# ---------------------------------------------------------------------------
# CLI helpers
# ---------------------------------------------------------------------------


def _parse_cli() -> argparse.Namespace:
    """Parse command-line arguments for the *db_isready* CLI wrapper.

    Returns:
        Parsed :class:`argparse.Namespace` holding all CLI options.

    Examples:
        >>> import sys
        >>> # Save original argv
        >>> original_argv = sys.argv
        >>>
        >>> # Test default values
        >>> sys.argv = ['db_isready.py']
        >>> args = _parse_cli()
        >>> args.database_url == DEFAULT_DB_URL
        True
        >>> args.max_tries == DEFAULT_MAX_TRIES
        True
        >>> args.interval == DEFAULT_INTERVAL
        True
        >>> args.timeout == DEFAULT_TIMEOUT
        True
        >>> args.log_level == DEFAULT_LOG_LEVEL
        True

        >>> # Test custom values
        >>> sys.argv = ['db_isready.py', '--database-url', 'postgresql://localhost/test',
        ...             '--max-tries', '5', '--interval', '1.5', '--timeout', '10',
        ...             '--log-level', 'DEBUG']
        >>> args = _parse_cli()
        >>> args.database_url
        'postgresql://localhost/test'
        >>> args.max_tries
        5
        >>> args.interval
        1.5
        >>> args.timeout
        10
        >>> args.log_level
        'DEBUG'

        >>> # Restore original argv
        >>> sys.argv = original_argv
    """

    parser = argparse.ArgumentParser(
        description="Wait until the configured database is ready.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--database-url",
        default=DEFAULT_DB_URL,
        help="SQLAlchemy URL (env DATABASE_URL)",
    )
    parser.add_argument("--max-tries", type=int, default=DEFAULT_MAX_TRIES, help="Maximum connection attempts")
    parser.add_argument("--interval", type=float, default=DEFAULT_INTERVAL, help="Delay between attempts in seconds")
    parser.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT, help="Per-attempt connect timeout in seconds")
    parser.add_argument("--max-backoff", type=float, default=DEFAULT_MAX_BACKOFF, help="Maximum backoff delay in seconds (jitter applied after)")
    parser.add_argument("--log-level", default=DEFAULT_LOG_LEVEL, help="Logging level (DEBUG, INFO, ...)")
    return parser.parse_args()


def main() -> None:  # pragma: no cover
    """CLI entry-point.

    * Parses command-line options.
    * Applies ``--log-level`` to the *db_isready* logger **before** the first
      message is emitted.
    * Delegates the actual probing to :func:`wait_for_db_ready`.
    * Exits with:

        * ``0`` - database became ready.
        * ``1`` - connection attempts exhausted.
        * ``2`` - SQLAlchemy missing (handled on import).
        * ``3`` - invalid parameter combination.
    """
    cli_args = _parse_cli()

    log = logging.getLogger("db_isready")
    log.setLevel(cli_args.log_level.upper())

    try:
        wait_for_db_ready(
            database_url=cli_args.database_url,
            max_tries=cli_args.max_tries,
            interval=cli_args.interval,
            timeout=cli_args.timeout,
            max_backoff=cli_args.max_backoff,
            sync=True,
            logger=log,
        )
    except RuntimeError as exc:
        log.error(f"Database unavailable: {exc}")
        sys.exit(1)

    sys.exit(0)


if __name__ == "__main__":  # pragma: no cover
    main()
