#!/usr/libexec/platform-python
import os
import sys
import subprocess
import re
import json

json_zfs = {
    "zfs_installed": False
}


def get_zfs_list():
    try:
        zfs_list_result = subprocess.Popen(
            ["zfs", "list"], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, universal_newlines=True).stdout
        # zfs_list_result = subprocess.Popen(
            # ["cat","/root/spoof/zfs_list"], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, universal_newlines=True).stdout
    except:
        return False

    zpools = []
    for line in zfs_list_result:
        regex = re.search("^(\S+)\s+(\S+)\s+(\S+)\s+\S+\s+(\S+).*$", line)
        if regex != None and regex.group(1) != "NAME" and not any(x in regex.group(1) for x in ["/","@"]):
            zpools.append(
                {
                    "name": regex.group(1),
                    "used": regex.group(2),
                    "avail": regex.group(3),
                    "mountpoint": regex.group(4)
                }
            )
    return zpools

def get_zpool_list():
    try:
        zpool_list_result = subprocess.Popen(
            ["zpool", "list"], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, universal_newlines=True).stdout
        # zpool_list_result = subprocess.Popen(
            # ["cat","/root/spoof/zpool_list"], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, universal_newlines=True).stdout
    except:
        return False
    #NAME              SIZE  ALLOC   FREE  CKPOINT  EXPANDSZ   FRAG    CAP  DEDUP    HEALTH  ALTROOT
    #stornado_pool_a  34.9T  2.52M  34.9T        -         -     0%     0%  1.00x    ONLINE  -
    #stornado_pool_b  1.30T  1.17M  1.30T        -         -     0%     0%  1.00x  SUSPENDED  -
    #stornado_pool_c  3.48T   636K  3.48T        -         -     0%     0%  1.00x    ONLINE  -

    zpools = []
    for line in zpool_list_result:
        regex = re.search("^(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+).*$", line)
        if regex != None and regex.group(1) != "NAME":
            zpools.append(
                {
                    "name": regex.group(1),
                    "raw_size": regex.group(2),
                    "raw_alloc": regex.group(3),
                    "raw_free": regex.group(4),
                    "ckpoint": regex.group(5),
                    "expandsz": regex.group(6),
                    "frag": regex.group(7),
                    "cap": regex.group(8),
                    "dedup": regex.group(9),
                    "health": regex.group(10),
                    "altroot": regex.group(11)
                }
            )
    zfs_list = get_zfs_list()
    for pool in zpools:
        for entry in zfs_list:
            if entry["name"] == pool["name"]:
                pool["used"] = entry["used"]
                pool["avail"] = entry["avail"]
                pool["mountpoint"] = entry["mountpoint"]
        if not all (key in pool for key in ("used","avail","mountpoint")):
            pool["used"] = "-"
            pool["avail"] = "-"
            pool["mountpoint"] = "-"

    return zpools

def zpool_status(pool_name):
    try:
        zpool_status_result = subprocess.Popen(
            ["zpool", "status", pool_name], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, universal_newlines=True).stdout.read()
        # zpool_status_result = subprocess.Popen(
            # ["cat","/root/spoof/zpool_status"], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, universal_newlines=True).stdout.read()
    except:
        print(f"failed to run 'zpool status {pool_name}'")
        exit(1)
    groupings = [(match.group(1),match.group(0)) for match in re.finditer(r"^\t{1}(\S+).*$\n(?:^\t{1} +.*$\n)+|^\t{1}(\S+).*$\n(?:^\t{1} +.*$\n)+",zpool_status_result,flags=re.MULTILINE)]
    groupings.append( [("state",match.group(1)) for match in re.finditer(r"^.*state\:\s+(\S+)",zpool_status_result,flags=re.MULTILINE)][0])
    return dict(groupings)

def zpool_iostat(pool_name):
    try:
        zpool_status_result = subprocess.Popen(
            ["zpool", "iostat","-v", pool_name], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, universal_newlines=True).stdout.read()
        # zpool_status_result = subprocess.Popen(
            # ["cat","/root/spoof/zpool_iostat_v"], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, universal_newlines=True).stdout.read()
    except:
        print(f"failed to run 'zpool iostat -v {pool_name}'")
        exit(1)
    groupings = [(match.group(1),match.group(0)) for match in re.finditer(r"^(\S+).*$\n(?:^ +.*$\n)+|^(\S+).*$\n(?:^ +.*$\n)+",zpool_status_result,flags=re.MULTILINE)]
    return dict(groupings)

def zpool_status_parse(zp_status_obj,key):
    if key not in zp_status_obj.keys():
        return [], []
    vdevs = [
            {
            "tag":key,
            "name": match.group(1),
            "state": match.group(2), 
            "read_errors": match.group(3), 
            "write_errors": match.group(4), 
            "checksum_errors": match.group(5)
            } for match in re.finditer(r"^\t  (\S+-\d+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+).*",zp_status_obj[key],flags=re.MULTILINE)
    ]
    disks = [
            {
            "tag":key,
            "name":match.group(1),
            "state": match.group(2), 
            "read_errors": match.group(3), 
            "write_errors": match.group(4), 
            "checksum_errors": match.group(5)
            } for match in re.finditer(r"^\t    (\S+)(?:-part[0-9])?\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+).*",zp_status_obj[key],flags=re.MULTILINE)
    ]

    exception_match = r"^(\d+-\d+)(?:-part[0-9])"
    for disk in disks:
            match = re.match(exception_match,disk["name"])
            if match:
                disk["name"] = match.group(1)

    counts = []
    disk_count = 0
    initial_disk = True
    for line in zp_status_obj[key].splitlines():
        regex_vdev = re.search("^\t  (\S+-\d+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+).*",line)
        regex_disk = re.search("^\t    (\S+)(?:-part[0-9])?\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+).*",line)
        if regex_vdev != None and not initial_disk:
            counts.append(disk_count)
            disk_count = 0
        if regex_disk != None:
            initial_disk = False
            disk_count = disk_count + 1
    counts.append(disk_count)

    return vdevs, disks, counts

def zpool_iostat_parse(zp_status_obj,key):
    if key not in zp_status_obj.keys():
        return [], []
    vdevs = [
            {
            "tag":key,
            "raid_level":match.group(1), 
            "alloc": match.group(2),
            "free": match.group(3),
            "read_ops": match.group(4),
            "write_ops": match.group(5),
            "read_bw": match.group(6),
            "write_bw": match.group(7)
            } for match in re.finditer(r"^  (\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+).*",zp_status_obj[key],flags=re.MULTILINE)
        ]
    disks = [
            {
            "tag":key,
            "name":match.group(1), 
            "alloc": match.group(2),
            "free": match.group(3),
            "read_ops": match.group(4),
            "write_ops": match.group(5),
            "read_bw": match.group(6),
            "write_bw": match.group(7)
            } for match in re.finditer(r"^    (\S+)(?:-part[0-9])?\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+).*",zp_status_obj[key],flags=re.MULTILINE)
        ]

    exception_match = r"^(\d+-\d+)(?:-part[0-9])"
    for disk in disks:
            match = re.match(exception_match,disk["name"])
            if match:
                disk["name"] = match.group(1)

    counts = []
    disk_count = 0
    for line in zp_status_obj[key].splitlines():
        regex_vdev = re.search("^  (\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+).*",line)
        regex_disk = re.search("^    (\S+)(?:-part[0-9])?\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+).*",line)
        if regex_vdev != None and disk_count > 0:
            counts.append(disk_count)
            disk_count = 0
        if regex_disk != None:
            disk_count = disk_count + 1
    counts.append(disk_count)
    return vdevs, disks, counts

def verify_zfs_device_format(zp_status_obj,pool_name):
    alert = []
    default_pattern = r"^\t    (\d+-\d+)(?:-part[0-9])?\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+).*"
    unsupported_pattern = r"^\t    (\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+).*"
    disks = [
            {
            "tag":pool_name,
            "name":match.group(1),
            "state": match.group(2), 
            "read_errors": match.group(3), 
            "write_errors": match.group(4), 
            "checksum_errors": match.group(5)
            } for match in re.finditer(default_pattern,zp_status_obj[pool_name],flags=re.MULTILINE)
    ]
    unsupported_disks = [
            {
                "tag":pool_name,
                "name":match.group(1),
                "state": match.group(2), 
                "read_errors": match.group(3), 
                "write_errors": match.group(4), 
                "checksum_errors": match.group(5)
            } for match in re.finditer(unsupported_pattern,zp_status_obj[pool_name],flags=re.MULTILINE)
    ]
    if len(unsupported_disks) > len(disks):
        for disk in disks:
            if disk in unsupported_disks:
                unsupported_disks.remove(disk)
        exception_match = r"^(\d+-\d+)(?:-part[0-9])"
        for disk in unsupported_disks:
            match = re.match(exception_match,disk["name"])
            if match:
                unsupported_disks.remove(disk)
        # default pattern didn't match.
        alert.append("ZFS status displayed by this module for zpool '{pn}' may be incomplete.\n\n".format(pn=pool_name))
        alert.append("This module can only display zfs status information for devices that are created using a device alias.\n\n")
        alert.append("This can be done using the 45Drives cockpit-zfs-manager package:\nhttps://github.com/45Drives/cockpit-zfs-manager/releases/\n\n")
        if unsupported_disks:
            alert.append("The following zfs devices do not conform:\n")
            for disk in unsupported_disks:
                alert.append("\t  {d}\n".format(d=disk["name"]))
        alert.append("\n")
        #sys.exit(1)
    return alert


def get_zpool_status():
    json_zfs["warnings"] = []
    for pool in json_zfs["zpools"]:
        # run 'zpool status <pool>' and 'zpool iostat -v <pool>' and group output by top level entry (<pool name>, special, cache etc.)
        status_output = zpool_status(pool["name"])
        iostat_output = zpool_iostat(pool["name"])
        pool["state"] = status_output["state"]
        pool["vdevs"] = []
        alert = verify_zfs_device_format(status_output,pool["name"])
        if alert:
            json_zfs["warnings"] = json_zfs["warnings"] + alert
        for key in status_output.keys():
            # parse the output of both commands by top level entry
            if key in iostat_output.keys():
                # get all parsed output as arrays of objects from each command. 
                status_vdevs, status_disks, status_disk_counts = zpool_status_parse(status_output,key)
                iostat_vdevs, iostat_disks, iostat_disk_counts = zpool_iostat_parse(iostat_output,key)
                if not status_disks or not iostat_disks or not status_disk_counts or not iostat_disk_counts:
                    print("/usr/share/cockpit/45drives-disks/scripts/zfs_info failed to interpret the following zfs information:")
                    print("zpool status {pn}:".format(pn=pool["name"]))
                    print(status_output[key])
                    print("zpool iostat -v {pn}:".format(pn=pool["name"]))
                    print(iostat_output[key])
                    print("Other Information: ")
                    print("status_vdevs",json.dumps(status_vdevs,indent=2))
                    print("status_disks",json.dumps(status_disks,indent=2))
                    print("status_disk_counts",json.dumps(status_disk_counts,indent=2))
                    print("iostat_vdevs",json.dumps(iostat_vdevs,indent=2))
                    print("iostat_disks",json.dumps(iostat_disks,indent=2))
                    print("iostat_disk_counts",json.dumps(iostat_disk_counts,indent=2))
                    exit(1)
                disk_index = 0
                for i in range(0,len(status_vdevs)):
                    #combine output of iostat command and zpool status for each vdev
                    status_vdevs[i]["raid_level"] = iostat_vdevs[i]["raid_level"]
                    status_vdevs[i]["alloc"] = iostat_vdevs[i]["alloc"]
                    status_vdevs[i]["free"] = iostat_vdevs[i]["free"]
                    status_vdevs[i]["read_ops"] = iostat_vdevs[i]["read_ops"]
                    status_vdevs[i]["write_ops"] = iostat_vdevs[i]["write_ops"]
                    status_vdevs[i]["read_bw"] = iostat_vdevs[i]["read_bw"]
                    status_vdevs[i]["write_bw"] = iostat_vdevs[i]["write_bw"]
                    status_vdevs[i]["disks"] = []
                    for j in range(disk_index,disk_index + status_disk_counts[i]):
                        # status_disk_counts stores the number of disks in current vdev
                        # combine the output of zpool status and zpool iostat for each disk in the current vdev
                        status_disks[j]["alloc"] = iostat_disks[j]["alloc"]
                        status_disks[j]["free"] = iostat_disks[j]["free"]
                        status_disks[j]["read_ops"] = iostat_disks[j]["read_ops"]
                        status_disks[j]["write_ops"] = iostat_disks[j]["write_ops"]
                        status_disks[j]["read_bw"] = iostat_disks[j]["read_bw"]
                        status_disks[j]["write_bw"] = iostat_disks[j]["write_bw"]
                        status_disks[j]["vdev_idx"] = len(pool["vdevs"])
                        status_vdevs[i]["disks"].append(status_disks[j])
                    pool["vdevs"].append(status_vdevs[i])
                    disk_index = disk_index + status_disk_counts[i]

def check_zfs():
    try:
        command_result = subprocess.run(
            ["command -v zfs"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, shell=True).returncode
    except:
        return False
    return (command_result == 0)


def create_disk_entries():
    disk_entries = {}
    for pool_index, pool in enumerate(json_zfs["zpools"]):
        for vdev in pool["vdevs"]:
            for disk in vdev["disks"]:
                disk_entries[disk["name"]] = {}
                disk_entries[disk["name"]]["zpool_name"] = pool["name"]
                disk_entries[disk["name"]]["zpool_used"] = pool["used"]
                disk_entries[disk["name"]]["zpool_avail"] = pool["avail"]
                disk_entries[disk["name"]]["zpool_mountpoint"] = pool["mountpoint"]
                disk_entries[disk["name"]]["zpool_state"] = pool["state"]
                disk_entries[disk["name"]]["zpool_idx"] = pool_index
                disk_entries[disk["name"]]["vdev_raid_level"] = vdev["raid_level"]
                disk_entries[disk["name"]]["vdev_alloc"] = vdev["alloc"]
                disk_entries[disk["name"]]["vdev_free"] = vdev["free"]
                disk_entries[disk["name"]]["vdev_read_ops"] = vdev["read_ops"]
                disk_entries[disk["name"]]["vdev_write_ops"] = vdev["write_ops"]
                disk_entries[disk["name"]]["vdev_read_bw"] = vdev["read_bw"]
                disk_entries[disk["name"]]["vdev_write_bw"] = vdev["write_bw"]
                disk_entries[disk["name"]]["name"] = disk["name"]
                disk_entries[disk["name"]]["alloc"] = disk["alloc"]
                disk_entries[disk["name"]]["free"] = disk["free"]
                disk_entries[disk["name"]]["read_ops"] = disk["read_ops"]
                disk_entries[disk["name"]]["write_ops"] = disk["write_ops"]
                disk_entries[disk["name"]]["read_bw"] = disk["read_bw"]
                disk_entries[disk["name"]]["write_bw"] = disk["write_bw"]
                disk_entries[disk["name"]]["vdev_idx"] = disk["vdev_idx"]
                disk_entries[disk["name"]]["state"] = disk["state"]
                disk_entries[disk["name"]]["read_errors"] = disk["read_errors"]
                disk_entries[disk["name"]]["write_errors"] = disk["write_errors"]
                disk_entries[disk["name"]]["checksum_errors"] = disk["checksum_errors"]
                disk_entries[disk["name"]]["tag"] = disk["tag"]

    json_zfs["zfs_disks"] = disk_entries


def main():
    if check_zfs():
        json_zfs["zfs_installed"] = True
        json_zfs["zpools"] = get_zpool_list()
        get_zpool_status()
        create_disk_entries()

    print(json.dumps(json_zfs, indent=4))


if __name__ == "__main__":
    main()
