from functools import lru_cache

from funcy import identity
from prometheus_client import CollectorRegistry, Summary

from pyinfra.server.operation_dispatcher import OperationDispatcher


class OperationDispatcherMonitoringDecorator:
    def __init__(self, operation_dispatcher: OperationDispatcher, naming_policy=identity):
        self.operation_dispatcher = operation_dispatcher
        self.operation2metric = {}
        self.naming_policy = naming_policy

    @property
    @lru_cache(maxsize=None)
    def registry(self):
        return CollectorRegistry(auto_describe=True)

    def make_summary_instance(self, op: str):
        return Summary(f"{self.naming_policy(op)}_seconds", f"Time spent on {op}.", registry=self.registry)

    def submit(self, operation, request):
        return self.operation_dispatcher.submit(operation, request)

    def pickup(self, operation):
        with self.get_monitor(operation):
            return self.operation_dispatcher.pickup(operation)

    def get_monitor(self, operation):
        monitor = self.operation2metric.get(operation, None) or self.register_operation(operation)
        return monitor.time()

    def register_operation(self, operation):
        summary = self.make_summary_instance(operation)
        self.operation2metric[operation] = summary
        return summary