#!/usr/bin/python3

import collections
import http.server
import io
import logging
import os
import re
import subprocess
import sys
import typing


logging.basicConfig(level=logging.INFO)


def get_error_counts(
    zpool_status_p_output: str,
) -> typing.Dict[str, typing.Dict[str, int]]:
    out = zpool_status_p_output.splitlines()

    error_counts: typing.Dict[str, typing.Dict[str, int]] = {}
    pool_name = ""
    summing = 0

    for line in out:
        m = re.match("^  pool: (.+)", line)
        if m:
            pool_name = str(m.groups(1)[0])
            error_counts[pool_name] = {
                "read": 0,
                "write": 0,
                "cksum": 0,
                "data_errors": 0,
            }
            summing = 0
            continue

        m = re.match("^errors: ([0-9]+) data errors", line)
        if m:
            (data_errors,) = m.groups()
            error_counts[pool_name]["data_errors"] += int(data_errors)

        m = re.match("^\tNAME", line)
        if m:
            summing = 1
            continue

        if summing:
            m = re.match("^\t *\\S+\\s+\\S+\\s+([0-9]+)\\s+([0-9]+)\\s+([0-9]+)", line)
            if m:
                read, write, cksum = m.groups()
                error_counts[pool_name]["read"] += int(read)
                error_counts[pool_name]["write"] += int(write)
                error_counts[pool_name]["cksum"] += int(cksum)

    return error_counts


def get_progress_status(line: str) -> typing.Tuple[float, float] | None:
    m = re.match(".*, ([0-9.]+)% done,.*", line)
    if m:
        progress = float(m.groups(1)[0]) / 100.0
        n = re.match(".*,(|[0-9]+ days) ([0-9]+):([0-9]+):([0-9]+) to go", line)
        if n:
            days_str, hours, minutes, seconds = (
                str(n.groups(1)[0]),
                int(n.groups(1)[1]),
                int(n.groups(1)[2]),
                int(n.groups(1)[3]),
            )
            if days_str.endswith(" days"):
                days = int(days_str[:-5])
            else:
                days = 0
            estimate = float(days * 86400 + hours * 3600 + minutes * 60 + seconds)
        else:
            estimate = float("inf")
        return (progress, estimate)
    return None


def get_scrub_progress(
    zpool_status_p_output: str,
) -> typing.Dict[str, typing.Tuple[float, float]]:
    progresses: typing.Dict[str, typing.Tuple[float, float]] = {}
    pool_name = ""

    for line in zpool_status_p_output.splitlines():
        m = re.match("^  pool: (.+)", line)
        if m:
            pool_name = str(m.groups(1)[0])
            scrubbing = False
            continue

        m = re.match(".*scrub repaired.*", line)
        if m:
            progresses[pool_name] = (1.0, 0.0)
            continue

        m = re.match(".*scrub in progress.*", line)
        if m:
            scrubbing = True
            continue

        if scrubbing:
            p = get_progress_status(line)
            if p is not None:
                progresses[pool_name] = p

    return progresses


def get_resilver_progress(
    zpool_status_p_output: str,
) -> typing.Dict[str, typing.Tuple[float, float]]:
    progresses: typing.Dict[str, typing.Tuple[float, float]] = {}
    pool_name = ""

    for line in zpool_status_p_output.splitlines():
        m = re.match("^  pool: (.+)", line)
        if m:
            pool_name = str(m.groups(1)[0])
            resilvering = False
            continue

        m = re.match(".*resilvered.*in.*with", line)
        if m:
            progresses[pool_name] = (1.0, 0.0)
            continue

        m = re.match(".*resilver in progress.*", line)
        if m:
            resilvering = True
            continue

        if resilvering:
            p = get_progress_status(line)
            if p is not None:
                progresses[pool_name] = p

    return progresses


def get_health(zpool_status_p_output: str) -> typing.Dict[str, bool]:
    healths: typing.Dict[str, bool] = {}
    pool_name = ""

    for line in zpool_status_p_output.splitlines():
        m = re.match("^  pool: (.+)", line)
        if m:
            pool_name = str(m.groups(1)[0])
            continue

        m = re.match("^ state: (.+)", line)
        if m:
            healthy = "ONLINE" == m.groups(1)[0]
            healths[pool_name] = healthy
            continue

    return healths


_module_version = ""


def get_kernel_module_version():
    """Get the kernel module version.

    This function caches it for the duration of the process.
    """
    global _module_version
    if _module_version == "":
        try:
            with open("/sys/module/zfs/version") as f:
                _module_version = f.read().strip()
        except Exception:
            _module_version = None
    return _module_version


class SimplePromHandler(http.server.BaseHTTPRequestHandler):
    def log_request(self, code=None, size=None):
        # FIXME: bump standard http_ metrics and send them back on request.
        pass

    def handle_non_metrics(self):
        if self.path == "/":
            self.send_response(200)
            self.send_header("Content-Type", "text/html; charset=utf-8")
            self.end_headers()
            self.wfile.write(
                (
                    "<html><body>%s exporter.  "
                    "<a href='/metrics'>Metrics.</a>"
                    "<body></html>" % os.path.basename(__file__)
                ).encode("utf-8")
            )
            return True
        elif self.path != "/metrics":
            self.send_response(404)
            self.end_headers()
            return True

    def do_GET(self):
        if self.handle_non_metrics():
            return

        self.send_response(200)
        self.send_header("Content-Type", "text/plain; version=0.0.4; charset=utf-8")
        self.end_headers()

        f = io.StringIO()

        module_version = get_kernel_module_version()
        if module_version:
            print(
                'zfs_info{version="%s"} %s' % (module_version, 1),
                file=f,
            )
        else:
            print(
                "zfs_info %s" % (0,),
                file=f,
            )
            f.seek(0)
            self.wfile.write(f.read().encode("utf-8"))
            return

        zpool_status_p_output = subprocess.check_output(
            ["zpool", "status", "-p"], universal_newlines=True
        )

        for pool, errors in get_error_counts(zpool_status_p_output).items():
            for klass, count in errors.items():
                print(
                    'zfs_pool_errors_total{zpool="%s", class="%s"} %s'
                    % (pool, klass, count),
                    file=f,
                )

        def plusinf(f: float) -> str | float:
            if f == float("inf"):
                return "+Inf"
            return f

        for pool, progress in get_scrub_progress(zpool_status_p_output).items():
            print(
                'zfs_pool_scrub_progress{zpool="%s"} %s' % (pool, progress[0]),
                file=f,
            )
            print(
                'zfs_pool_scrub_remaining_time_seconds{zpool="%s"} %s'
                % (pool, plusinf(progress[1])),
                file=f,
            )

        for pool, progress in get_resilver_progress(zpool_status_p_output).items():
            print(
                'zfs_pool_resilver_progress{zpool="%s"} %s' % (pool, progress[0]),
                file=f,
            )
            print(
                'zfs_pool_resilver_remaining_time_seconds{zpool="%s"} %s'
                % (pool, plusinf(progress[1])),
                file=f,
            )

        for pool, healthy in get_health(zpool_status_p_output).items():
            print(
                'zfs_pool_healthy{zpool="%s"} %s' % (pool, 1 if healthy else 0),
                file=f,
            )

        poolitem = collections.namedtuple("poolitem", "name size free")
        poolitems = {}
        for r in subprocess.check_output(
            ["zpool", "list", "-Hp", "-o", "name,size,free"]
        ).splitlines():
            if not r:
                continue
            r = r.split(b"\t")
            r = poolitem(
                r[0].decode("utf-8").replace('"', "_"), *(int(x) for x in r[1:])
            )
            poolitems[r.name] = r
            print('zfs_pool_avail_bytes{zpool="%s"} %s' % (r.name, r.free), file=f)
            print('zfs_pool_size_bytes{zpool="%s"} %s' % (r.name, r.size), file=f)
        datasetitem = collections.namedtuple(
            "datasetitem", "name type avail refer used compressratio"
        )
        for r in subprocess.check_output(
            [
                "zfs",
                "list",
                "-t",
                "filesystem,snapshot,volume",
                "-Hp",
                "-o",
                "name,type,avail,refer,used,compressratio",
            ]
        ).splitlines():
            if not r:
                continue

            def autoint(x):
                if b"." in x:
                    return float(x)
                try:
                    return int(x)
                except Exception:
                    return None

            r = r.split(b"\t")
            r = datasetitem(
                r[0].decode("utf-8").replace('"', "_"),
                r[1].decode("utf-8").replace('"', "_"),
                *(autoint(x) for x in r[2:]),
            )
            if r.avail is not None:
                print(
                    'zfs_dataset_avail_bytes{dataset="%s", type="%s"} %s'
                    % (r.name, r.type, r.avail),
                    file=f,
                )
            if r.refer is not None:
                print(
                    'zfs_dataset_refer_bytes{dataset="%s", type="%s"} %s'
                    % (r.name, r.type, r.refer),
                    file=f,
                )
            if r.used is not None:
                print(
                    'zfs_dataset_used_bytes{dataset="%s", type="%s"} %s'
                    % (r.name, r.type, r.used),
                    file=f,
                )
            if r.compressratio is not None:
                print(
                    'zfs_dataset_compress_ratio{dataset="%s", type="%s"} %.2f'
                    % (r.name, r.type, r.compressratio),
                    file=f,
                )
            try:
                print(
                    'zfs_dataset_size_bytes{dataset="%s", type="%s"} %s'
                    % (r.name, r.type, poolitems[r.name.split("/")[0]].size),
                    file=f,
                )
            except KeyError:
                pass
        f.seek(0)
        self.wfile.write(f.read().encode("utf-8"))


server = http.server.HTTPServer(("", int(sys.argv[1])), SimplePromHandler)
logging.info("Serving on TCP port %s", sys.argv[1])
server.serve_forever()
#    httpd.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
