from typing import Callable
from prometheus_fastapi_instrumentator.metrics import Info
from prometheus_client import Gauge
from sqlalchemy import select, func

from app.core.db import async_session_maker
from app.models.transaction import Transaction
from app.models.user import User


def number_of_users() -> Callable[[Info], None]:
    METRIC = Gauge(
        "number_of_users_total",
        "Number of registered users.",
        labelnames=("users",)
    )

    async def instrumentation(info: Info) -> None:
        try:
            async with async_session_maker() as session:
                result = await session.execute(select(func.count(User.id)))
                user_count = result.scalar_one() or 0
        except Exception:
            # In case of DB errors, avoid crashing metrics endpoint
            user_count = 0
        METRIC.labels(users="total").set(user_count)

    return instrumentation


def number_of_transactions() -> Callable[[Info], None]:
    METRIC = Gauge(
        "number_of_transactions_total",
        "Number of transactions stored.",
        labelnames=("transactions",)
    )

    async def instrumentation(info: Info) -> None:
        try:
            async with async_session_maker() as session:
                result = await session.execute(select(func.count()).select_from(Transaction))
                transaction_count = result.scalar_one() or 0
        except Exception:
            # In case of DB errors, avoid crashing metrics endpoint
            transaction_count = 0
        METRIC.labels(transactions="total").set(transaction_count)

    return instrumentation