#  Copyright (C) 2024
#  ABM, Moscow
#
#  UNPUBLISHED PROPRIETARY MATERIAL.
#  ALL RIGHTS RESERVED.
#
#  Authors: Vasya Svintsov <v.svintsov@techokert.ru>
import re
from dataclasses import dataclass
from typing import Awaitable, Final

import aiohttp
from aiohttp import web, hdrs
from http_tools import Answer, HttpStatusCode, HttpServer, IncomingRequest
from http_tools.mime_types import ContentType

from .monitored_components.monitored_client_session import MonitoredClientSession
from .metrics_registry import MetricsRegistry
from .request_id_exemplar import get_request_id_exemplar


OPENMETRICS_CONTENT_TYPE: Final[str] = 'application/openmetrics-text'  # TODO need move to http-tools-abm


class PrometheusController(MetricsRegistry):
    @dataclass
    class Config(MetricsRegistry.Config):
        metrics_path = '/metrics'
        liveness_path = '/liveness'

    @dataclass
    class Context:
        http_server: HttpServer

    def __init__(self, config: Config, context: Context) -> None:
        super().__init__(config)
        self.config = config
        self.context = context
        self.context.http_server.register_handler(self.config.metrics_path, self.get_metrics)
        self.context.http_server.register_handler(self.config.liveness_path, self.check_liveness)
        self.context.http_server.middlewares.append(self.incoming_http_requests_middleware)

    def get_metrics(self, incoming_request: IncomingRequest) -> Answer:
        if (accept_hdr := incoming_request.metadata.get_header(hdrs.ACCEPT)) and OPENMETRICS_CONTENT_TYPE in accept_hdr:
            return Answer(self.render_metrics(openmetrics_format=True), HttpStatusCode.OK, OPENMETRICS_CONTENT_TYPE)
        return Answer(self.render_metrics(), HttpStatusCode.OK, ContentType.Text)

    @staticmethod
    def check_liveness(_: IncomingRequest) -> Answer:
        return Answer(None, HttpStatusCode.OK)

    def get_monitored_client_session(
            self, url_pattern_to_replacer: dict[re.Pattern | str, str] | None = None, **kwargs
    ) -> MonitoredClientSession:
        return MonitoredClientSession(self, url_pattern_to_replacer=url_pattern_to_replacer, **kwargs)

    @web.middleware
    async def incoming_http_requests_middleware(
            self, request: aiohttp.web.Request, handler: aiohttp.typedefs.Handler, *args, **kwargs
    ) -> web.StreamResponse | Awaitable[web.StreamResponse]:
        exemplar = get_request_id_exemplar()
        labels = {'method': request.method, 'path': request.path}
        self.counter('http_server__accepted_requests', labels).inc(exemplar=exemplar)
        with self.track('http_server__requests', labels, exemplar=exemplar) as tracker:
            response = await handler(request)
            tracker.labels['status_code'] = response.status
        self.counter('http_server__incoming_traffic', labels).inc(request.content_length or 0, exemplar=exemplar)
        self.counter('http_server__outgoing_traffic', labels).inc(response.content_length or 0, exemplar=exemplar)
        return response
