# Copyright © The Debusine Developers
# See the AUTHORS file at the top-level directory of this distribution
#
# This file is part of Debusine. It is subject to the license terms
# in the LICENSE file found in the top-level directory of this
# distribution. No part of Debusine, including this file, may be copied,
# modified, propagated, or distributed except according to the terms
# contained in the LICENSE file.

"""Utilities for emitting metrics."""

from collections.abc import Iterable
from logging import getLogger
from typing import Literal, Protocol, assert_never, override

from celery import current_task
from prometheus_client import Counter, Gauge, Histogram, Metric, Summary
from prometheus_client.registry import REGISTRY

from debusine.project.celery import get_celery_token

logger = getLogger(__name__)


class Collector(Protocol):
    """An object that supports metric collection."""

    def collect(self) -> Iterable[Metric]:
        """Collect metrics."""


class MetricNotFound(Exception):
    """The named metric was not found."""

    def __init__(self, name: str) -> None:
        """Construct the exception."""
        self.name = name

    @override
    def __str__(self) -> str:
        """Render the exception as a string."""
        return f"Metric {self.name!r} not found"


class BadMetricType(Exception):
    """The named metric has an unexpected type."""

    def __init__(self, name: str, metric_class_name: str) -> None:
        """Construct the exception."""
        self.name = name
        self.metric_class_name = metric_class_name

    @override
    def __str__(self) -> str:
        """Render the exception as a string."""
        return f"Metric {self.name!r} is not a {self.metric_class_name}"


def _check_metric_type[C: Collector](
    name: str, metric: Collector, metric_class: type[C]
) -> C:
    if not isinstance(metric, metric_class):
        raise BadMetricType(name, metric_class.__name__)
    return metric


def emit_metric(
    *,
    metric_type: Literal["counter", "gauge", "summary", "histogram"],
    name: str,
    labels: dict[str, str],
    value: float,
) -> None:
    """
    Emit a metric.

    If we're running in a Celery task, ask the server to emit the metric.
    Otherwise, do so directly.
    """
    if current_task and (celery_token := get_celery_token()) is not None:
        try:
            celery_token.client.emit_metric(
                metric_type=metric_type, name=name, labels=labels, value=value
            )
        except Exception as e:
            logger.warning(  # noqa: G200
                "Failed to report metric back to the server: %s", e
            )
    else:
        metric = REGISTRY._names_to_collectors.get(name)
        if metric is None:
            raise MetricNotFound(name)

        match metric_type:
            case "counter":
                _check_metric_type(name, metric, Counter).labels(**labels).inc(
                    value
                )
            case "gauge":
                _check_metric_type(name, metric, Gauge).labels(**labels).inc(
                    value
                )
            case "summary":
                _check_metric_type(name, metric, Summary).labels(
                    **labels
                ).observe(value)
            case "histogram":
                _check_metric_type(name, metric, Histogram).labels(
                    **labels
                ).observe(value)
            case _ as unreachable:
                assert_never(unreachable)
