#!/usr/bin/env python3

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

import requests
import nagiosplugin

_log = logging.getLogger("nagiosplugin")


@dataclasses.dataclass
class Entities(nagiosplugin.Resource):
    url: str
    token: str
    device_class: str
    min: float
    max: float
    filters: list[str]
    attribute: Optional[str]
    friendly_name: bool

    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") == self.device_class and all(
                state["attributes"].get(k) == v for k, v in self.filters
            ):
                if self.attribute is not None:
                    value = state["attributes"].get(self.attribute, -1)
                    uom = None
                else:
                    value = float(state["state"]) if state["state"].isnumeric() else -1
                    uom = state["attributes"].get("unit_of_measurement")

                if self.friendly_name:
                    name = (
                        state["attributes"]
                        .get("friendly_name", state["entity_id"])
                        .translate({ord(c): "_" for c in "'-"})
                    )
                else:
                    name = state["entity_id"]

                if state["state"] == "unavailable":
                    _log.info(f"{state['entity_id']} unavailable")
                yield nagiosplugin.Metric(
                    name,
                    value,
                    uom,
                    context=self.device_class,
                    min=self.min,
                    max=self.max,
                )


def key_value(arg: str):
    k, _, v = arg.partition("=")
    return k, v


@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(
        "-d",
        "--device-class",
        type=str,
        required=True,
        help="Device class of entities to monitor",
    )
    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",
    )
    argp.add_argument(
        "--min",
        type=float,
        help="Min for performance data",
    )
    argp.add_argument(
        "--max",
        type=float,
        help="Max for performance data",
    )
    argp.add_argument(
        "-a", "--attribute", type=str, help="Check attribute instead of value"
    )
    argp.add_argument(
        "-f",
        "--filter",
        type=key_value,
        default=[],
        nargs="*",
        help="Filter by 'attribute=value'",
    )
    argp.add_argument(
        "--friendly",
        action="store_true",
        help="Use friendly name, when available",
    )

    args = argp.parse_args()

    check = nagiosplugin.Check(
        Entities(
            args.url,
            args.token,
            args.device_class,
            args.min,
            args.max,
            args.filter,
            args.attribute,
            args.friendly,
        ),
        nagiosplugin.ScalarContext(args.device_class, args.warning, args.critical),
    )
    check.main(args.verbose)


if __name__ == "__main__":
    main()