from ..base import BaseExporter
from ..utils import podman_batch_command, get_podman_names
import re
import json

# copied from node_exporter
UNIT_STATES = ["active", "activating", "deactivating", "inactive", "failed"]

class CustomExporter(BaseExporter):
    def metric_config(self):
        return {
            "node_systemd_unit_state": {"type": "Gauge", "labels": ["machine", "name", "state", "type"]},
            "podman_exec_errors": {"type": "Gauge", "labels": ["machine"]},
        }

    def generate(self):
        if len(self.args) == 0:
            raise Exception("services (comma seperated) must be passed as first arg")

        pod_filter = None
        if len(self.args) >= 2:
             pod_filter = self.args[1]

        services = self.args[0].split(",")
        service_state = self.metrics.get('node_systemd_unit_state')
        podman_exec_errors = self.metrics.get('podman_exec_errors')
        pod_names = get_podman_names(pod_filter)

        service_results = podman_batch_command(
            pool_size=5,
            pods=pod_names,
            command=["systemctl", "list-units", "-a", "-o", "json", "--type=service"]
        )
        for pod, sr in service_results.items():
            if sr.returncode == 0:
                sj = json.loads(sr.stdout)
                for service in services:
                    matched_units = []
                    try:
                        service_re = re.compile(service + "\.service$")
                        for unit in sj:
                            if service_re.match(unit['unit']):
                                matched_units.append(unit)
                    except Exception:
                        continue

                    if matched_units:
                        for unit in matched_units:
                            for state in UNIT_STATES:
                                state_value = 0
                                if unit['active'] == state:
                                    state_value = 1
                                service_state.labels(pod, unit['unit'], state, 'unknown').set(state_value)
            else:
                podman_exec_errors.labels(pod).inc()