"""

    AUTHOR:

        Khushal P Soonderji
        Sharvil J Daiya

    DATE:

        Saturday, 28th Sept., 2024

    OBJECTIVE:

        To have one place from where several metrics are measured using easy to use context managers.

    REFERENCES:

        N/A

    DOWNLOADS:

        N/A

"""


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                   IMPORT                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# To make sibling directories accessible for imports:
import sys
sys.path.append(".")
sys.path.append("..")

# To measure time:
import time

# To capture metrics for Prometheus:
from prometheus_client import Counter, Summary, Gauge


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                           MACROS / ONE-TIME INIT                                              ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# Define the Prometheus metrics.
# FOR THE MICROSERVICE AS A WHOLE:
WORKER_COUNT = Counter(
    name = "ms_workers_active_total",
    documentation = "The number of threads for the microservice being monitored.",
    labelnames = ["project_name", "service_name", "host_name"]
)

# Define the Prometheus metrics.
# FOR INDIVIDUAL API ENDPOINTS:
REQUEST_LATENCY = Summary(
    name = "http_request_latency_seconds",
    documentation = "Latency of HTTP requests in seconds.",
    labelnames = ["method", "endpoint", "http_status"]
)
TOTAL_REQUEST_COUNT = Counter(
    name = "http_requests_total",
    documentation = "Total HTTP requests.",
    labelnames = ["method", "endpoint", "http_status"]
)
LIVE_REQUEST_COUNT = Gauge(
    name = "http_requests_live_total",
    documentation = "To check if an API endpoint is being served right now.",
    labelnames = ["endpoint"]
)


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                VARIABLES                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                  CLASSES                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


class MetricsAPI:

    """
    Use this class through its context manager to automatically measure all the metrics in one place.
    This was originally created to measure the performance of API endpoints made in Quart, but it should work with
    other frameworks as well.
    """

    def __init__(
            self,
            method = None,
            endpoint = None,
            raise_exception = False
    ):

        # Make provisions for things to note.
        # NOTE: THESE MUST BE SET FROM OUTSIDE:
        self.method = method
        self.endpoint = endpoint
        self.http_code = None
        self.__raise_exception = raise_exception

    async def __aenter__(self):

        # Note down the start time immediately:
        self.__start_ts = time.perf_counter()
        self.__cpu_start_ts = time.process_time()

        # Note down the metrics:
        LIVE_REQUEST_COUNT.labels(self.endpoint).inc(1)

        # Setup done:
        return self

    async def __aexit__(self, exc_type, exc_value, traceback):

        # Note down the metrics:
        LIVE_REQUEST_COUNT.labels(self.endpoint).dec(1)

        TOTAL_REQUEST_COUNT.labels(
            self.method,
            self.endpoint,
            self.http_code
        ).inc()

        REQUEST_LATENCY.labels(
            self.method,
            self.endpoint,
            self.http_code
        ).observe(time.perf_counter() - self.__start_ts)

        # Handle the exception as per the user's preference:
        return False if self.__raise_exception else True


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                  FUNCTIONS                                                    ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                               MAIN PROGRAM                                                    ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


if __name__ == "__main__":

    import asyncio
    import random
    from prometheus_client import generate_latest

    async def simulate_endpoint():

        async with MetricsAPI(
                method = random.choice(["GET", "POST"]),
                endpoint = f"https://my.domain.com/api/{random.choice([0, 1, 2, 3])}"
        ) as metrics:

            # Simulate some action on some endpoint:
            await asyncio.sleep(1.0)

            # Note down the values:
            metrics.http_code = 200

    async def main():

        print("Simulating endpoints...")
        tasks = [simulate_endpoint() for _ in range(250)]
        await asyncio.gather(*tasks)
        print("Done!")

        print("METRICS:")
        print(generate_latest().decode())

    asyncio.run(main())