#!/usr/bin/env python3
import json
import subprocess
from typing import List

import click


def _get_failed_units(cmd) -> List[str]:
    res = subprocess.run(cmd, capture_output=True, check=True)
    if res.returncode == 0:
        units = json.loads(res.stdout)
        return [unit.get("unit") for unit in units]
    else:
        click.echo(f"failed to run {cmd}", err=True)
        return []


@click.command()
@click.version_option(version="0.1.0")
def cli() -> int:
    """
    Get a list of systemd units (system and users) that have
    failed, and print the output in JSON, in a format compatible to
    waybar.
    """
    failed_system_units = _get_failed_units(
        [
            "systemctl",
            "list-units",
            "-o",
            "json",
            "--state=failed",
        ],
    )

    failed_user_units = _get_failed_units(
        [
            "systemctl",
            "--user",
            "list-units",
            "-o",
            "json",
            "--state=failed",
        ],
    )

    failed_units = len(failed_user_units) + len(failed_system_units)

    # The output format documentation:
    # https://github.com/Alexays/Waybar/wiki/Module:-Custom
    output = {"text": failed_units, "class": "success"}

    if failed_units > 0:
        output["class"] = "errors"

        tooltip = []

        if len(failed_user_units) > 0:
            tooltip.append("failed user units: {}".format(", ".join(failed_user_units)))

        if len(failed_system_units) > 0:
            tooltip.append(
                "failed system units: {}".format(", ".join(failed_system_units))
            )

        output["tooltip"] = "\n".join(tooltip)

    click.echo(json.dumps(output))

    return 0


if __name__ == "__main__":
    cli()