from __future__ import annotations

from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass
import time
import weakref
from collections.abc import Iterator

from sqlalchemy import event
from sqlalchemy.engine import Connection, Engine
from sqlalchemy.engine.interfaces import ExecutionContext


@dataclass(slots=True)
class QueryMetrics:
    query_count: int = 0
    total_ms: float = 0.0
    slowest_ms: float = 0.0
    error_count: int = 0

    def record(self, elapsed_ms: float, *, failed: bool = False) -> None:
        self.query_count += 1
        self.total_ms += elapsed_ms
        self.slowest_ms = max(self.slowest_ms, elapsed_ms)
        if failed:
            self.error_count += 1


_current_metrics: ContextVar[QueryMetrics | None] = ContextVar("govoplan_query_metrics", default=None)
_instrumented_engines: weakref.WeakSet[Engine] = weakref.WeakSet()
_START_STACK_KEY = "govoplan_query_metric_starts"


@contextmanager
def collect_query_metrics() -> Iterator[QueryMetrics]:
    metrics = QueryMetrics()
    token = _current_metrics.set(metrics)
    try:
        yield metrics
    finally:
        _current_metrics.reset(token)


def current_query_metrics() -> QueryMetrics | None:
    return _current_metrics.get()


def instrument_engine(engine: Engine) -> Engine:
    if engine in _instrumented_engines:
        return engine

    event.listen(engine, "before_cursor_execute", _before_cursor_execute)
    event.listen(engine, "after_cursor_execute", _after_cursor_execute)
    event.listen(engine, "handle_error", _handle_error)
    _instrumented_engines.add(engine)
    return engine


def _before_cursor_execute(
    conn: Connection,
    _cursor: object,
    _statement: str,
    _parameters: object,
    _context: ExecutionContext,
    _executemany: bool,
) -> None:
    if current_query_metrics() is None:
        return
    starts = conn.info.setdefault(_START_STACK_KEY, [])
    starts.append(time.perf_counter())


def _after_cursor_execute(
    conn: Connection,
    _cursor: object,
    _statement: str,
    _parameters: object,
    _context: ExecutionContext,
    _executemany: bool,
) -> None:
    _record_elapsed(conn, failed=False)


def _handle_error(exception_context: object) -> None:
    conn = getattr(exception_context, "connection", None)
    if isinstance(conn, Connection):
        _record_elapsed(conn, failed=True)


def _record_elapsed(conn: Connection, *, failed: bool) -> None:
    metrics = current_query_metrics()
    if metrics is None:
        return
    starts = conn.info.get(_START_STACK_KEY)
    if not starts:
        return
    started_at = starts.pop()
    elapsed_ms = (time.perf_counter() - started_at) * 1000
    metrics.record(elapsed_ms, failed=failed)
