#!/usr/bin/python3
import time
import os
import asyncio
import sys

sys.path.append('/usr/local/lib/python-aioruckus')

from prometheus_client import start_http_server, PROCESS_COLLECTOR, PLATFORM_COLLECTOR, GC_COLLECTOR
from prometheus_client.core import GaugeMetricFamily, CounterMetricFamily, StateSetMetricFamily, InfoMetricFamily, REGISTRY
from prometheus_client.registry import Collector

from aioruckus import AjaxSession


async def collect_async(host, user, password):
  async with AjaxSession.async_create(host, user, password) as session:
    metrics = []
    def reg(metric):
      metrics.append(metric)
      return metric

    base_label_names = ('source',)
    base_label_values = (host,)
    client_label_names = base_label_names + ('ap_mac', 'ap_name', 'bssid', 'channel', 'essid', 'ip', 'mac', 'name', 'vlan', 'radio_type')
    ruckus_client_radio_signal_db = reg(GaugeMetricFamily('ruckus_client_radio_signal_db', 'Client radio signal in dB', labels=client_label_names))
    ruckus_client_noise_db = reg(GaugeMetricFamily('ruckus_client_noise_db', 'Client noise floor in dB', labels=client_label_names))
    ruckus_client_receive_bytes_total = reg(CounterMetricFamily('ruckus_client_receive_bytes_total', 'Client total received bytes', labels=client_label_names))
    ruckus_client_receive_packets_total = reg(CounterMetricFamily('ruckus_client_receive_packets_total', 'Client total received packets', labels=client_label_names))
    ruckus_client_transmit_bytes_total = reg(CounterMetricFamily('ruckus_client_transmit_bytes_total', 'Client total transmitted bytes', labels=client_label_names))
    ruckus_client_transmit_packets_total = reg(CounterMetricFamily('ruckus_client_transmit_packets_total', 'Client total transmitted packets', labels=client_label_names))

    device_info_label_names = base_label_names + ('name', 'mac', 'model', 'serial', 'ip', 'version')
    ruckus_device_info = reg(InfoMetricFamily('ruckus_device', 'Device info', labels=device_info_label_names))
    device_label_names = base_label_names + ('name', 'mac')
    ruckus_device_uptime_seconds = reg(CounterMetricFamily('ruckus_device_uptime_seconds', 'Device uptime in seconds', labels=device_label_names))

    clients = await session.api.get_active_clients(interval_stats=True)
    for client in clients:
      client_label_values = base_label_values + (client['ap'], client['ap-name'], client['vap-mac'], client['channel'], client['ssid'], client.get('ip', ''), client['mac'], client.get('hostname', ''), client.get('vlan', ''), client['radio-type'])
      ruckus_client_radio_signal_db.add_metric(client_label_values, float(client['received-signal-strength']))
      ruckus_client_noise_db.add_metric(client_label_values, float(client['noise-floor']))
      ruckus_client_receive_bytes_total.add_metric(client_label_values, float(client['total-rx-bytes']))
      ruckus_client_receive_packets_total.add_metric(client_label_values, float(client['total-rx-pkts']))
      ruckus_client_transmit_bytes_total.add_metric(client_label_values, float(client['total-tx-bytes']))
      ruckus_client_transmit_packets_total.add_metric(client_label_values, float(client['total-tx-pkts']))

    devices = await session.api.get_ap_stats()
    for device in devices:
      device_info_label_values = base_label_values + (device['ap-name'], device['mac'], device['model'], device['serial-number'], device['ip'], device['firmware-version'])
      ruckus_device_info.add_metric(device_info_label_values, {})
      device_label_values = base_label_values + (device['ap-name'], device['mac'])
      ruckus_device_uptime_seconds.add_metric(device_label_values, float(device['uptime']))

    return metrics


class CustomCollector(Collector):
    def collect(self):
      for host in os.environ['RUCKUS_EXPORTER_HOSTS'].split(' '):
        yield from asyncio.run(collect_async(host, os.environ['RUCKUS_EXPORTER_USER'], os.environ['RUCKUS_EXPORTER_PASSWORD']))


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=int(os.environ['RUCKUS_EXPORTER_PORT']))
  while True:
    time.sleep(100)