#!/usr/bin/python3

import argparse
import logging
import json
import subprocess
import re
import socket
import requests

# getting rid of udev creates a need for fstab backups (or export to zabbix or something else)
# swift03-swiftstore-66 for raid0 storcli


class SystemTool:
    def __init__(self):
        self.entries = self.run()

    def run(self) -> list[dict]:
        """Meant to be implemented in parent classes"""
        return []

    def get(self, **kwargs) -> dict:
        """Return one item from entries matching arbitrary keyword arguments"""
        lookup = []
        result = {}
        for item in self.entries:
            if any(item.get(key) == value for key, value in kwargs.items()):
                lookup.append(item)
        if not lookup:
            logger.info(f"{self.__class__.__name__} found no entry matching '{kwargs}'")
        elif len(lookup) > 1:
            logger.warning(f"{self.__class__.__name__} found more than 1 entry matching '{kwargs}', skipping: {lookup}")
        else:
            result = lookup[0]
        return result


class FStab(SystemTool):
    def run(self) -> list[dict]:
        regex = re.compile(r"^\s*(#?)\s*UUID=(\S{36})\s+(/srv/node/\S+)", re.MULTILINE)
        with open("/etc/fstab", "r") as file:
            fstab = file.read()
        data = set(match.groups((1, 2, 3)) for match in regex.finditer(fstab))
        return [{"mountpoint": x[2], "uuid": x[1], "commented": bool(x[0])} for x in data]


class LSblk(SystemTool):
    def run(self) -> list[dict]:
        cmd = "lsblk --json --bytes --nodeps --output MOUNTPOINT,UUID,PATH,SIZE,SERIAL,REV".split()
        result = subprocess.run(cmd, capture_output=True, check=True, text=True, timeout=3)
        return json.loads(result.stdout).get("blockdevices")


class Smartctl(SystemTool):
    def run(self) -> list[dict]:
        cmd = ["smartctl", "--scan-open", "--", "--info", "--json"]
        result = subprocess.run(cmd, capture_output=True, check=True, text=True, timeout=3)
        data = []
        for command in result.stdout.splitlines():
            if "scsi" not in command:
                continue
            cmd = f"smartctl {command}".split()
            result = subprocess.run(cmd, capture_output=True, check=True, text=True, timeout=3)
            smart = json.loads(result.stdout)
            device = smart.get("device", {}).get("name")
            serial = smart.get("serial_number")
            firmware = smart.get("revision")
            data.append({"device": device, "serial": serial, "firmware": firmware})
        return data


class Storcli(SystemTool):
    def run(self) -> list[dict]:
        cmd = ["storcli64", "/call/sall", "show", "all", "J"]
        try:
            result = subprocess.run(cmd, capture_output=True, check=True, text=True, timeout=3)
            data = json.loads(result.stdout)
        except Exception as e:
            logger.error(f"failed to get storcli data: {e}")
            return []
        drives = []
        for controller in data.get("Controllers"):
            response_data = controller["Response Data"]
            iterator = iter(response_data.items())
            for key, value in iterator:
                slot = key.split()[1]
                state = value[0]["State"]
                key, value = next(iterator)
                serial = value[f"Drive {slot} Device attributes"]["SN"]
                drives.append({"slot": slot, "serial": serial, "state": state})
        return drives


class SwiftRecon:
    def __init__(self):
        hostname = socket.getfqdn()
        self._recon = f"http://{hostname}:6200/recon"
        self.unmounted = self.get_unmounted()
        self.devices = self.get_devices()

    def get_unmounted(self) -> list[str]:
        response = requests.get(f"{self._recon}/unmounted", timeout=3)
        response.raise_for_status()
        return [x["device"] for x in response.json() if x["mounted"] is False]

    def get_devices(self) -> list[str]:
        response = requests.get(f"{self._recon}/devices", timeout=3)
        response.raise_for_status()
        data = response.json()
        if len(data) > 1:
            logger.warning("more than 1 key in response")
        if "/srv/node" not in data:
            logger.warning("/srv/node not in response")
        return sorted(data["/srv/node"])


def bytes_to_human_readable(bytes: int):
    if not isinstance(bytes, int) or bytes <= 0:
        raise ValueError("input must be a positive integer")
    units = ["B", "KB", "MB", "GB", "TB", "PB"]
    index = 0
    size = bytes
    while size >= 1000 and index < len(units) - 1:
        size /= 1000
        index += 1
    return f"{int(size)}{units[index]}"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="basic argparse and logging setup")
    parser.add_argument("-v", "--verbose", action="count", default=0, help="increase log level")
    return parser.parse_args()


def setup_logging(verbosity) -> logging.Logger:
    format = "%(levelname)s: %(message)s"
    level = logging.WARNING
    if verbosity == 1:
        level = logging.INFO
    elif verbosity >= 2:
        level = logging.DEBUG
    logging.basicConfig(level=level, format=format)
    return logging.getLogger(__name__)


def main():
    fstab = FStab()
    lsblk = LSblk()
    storcli = Storcli()
    recon = SwiftRecon()
    smart = Smartctl()
    logger.info("swift devices:", recon.devices)
    logger.info("swift unmounted:", recon.unmounted)
    for swift_device in recon.devices:
        mountpoint = f"/srv/node/{swift_device}"
        unmounted = swift_device in recon.unmounted
        if unmounted:
            logger.warning(f"{mountpoint} is unmounted")

        fstab_info = fstab.get(mountpoint=mountpoint)
        if not fstab_info:
            logger.warning(f"{mountpoint} not found in fstab")
        uuid = fstab_info.get("uuid")

        lsblk_info = lsblk.get(mountpoint=mountpoint)
        if not lsblk_info and not unmounted:
            logger.warning(f"{mountpoint} is mounted according to swift but could not find it with lsblk")
        elif not lsblk_info and uuid:
            lsblk_info = lsblk.get(uuid=uuid)
        if not lsblk_info:
            logger.warning(f"{mountpoint} could not find any matching block device")
        lsblk_uuid = lsblk_info.get("uuid")
        if lsblk_uuid and uuid != lsblk_uuid:
            logger.error(f"{mountpoint} UUID mismatch between mount '{lsblk_uuid}' and fstab '{uuid}'")
            uuid = lsblk_uuid
        device = lsblk_info.get("path")
        bytes = lsblk_info.get("size")
        size = bytes_to_human_readable(int(bytes)) if bytes else None
        serial = lsblk_info.get("serial")
        firmware = lsblk_info.get("rev")
        storcli_info = storcli.get(serial=serial) if serial else {}
        slot = storcli_info.get("slot")

        smart_info = smart.get(device=device) if device else {}
        smart_serial = smart_info.get("serial")
        if smart_serial and serial != smart_serial:
            logger.error(f"{device} serial mismatch between lsblk '{serial}' and smartctl '{smart_serial}'")
            serial = smart_serial

        smart_firmware = smart_info.get("revision")
        if smart_firmware and firmware != smart_firmware:
            logger.error(f"{device} firmware mismatch between lsblk '{firmware}' and smartctl '{smart_firmware}'")
            firmware = smart_firmware

        bla = {
            "mountpoint": mountpoint,
            "device": device,
            "size": size,
            "unmounted": unmounted,
            "firmware": firmware,
            "slot": slot,
            "serial": serial,
            "uuid": uuid,
        }
        print(bla)


if __name__ == "__main__":

    args = parse_args()
    logger = setup_logging(args.verbose)
    main()
