import subprocess
import re
import csv
import os

# Define your shell command
command = ["zpool", "iostat", "-Tu", "-l", "-p", "-v", "-y", "1", "1"]

headerRow = [
    "time",
    "name",
    "capacity_alloc",
    "capacity_free",
    "ops_read",
    "ops_write",
    "bandwidth_read",
    "bandwidth_write",
    "total_wait_read",
    "total_wait_write",
    "disk_wait_read",
    "disk_wait_write",
    "syncq_wait_read",
    "syncq_wait_write",
    "asyncq_wait_read",
    "asyncq_wait_write",
    "scrub_wait",
    "trim_wait",
]


def runZpoolIostat():
    # Run the command and capture output
    result = subprocess.run(command, capture_output=True, text=True)

    # Check the return code (0 for success)
    if result.returncode == 0:
        # Access the captured output as a string
        output = result.stdout

        # split output by pools
        pools = re.split(r"^-", output, flags=re.MULTILINE)

        i = 0
        for pool in pools:
            if i == 0:
                # get timestamp
                timestamp = pool.splitlines()[0].strip()

            else:
                # remove ------ lines
                pool = re.sub("^---.*", "", pool).strip()

                # turn - values into 0
                pool = re.sub(" -", "0", pool)

                y = 0

                poolbyline = re.split("\n", pool)

                for line in poolbyline:
                    lineData = line.split()

                    if not lineData:
                        break

                    if y == 0:
                        poolname = lineData[0]
                        filename = poolname + "-zio.csv"

                        # Check if the file exists
                        if not os.path.exists(filename):
                            # Create the file in write mode ('w') if it doesn't exist
                            with open(filename, "w", newline="") as csvfile:
                                csv_writer = csv.writer(csvfile)
                                # Write the header row (optional)
                                csv_writer.writerow(headerRow)

                    # Now open the file in append mode ('a')
                    with open(filename, "a", newline="") as csvfile:
                        csv_writer = csv.writer(csvfile)
                        lineData.insert(0, timestamp)
                        csv_writer.writerow(lineData)

                    y += 1

            i += 1

    else:
        print(f"Error running command: {result.stderr}")


def main():
    runZpoolIostat()


if __name__ == "__main__":
    main()