#!/usr/bin/python3
import time
import subprocess
import json
from prometheus_client import start_http_server, PROCESS_COLLECTOR, PLATFORM_COLLECTOR, GC_COLLECTOR
from prometheus_client.core import GaugeMetricFamily, CounterMetricFamily, REGISTRY
from prometheus_client.registry import Collector

class CustomCollector(Collector):
    def collect(self):
        service_status = json.loads(subprocess.run(['doveadm', '-f', 'json', 'service', 'status'], check=True, capture_output=True).stdout.decode())
        process_status = json.loads(subprocess.run(['doveadm', '-f', 'json', 'process', 'status'], check=True, capture_output=True).stdout.decode())
        dovecot_processes_count = GaugeMetricFamily('dovecot_processes_count', 'Number of running processes per service', labels=['service'])
        dovecot_processes_total = CounterMetricFamily('dovecot_processes_total', 'Total count of started processes per service', labels=['service'])
        dovecot_processes_limit = GaugeMetricFamily('dovecot_processes_limit', 'Process limit per service', labels=['service'])
        dovecot_clients_count = GaugeMetricFamily('dovecot_clients_count', 'Number of connected clients per service', labels=['service'])
        dovecot_clients_total = CounterMetricFamily('dovecot_clients_total', 'Total count of clients per service', labels=['service'])
        dovecot_clients_per_process_limit = GaugeMetricFamily('dovecot_clients_per_process_limit', 'Client limit per process', labels=['service'])
        for line in service_status:
          dovecot_processes_count.add_metric([line['name']], int(line['process_count']))
          dovecot_processes_total.add_metric([line['name']], int(line['process_total']))
          dovecot_processes_limit.add_metric([line['name']], int(line['process_limit']))
          dovecot_clients_count.add_metric(
            [line['name']],
            sum(int(line['client_limit'])-int(proc['available_count']) for proc in process_status if proc['name'] == line['name'])
          )
          dovecot_clients_total.add_metric(
            [line['name']],
            sum(int(proc['total_count']) for proc in process_status if proc['name'] == line['name'])
          )
          dovecot_clients_per_process_limit.add_metric([line['name']], int(line['client_limit']))
        yield dovecot_processes_count
        yield dovecot_processes_total
        yield dovecot_processes_limit
        yield dovecot_clients_count
        yield dovecot_clients_total
        yield dovecot_clients_per_process_limit

REGISTRY.unregister(PROCESS_COLLECTOR)
REGISTRY.unregister(PLATFORM_COLLECTOR)
REGISTRY.unregister(GC_COLLECTOR)
REGISTRY.register(CustomCollector())

if __name__ == '__main__':
  start_http_server(addr='127.0.0.1', port=9162)
  while True:
    time.sleep(100)