#!/usr/bin/python3
import argparse
import re
import subprocess
import time

import ltsdb_record

ap = argparse.ArgumentParser()
ap.add_argument("device")
args = ap.parse_args()

p = subprocess.run(["/usr/sbin/smartctl", "-A", args.device],
                   stdout=subprocess.PIPE, universal_newlines=True)
report0 = []
state = 0
for ln in p.stdout.splitlines():
    if state == 0 and ln.startswith("ID# ATTRIBUTE_NAME"):
        state = 1
    elif state == 1 and ln == "":
        state = 2
    elif state == 1:
        (id, attribute_name, flag, value, worst, thresh, type, updated, when_failed, raw_value) = ln.split(None, 9)

        if attribute_name == "Command_Timeout":
            # This is a tuple of three values and I don't know what they mean
            # so I just skip them.
            # I guess I could just record them as smart_command_timeout_1,
            # smart_command_timeout_2 and smart_command_timeout_3 ...
            continue
        if "_Ct" in attribute_name or "_Count" in attribute_name or "_Cnt" in attribute_name:
            unit = "count"
        elif "_Hours" in attribute_name:
            unit = "hours"
        elif "Total_LBAs_Written" in attribute_name:
            unit = "blocks"
        elif "Temperature_Cel" in attribute_name:
            unit = "°C"
        else:
            unit = "unknown"

        if unit == "°C":
            # Sometimes there is extra information included - just ignore that.
            value = int(raw_value.split()[0])
        elif unit == "hours":
            if m := re.match(r"([0-9]+)h\+([0-9]+)m\+([0-9.]+)s", raw_value):
                # e.g. 60633h+54m+11.557s
                value = (int(m.group(1)) * 3600 + int(m.group(2)) * 60 + float(m.group(2))) / 3600
            else:
                value = int(raw_value)
        else:
            value = int(raw_value)

        report0.append(
            {
                "measure": "smart_" + attribute_name.lower(),
                "unit": unit,
                "value": value,
            })
now = time.time()
report = [
    {
        "description": {
            "hostname": ltsdb_record.node,
            "device": args.device,
            "measure": r["measure"],
            "unit": r["unit"],
        },
        "data": [
            [now, r["value"]]
        ]
    }
    for r in report0
]

success = ltsdb_record.record_observations(report)
exit(1 - success)