#!/usr/bin/python3
"""
Icinga/Nagios plugin to check ESPHome devices status using the ESPHome
Dashboard API.

Copyright (c) 2022 Benjamin Renard <brenard@zionetrix.net>

This program is free software; you can redistribute it and/or
modify it under the terms of the GNU General Public License version 3
as published by the Free Software Foundation.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
"""

import argparse
import sys
import time
import logging
import re

import requests


# nagios exit code
STATUS = {
    'OK': 0,
    'WARNING': 1,
    'CRITICAL': 2,
    'UNKNOWN': 3
}

DEFAULT_HOST = 'http://127.0.0.1:6052'
DEFAULT_RETRY_COUNT = 4
DEFAULT_RETRY_DELAY = 1

parser = argparse.ArgumentParser()
parser.add_argument(
    '-d', '--debug',
    action="store_true",
    dest="debug",
    default=False
)

parser.add_argument(
    '-H', '--host',
    action="store",
    dest="host",
    help=f'ESPHome dashboard URL (default: {DEFAULT_HOST})',
    type=str,
    default=DEFAULT_HOST
)

parser.add_argument(
    '-r', '--retry',
    action="store",
    dest="retry_count",
    help=(
        'Number of retry to retreive device status '
        f'(default: {DEFAULT_RETRY_COUNT})'),
    type=int,
    default=DEFAULT_RETRY_COUNT
)

parser.add_argument(
    '-D', '--delay',
    action="store",
    dest="retry_delay",
    help=(
        'Delay in second between two retry to retreive device status '
        f'(default: {DEFAULT_RETRY_DELAY}s)'),
    type=int,
    default=DEFAULT_RETRY_DELAY
)


def exclude_pattern(value):
    """ Check and compile exclusion pattern parameter """
    return re.compile(value)


parser.add_argument(
    '-x', '--exclude',
    action="append",
    dest="exclude",
    help='Regex exclude pattern(s)',
    type=exclude_pattern,
    default=[]
)
options = parser.parse_args()

logging.basicConfig(
    level=logging.DEBUG if options.debug else logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s')

if options.host[-1] == '/':
    options.host = options.host[-1]


def is_excluded(name):
    """ Check if device is excluded """
    for pattern in options.exclude:
        if pattern.search(name):
            logging.debug('Device %s is excluded', name)
            return True
    logging.debug('Device %s is not excluded', name)
    return False


r = requests.get(f'{options.host}/devices')
devices_data = r.json()
logging.debug('Devices data: %s (%s)', devices_data, type(devices_data))

if not devices_data:
    print('UNKNOWN - Fail to retreive devices using ESPHome Dashboard API')
    sys.exit(STATUS['UNKNOWN'])

COUNT = 0
while COUNT < options.retry_count:
    r = requests.get(f'{options.host}/ping')
    COUNT += 1
    ping_data = r.json()
    logging.debug('Ping data: %s (%s)', ping_data, type(ping_data))

    if ping_data:
        UNREACHABLE = False
        for dev in ping_data:
            if (
                not ping_data[dev]
                and not is_excluded(dev.replace('.yaml', ''))
            ):
                UNREACHABLE = True
                break

        if not UNREACHABLE:
            break

    logging.debug('Wait %d seconds before retry...', options.retry_delay)
    time.sleep(options.retry_delay)

if not ping_data:
    print(
        'UNKNOWN - Fail to retreive devices status '
        'using ESPHome Dashboard API')
    sys.exit(STATUS['UNKNOWN'])

UPDATE_AVAILABLE = 0
UNREACHABLE_DEVICES = 0
NO_PING_DATA = 0
errors = []
devices = {}
for dev in devices_data['configured']:
    devices[dev['name']] = dev
    logging.debug('Device %s: %s', dev['name'], dev)
    if is_excluded(dev['name']):
        continue
    if dev['deployed_version'] != dev['current_version']:
        UPDATE_AVAILABLE += 1
        errors.append(
            f'Update available for device {dev["name"]} '
            f'({dev["deployed_version"]} => {dev["current_version"]})')

    if dev['configuration'] not in ping_data:
        NO_PING_DATA += 1
        errors.append(
            f'No ping data found for device {dev["name"]} '
            f'({dev["configuration"]})')
    elif not ping_data[dev['configuration']]:
        UNREACHABLE_DEVICES += 1
        errors.append(f'Device {dev["name"]} is unreachable')

if not errors:
    print(f'OK - no problem detected on the {len(devices)} devices')
    EXIT_STATUS = STATUS['OK']
else:
    msg = []
    if UNREACHABLE_DEVICES:
        msg.append(f'{UNREACHABLE_DEVICES} unreachable devices')
    if NO_PING_DATA:
        msg.append(f'{NO_PING_DATA} missing ping device status')
    if UPDATE_AVAILABLE:
        msg.append(f'{UPDATE_AVAILABLE} update available')
    print(f'WARNING - {", ".join(msg)}')
    print("\n".join([f'- {error}' for error in errors]))
    EXIT_STATUS = STATUS['WARNING']

print("\nDevices:\n" + "\n".join([
    (
        f'- {name} (version = '
        f'{dev["deployed_version"] if dev["deployed_version"] else "unknown"}'
        f', address = {dev["address"] if dev["address"] else "unknown"})'
    )
    for name, dev in devices.items()
]))
sys.exit(EXIT_STATUS)