from prometheus_client import Gauge
from django.db.models import Count
from django.utils import timezone

from labs.models import (
    LabInstance,
    LabSSHServer,
)


lab_instance_status = Gauge(
    name="lab_instance_status",
    documentation="The status of the lab instance",
    labelnames=["status"],
)

lab_instance_creating_duration_gauge = Gauge(
    name="lab_instance_creating_duration",
    documentation="The duration of the lab instance creating",
    labelnames=["lab_instance_id"],
)

lab_instance_terminating_duration_gauge = Gauge(
    name="lab_instance_terminating_duration",
    documentation="The duration of the lab instance terminating",
    labelnames=["lab_instance_id"],
)

lab_ssh_server_status_gauge = Gauge(
    name="lab_ssh_server_status",
    documentation="The status of the lab ssh server",
    labelnames=["domain"],
)

lab_ssh_server_active_session_gauge = Gauge(
    name="lab_ssh_server_active_session",
    documentation="The number of active sessions of the lab ssh server",
    labelnames=["domain"],
)


def update_metrics():
    now = timezone.now()

    lab_instance_status_counts = dict(
        LabInstance.objects.values("status")
        .annotate(count=Count("status"))
        .values_list("status", "count")
    )
    for status in LabInstance.LabInstanceStatus:
        lab_instance_status.labels(status.value).set(
            lab_instance_status_counts.get(status.value, 0)
        )

    creating_lab_instances = LabInstance.objects.filter(
        status=LabInstance.LabInstanceStatus.CREATING
    ).values("id", "created_at")
    creating_lab_instance_ids = {
        lab_instance["id"] for lab_instance in creating_lab_instances
    }

    # Remove labels for instances that are no longer in the CREATING state
    for label in lab_instance_creating_duration_gauge.collect()[0].samples:
        if label.labels["lab_instance_id"] not in creating_lab_instance_ids:
            lab_instance_creating_duration_gauge.remove(label.labels["lab_instance_id"])

    for lab_instance in creating_lab_instances:
        lab_instance_creating_duration_gauge.labels(lab_instance["id"]).set(
            (now - lab_instance["created_at"]).total_seconds()
        )

    terminating_lab_instances = LabInstance.objects.filter(
        status=LabInstance.LabInstanceStatus.TERMINATING
    ).values("id", "start_terminating_at")
    terminating_lab_instance_ids = {
        lab_instance["id"] for lab_instance in terminating_lab_instances
    }

    # Remove labels for instances that are no longer in the TERMINATING state
    for label in lab_instance_terminating_duration_gauge.collect()[0].samples:
        if label.labels["lab_instance_id"] not in terminating_lab_instance_ids:
            lab_instance_terminating_duration_gauge.remove(
                label.labels["lab_instance_id"]
            )

    for lab_instance in terminating_lab_instances:
        lab_instance_terminating_duration_gauge.labels(lab_instance["id"]).set(
            (now - lab_instance["start_terminating_at"]).total_seconds()
        )

    ssh_servers = LabSSHServer.objects.all().values(
        "domain", "active_session", "status"
    )
    for ssh_server in ssh_servers:
        lab_ssh_server_status_gauge.labels(ssh_server["domain"]).set(
            int(
                LabSSHServer.Status(ssh_server["status"]) == LabSSHServer.Status.HEALTHY
            )
        )
        lab_ssh_server_active_session_gauge.labels(ssh_server["domain"]).set(
            ssh_server["active_session"]
        )