#!/usr/bin/env python3

import argparse
import dataclasses
import logging
import sys
from urllib.parse import urljoin

import requests
import nagiosplugin

_log = logging.getLogger("nagiosplugin")


@dataclasses.dataclass
class Entities(nagiosplugin.Resource):
    url: str
    token: str

    def hass_get(self, endpoint: str) -> requests.Response:
        headers = {
            "Authorization": "Bearer " + self.token,
            "Content-Type": "application/json",
        }

        r = requests.get(urljoin(self.url, endpoint), headers=headers)
        if not r.ok:
            raise Exception("Failed to query Home Assistant API: " + r.text)

        return r.json()

    def check_api(self):
        message = self.hass_get("/api/").get("message")
        if message != "API running.":
            print("ERROR: " + message)
            sys.exit(1)
        else:
            print("OK: " + message)

    def probe(self):
        response = self.hass_get("/api/states")
        for state in response:
            if state["attributes"].get("device_class") == "battery":
                if state["state"] == "unavailable":
                    _log.info(f"{state['entity_id']} unavailable")
                yield nagiosplugin.Metric(
                    state["entity_id"],
                    float(state["state"]) if state["state"].isnumeric() else -1,
                    uom=state["attributes"].get("unit_of_measurement"),
                    min=0,
                    max=100,
                    context="battery",
                )


@nagiosplugin.guarded
def main():
    argp = argparse.ArgumentParser(description=__doc__)
    argp.add_argument(
        "-t", "--token", required=True, type=str, help="API token for Home Assistant"
    )
    argp.add_argument(
        "-u", "--url", required=True, type=str, help="URL for Home Assistant"
    )
    argp.add_argument("-v", "--verbose", action="count", default=0)

    argp.add_argument(
        "-w",
        "--warning",
        metavar="RANGE",
        required=True,
        help="return warning if battery percentage is outside RANGE",
    )
    argp.add_argument(
        "-c",
        "--critical",
        metavar="RANGE",
        required=True,
        help="return critical if battery percentage is outside RANGE",
    )

    args = argp.parse_args()

    check = nagiosplugin.Check(
        Entities(args.url, args.token),
        nagiosplugin.ScalarContext("battery", args.warning, args.critical),
    )
    check.main(args.verbose)


if __name__ == "__main__":
    main()