# 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.

"""Views for the server application's OpenMetrics statistics."""

import os
from collections.abc import Callable, Generator, Mapping
from enum import StrEnum
from typing import Any

from celery import current_task
from django.utils.encoding import smart_str
from prometheus_client import Metric
from prometheus_client import exposition as prometheus
from prometheus_client.multiprocess import MultiProcessCollector
from prometheus_client.openmetrics import exposition as openmetrics
from prometheus_client.registry import CollectorRegistry, REGISTRY
from rest_framework import renderers, status
from rest_framework.request import Request
from rest_framework.response import Response

from debusine.db.metrics import (
    BadMetricType,
    Collector,
    MetricNotFound,
    emit_metric,
)
from debusine.db.models import Token
from debusine.db.models.permissions import anonymous_access
from debusine.server.exceptions import DebusineAPIException
from debusine.server.open_metrics import DebusineCollector
from debusine.server.serializers import EmitMetricSerializer
from debusine.server.views.base import BaseAPIView
from debusine.server.views.rest import IsWorkerTokenAuthenticated
from debusine.tasks.models import WorkerType


class FixedMultiProcessCollector(MultiProcessCollector):
    """
    Multi-process collector with fixed histogram sample ordering.

    Works around https://github.com/prometheus/client_python/issues/1147.
    """

    def collect(self) -> Generator[Metric]:
        """Reorder collected samples to avoid interleaving metrics."""
        for metric in super().collect():  # type: ignore[no-untyped-call]
            metric.samples.sort(
                key=lambda sample: (
                    tuple(
                        (key, value)
                        for key, value in sample.labels.items()
                        if key != "le"
                    ),
                    sample.name,
                )
            )
            yield metric


class FilteredCollector:
    """A metric collector that only exposes some metrics from a parent."""

    def __init__(
        self, parent: Collector, accept: Callable[[Metric], bool]
    ) -> None:
        """Initialize the collector."""
        self.parent = parent
        self.accept = accept

    def collect(self) -> Generator[Metric]:
        """Yield each accepted metric."""
        for metric in self.parent.collect():
            if self.accept(metric):
                yield metric


def extract_media_type(content_type: str) -> str:
    """Remove the charset from a fully specified Content-Type."""
    return content_type.split(";", 1)[0].strip()


class RendererFormats(StrEnum):
    """Supported output formats."""

    PROMETHEUS = "prometheus"
    OPENMETRICS = "openmetrics"


class BaseTextRenderer(renderers.BaseRenderer):
    """Base class for returning text responses."""

    def render(
        self,
        data: str,
        accepted_media_type: str | None = None,
        renderer_context: Mapping[str, Any] | None = None,
    ) -> str:
        """Render data to bytes."""
        accepted_media_type, renderer_context  # fake usage for vulture
        assert self.charset
        return smart_str(data, encoding=self.charset)


class PrometheusRenderer(BaseTextRenderer):
    """Renderer for Prometheus native metrics."""

    media_type = extract_media_type(prometheus.CONTENT_TYPE_LATEST)
    format = RendererFormats.PROMETHEUS


class OpenMetricsRenderer(BaseTextRenderer):
    """Renderer for OpenMetrics format metrics."""

    media_type = extract_media_type(openmetrics.CONTENT_TYPE_LATEST)
    format = RendererFormats.OPENMETRICS


class OpenMetricsView(BaseAPIView):
    """View used to get the OpenMetrics statistics."""

    renderer_classes = [PrometheusRenderer, OpenMetricsRenderer]

    @property
    def _is_multiprocess(self) -> bool:
        """Return True if we are using multi-process collection."""
        return "PROMETHEUS_MULTIPROC_DIR" in os.environ

    def get(self, request: Request) -> Response:
        """Return OpenMetrics statistics."""
        self.enforce(anonymous_access)
        # Registry of metrics collected separately by each process.
        non_database_registry: Collector
        if self._is_multiprocess:
            non_database_registry = CollectorRegistry(auto_describe=True)
            FixedMultiProcessCollector(non_database_registry)  # type: ignore[no-untyped-call]
        else:
            # To accommodate people running `./manage.py runserver` by hand,
            # if multi-process collection isn't set up, just use the global
            # registry.
            non_database_registry = REGISTRY
        non_database_registry = FilteredCollector(
            non_database_registry,
            lambda metric: metric.name.startswith(("django_", "debusine_")),
        )

        # Registry of metrics based on the database.
        database_registry = CollectorRegistry(auto_describe=True)
        database_registry.register(DebusineCollector())

        match request.accepted_renderer.format:
            case RendererFormats.PROMETHEUS:
                data = prometheus.generate_latest(non_database_registry)
                assert not data or data.endswith(b"\n")
                data += prometheus.generate_latest(database_registry)
                content_type = prometheus.CONTENT_TYPE_LATEST
            case RendererFormats.OPENMETRICS:
                data = openmetrics.generate_latest(non_database_registry)  # type: ignore[no-untyped-call]
                data = data.removesuffix(b"# EOF\n")
                assert not data or data.endswith(b"\n")
                data += openmetrics.generate_latest(database_registry)  # type: ignore[no-untyped-call]
                content_type = openmetrics.CONTENT_TYPE_LATEST
            case _:  # pragma: no cover
                raise AssertionError("An unexpected renderer was requested")
        return Response(
            data, status=status.HTTP_200_OK, content_type=content_type
        )


class OpenMetricsEmitView(BaseAPIView):
    """Emit an additional metric from a separate worker."""

    permission_classes = [IsWorkerTokenAuthenticated]

    def post(self, request: Request) -> Response:
        """Emit a metric from a worker."""
        serializer = EmitMetricSerializer(data=request.data)
        serializer.is_valid(raise_exception=True)

        # Already validated by IsWorkerTokenAuthenticated
        assert isinstance(self.request.auth, Token)
        assert self.request.auth.worker
        if self.request.auth.worker.worker_type not in (
            WorkerType.CELERY,
            WorkerType.SCHEDULER,
        ):
            raise DebusineAPIException(
                title="Error",
                detail=(
                    f"Endpoint is not accessible to "
                    f"{self.request.auth.worker.worker_type} workers"
                ),
                status_code=status.HTTP_403_FORBIDDEN,
            )
        self.enforce(anonymous_access)  # restricted to workers

        metric_type = serializer.validated_data["metric_type"]
        name = serializer.validated_data["name"]
        labels = serializer.validated_data["labels"]
        value = serializer.validated_data["value"]

        # Make sure emit_metric doesn't recursively call back to this view.
        assert not current_task

        try:
            emit_metric(
                metric_type=metric_type, name=name, labels=labels, value=value
            )
        except MetricNotFound as e:
            raise DebusineAPIException(
                title="Metric not found",
                detail=str(e),
                status_code=status.HTTP_404_NOT_FOUND,
            )
        except BadMetricType as e:
            raise DebusineAPIException(
                title="Bad metric type",
                detail=str(e),
                status_code=status.HTTP_400_BAD_REQUEST,
            )

        return Response(status=status.HTTP_204_NO_CONTENT)
