#!/omd/sites/cmk7309/local/lib/python3/bin/python3

from cmk.agent_based.v2 import (
    CheckPlugin, CheckResult, Result, State, Service, AgentSection
)
from datetime import datetime

# Logging function to append messages to the log file
def log_message(message):
    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    with open('/omd/sites/cmk7309/slurm_py.log', 'a') as log_file:
        log_file.write(f"{timestamp} - {message}\n")

# Log start of script execution
log_message("Starting slurm.py execution")

def parse_slurm(string_table):
    log_message("Parsing slurm section")
    parsed = {}
    for line in string_table:
        state, svc_name, _, status, *message = line
        parsed[svc_name] = {
            "state": int(state),
            "status": status,
            "message": " ".join(message)
        }
    return parsed

def discover_slurm(section):
    log_message("Discovering services")
    for svc_name in section:
        yield Service(item=svc_name)

def check_slurm(item, section):
    log_message(f"Checking service: {item}")
    svc_data = section.get(item)
    if not svc_data:
        return CheckResult(
            state=State.UNKNOWN,
            summary="Item not found"
        )

    status_map = {"OK": State.OK, "WARN": State.WARN, "CRIT": State.CRIT}
    state = status_map.get(svc_data.get("status"), State.UNKNOWN)

    return CheckResult(
        state=state,
        summary=svc_data.get("message")
    )

AgentSection(
    name="slurm",
    parse_function=parse_slurm,
)

CheckPlugin(
    name="slurm",
    service_name="Slurm Service %s",
    sections=["slurm"],
    discovery_function=discover_slurm,
    check_function=check_slurm,
)

# Log end of script execution
log_message("Ending slurm.py execution")