from flask import Flask, request
from prometheus_client import Counter, Gauge, Info, start_http_server
import json
app = Flask(__name__)


class TelemetryInfo:
    def __init__(self, data):
        self.node_id = data['chain']['node_id']
        self.num_peers = data['chain']['num_peers']
        self.status = data['chain']['status']
        self.cpu_usage = data['system']['cpu_usage']
        self.memory_usage = data['system']['memory_usage']
        self.bandwidth_upload = data['system']['bandwidth_upload']
        self.bandwidth_download = data['system']['bandwidth_download']


class Node:
    def __init__(self, node_id):
        def _(s): return f"{s}_{node_id}"

        self.updates = Counter(_('updates'), 'Total updates received')
        self.num_peers = Gauge(_('num_peers'), 'Number of active peers')
        self.cpu_usage = Gauge(_('cpu_usage'), "CPU usage")
        self.memory_usage = Gauge(_('memory_usage'), "Memory usage")
        self.bandwidth_upload = Gauge(
            _('bandwidth_upload'), "Bandwidth upload in bytes")
        self.bandwidth_download = Gauge(
            _('bandwidth_download'), "Bandwidth download in bytes"
        )
        self.status = Info(_('status'), "Node status")

    def submit(self, info: TelemetryInfo):
        self.updates.inc()
        self.num_peers.set(info.num_peers)
        self.cpu_usage.set(info.cpu_usage)
        self.memory_usage.set(info.memory_usage)
        self.bandwidth_upload.set(info.bandwidth_upload)
        self.bandwidth_download.set(info.bandwidth_download)
        self.status.info({'status': info.status})


class Prometheus:
    def __init__(self):
        self.dic = {}

    def get(self, node_id):
        if not node_id in self.dic:
            self.dic[node_id] = Node(node_id)
        return self.dic[node_id]

    def submit(self, data):
        self.get(data.node_id).submit(data)


prom = Prometheus()


@app.route('/', methods=['POST'])
def hello_world():
    data = json.loads(request.data.decode())
    info = TelemetryInfo(data)
    print("Received data from:", info.node_id)
    prom.submit(info)
    return ''


if __name__ == 'main':
    print("Starting prometheus server")
    start_http_server(3389)