import time
import docker
from prometheus_client import start_http_server, Gauge

client = docker.from_env()

cpu_gauge = Gauge("container_cpu_percent", "CPU usage %", ["name", "container_label_coolify_projectName", "container_label_coolify_serviceName"])
mem_gauge = Gauge("container_mem_bytes", "Memory usage bytes", ["name", "container_label_coolify_projectName", "container_label_coolify_serviceName"])
net_rx_gauge = Gauge("container_net_rx_bytes", "Network RX bytes", ["name", "container_label_coolify_projectName", "container_label_coolify_serviceName"])
net_tx_gauge = Gauge("container_net_tx_bytes", "Network TX bytes", ["name", "container_label_coolify_projectName", "container_label_coolify_serviceName"])

def calc_cpu(stats):
    cd = stats["cpu_stats"]["cpu_usage"]["total_usage"] - stats["precpu_stats"]["cpu_usage"]["total_usage"]
    sd = stats["cpu_stats"]["system_cpu_usage"] - stats["precpu_stats"]["system_cpu_usage"]
    cpus = stats["cpu_stats"].get("online_cpus", 1)
    return (cd / sd) * cpus * 100 if sd > 0 else 0

def collect():
    for c in client.containers.list():
        stats = c.stats(stream=False)
        name = c.name
        labels = c.labels
        project = labels.get("coolify.projectName", labels.get("com.docker.compose.project", ""))
        service = labels.get("coolify.service.subName", labels.get("com.docker.compose.service", ""))
        cpu_gauge.labels(name, project, service).set(calc_cpu(stats))
        mem_gauge.labels(name, project, service).set(stats["memory_stats"].get("usage", 0))
        net = stats.get("networks", {})
        rx = sum(v["rx_bytes"] for v in net.values())
        tx = sum(v["tx_bytes"] for v in net.values())
        net_rx_gauge.labels(name, project, service).set(rx)
        net_tx_gauge.labels(name, project, service).set(tx)

start_http_server(9338)
while True:
    collect()
    time.sleep(15)