import logging
import time
from typing import Optional, Iterable, Dict, Any

from opentelemetry.metrics import Meter, Observation, CallbackOptions

logger = logging.getLogger(__name__)

_collector: Optional["DatabaseMetricsCollector"] = None

_STATS_QUERY = """
SELECT
    count(*) AS total,
    count(*) FILTER (WHERE state = 'active') AS active,
    count(*) FILTER (WHERE state = 'idle') AS idle,
    count(*) FILTER (WHERE state = 'idle in transaction') AS idle_in_tx,
    count(*) FILTER (WHERE wait_event IS NOT NULL AND state = 'active') AS waiting,
    (SELECT setting::int FROM pg_settings WHERE name = 'max_connections') AS max_conn,
    pg_database_size(current_database()) AS db_size
FROM pg_stat_activity
WHERE datname = current_database()
"""


class DatabaseMetricsCollector:
    """
    Collects database connection metrics from Django's connection handler
    and PostgreSQL's pg_stat_activity.

    All PostgreSQL stats are fetched in a **single query** per database alias,
    then cached for the duration of one collection cycle (default 60s) to
    avoid redundant round-trips.
    """

    def __init__(self, meter: Meter, db_aliases: Optional[list] = None):
        self._meter = meter
        self._db_aliases = db_aliases
        self._cache: Dict[str, Dict[str, Any]] = {}
        self._cache_ts: float = 0
        self._cache_ttl: float = 10.0
        self._register_metrics()
        logger.info("Database metrics collector initialised")

    def _get_aliases(self) -> list:
        if self._db_aliases:
            return self._db_aliases
        try:
            from django.db import connections
            return list(connections)
        except Exception:
            return []

    # ------------------------------------------------------------------
    # Single-query snapshot, cached per collection cycle
    # ------------------------------------------------------------------

    def _refresh_cache(self) -> None:
        now = time.monotonic()
        if now - self._cache_ts < self._cache_ttl:
            return

        new_cache: Dict[str, Dict[str, Any]] = {}
        for alias in self._get_aliases():
            row = self._query_pg(alias, _STATS_QUERY)
            if row:
                total, active, idle, idle_in_tx, waiting, max_conn, db_size = row[0]
                new_cache[alias] = {
                    "total": total,
                    "active": active,
                    "idle": idle,
                    "idle_in_tx": idle_in_tx,
                    "waiting": waiting,
                    "max_conn": max_conn,
                    "db_size": db_size,
                    "utilization": round((total / max_conn) * 100, 2) if max_conn else 0,
                }
        self._cache = new_cache
        self._cache_ts = now

    def _stat(self, alias: str, key: str):
        return self._cache.get(alias, {}).get(key)

    # ------------------------------------------------------------------
    # Low-level query helper
    # ------------------------------------------------------------------

    def _query_pg(self, alias: str, sql: str):
        try:
            from django.db import connections

            conn = connections[alias]
            if conn.connection is None or not conn.is_usable():
                conn.ensure_connection()
            with conn.connection.cursor() as cur:
                cur.execute(sql)
                return cur.fetchall()
        except Exception:
            return None

    # ------------------------------------------------------------------
    # Metric registration
    # ------------------------------------------------------------------

    def _register_metrics(self):
        self._meter.create_observable_gauge(
            name="db.django.connections.open",
            callbacks=[self._observe_django_connections],
            unit="{connections}",
            description="Number of open Django database connections in this process",
        )

        self._meter.create_observable_gauge(
            name="db.postgresql.connections.total",
            callbacks=[self._observe_pg_total],
            unit="{connections}",
            description="Total connections to the PostgreSQL server",
        )

        self._meter.create_observable_gauge(
            name="db.postgresql.connections.active",
            callbacks=[self._observe_pg_active],
            unit="{connections}",
            description="Active (running a query) connections",
        )

        self._meter.create_observable_gauge(
            name="db.postgresql.connections.idle",
            callbacks=[self._observe_pg_idle],
            unit="{connections}",
            description="Idle connections on the PostgreSQL server",
        )

        self._meter.create_observable_gauge(
            name="db.postgresql.connections.idle_in_transaction",
            callbacks=[self._observe_pg_idle_in_tx],
            unit="{connections}",
            description="Connections stuck in 'idle in transaction' state",
        )

        self._meter.create_observable_gauge(
            name="db.postgresql.connections.waiting",
            callbacks=[self._observe_pg_waiting],
            unit="{connections}",
            description="Connections waiting on a lock",
        )

        self._meter.create_observable_gauge(
            name="db.postgresql.connections.max",
            callbacks=[self._observe_pg_max],
            unit="{connections}",
            description="PostgreSQL max_connections setting",
        )

        self._meter.create_observable_gauge(
            name="db.postgresql.connections.utilization",
            callbacks=[self._observe_pg_utilization],
            unit="%",
            description="Percentage of max_connections currently in use",
        )

        self._meter.create_observable_gauge(
            name="db.postgresql.database.size",
            callbacks=[self._observe_pg_db_size],
            unit="By",
            description="Database size in bytes",
        )

    # ------------------------------------------------------------------
    # Django-level connections (no DB query needed)
    # ------------------------------------------------------------------

    def _observe_django_connections(
        self, options: CallbackOptions
    ) -> Iterable[Observation]:
        try:
            from django.db import connections

            for alias in self._get_aliases():
                conn = connections[alias]
                is_open = conn.connection is not None and conn.is_usable()
                yield Observation(1 if is_open else 0, {"db.alias": alias})
        except Exception:
            pass

    # ------------------------------------------------------------------
    # PostgreSQL stats — each callback reads from the shared cache
    # ------------------------------------------------------------------

    def _observe_pg_total(self, options: CallbackOptions) -> Iterable[Observation]:
        self._refresh_cache()
        for alias in self._cache:
            val = self._stat(alias, "total")
            if val is not None:
                yield Observation(val, {"db.alias": alias})

    def _observe_pg_active(self, options: CallbackOptions) -> Iterable[Observation]:
        self._refresh_cache()
        for alias in self._cache:
            val = self._stat(alias, "active")
            if val is not None:
                yield Observation(val, {"db.alias": alias})

    def _observe_pg_idle(self, options: CallbackOptions) -> Iterable[Observation]:
        self._refresh_cache()
        for alias in self._cache:
            val = self._stat(alias, "idle")
            if val is not None:
                yield Observation(val, {"db.alias": alias})

    def _observe_pg_idle_in_tx(self, options: CallbackOptions) -> Iterable[Observation]:
        self._refresh_cache()
        for alias in self._cache:
            val = self._stat(alias, "idle_in_tx")
            if val is not None:
                yield Observation(val, {"db.alias": alias})

    def _observe_pg_waiting(self, options: CallbackOptions) -> Iterable[Observation]:
        self._refresh_cache()
        for alias in self._cache:
            val = self._stat(alias, "waiting")
            if val is not None:
                yield Observation(val, {"db.alias": alias})

    def _observe_pg_max(self, options: CallbackOptions) -> Iterable[Observation]:
        self._refresh_cache()
        for alias in self._cache:
            val = self._stat(alias, "max_conn")
            if val is not None:
                yield Observation(val, {"db.alias": alias})

    def _observe_pg_utilization(self, options: CallbackOptions) -> Iterable[Observation]:
        self._refresh_cache()
        for alias in self._cache:
            val = self._stat(alias, "utilization")
            if val is not None:
                yield Observation(val, {"db.alias": alias})

    def _observe_pg_db_size(self, options: CallbackOptions) -> Iterable[Observation]:
        self._refresh_cache()
        for alias in self._cache:
            val = self._stat(alias, "db_size")
            if val is not None:
                yield Observation(val, {"db.alias": alias})


def start_db_metrics(
    meter: Meter,
    db_aliases: Optional[list] = None,
) -> "DatabaseMetricsCollector":
    """
    Create and start the database metrics collector.

    Args:
        meter: OpenTelemetry Meter to register instruments on.
        db_aliases: List of Django DB aliases to monitor. Defaults to all
            configured aliases.

    Returns:
        The initialised DatabaseMetricsCollector instance.
    """
    global _collector
    if _collector is not None:
        logger.debug("Database metrics collector already running")
        return _collector

    _collector = DatabaseMetricsCollector(meter, db_aliases=db_aliases)
    return _collector


def get_db_metrics_collector() -> Optional["DatabaseMetricsCollector"]:
    """Return the active collector, or None if not started."""
    return _collector
