#! /usr/bin/python3

from SmartDevice import SmartDevice
from pySMART import Device, DeviceList

import jsonpickle
import argparse
import sys

actionList = ["list","status","short","long","abort"]

parser = argparse.ArgumentParser(description='Handle smart disks.')
parser.add_argument('action', help="What to do: {}".format(actionList))
parser.add_argument('-d','--disk', required=False, help="Use a single disk ex: 'sdb'")
parser.add_argument('-e','--exclude', required=False, help="Exclude a single disk ex: 'sda'")

args = parser.parse_args()
disk = args.disk
exclude = args.exclude
action = args.action 

devices = {}



if disk:
  dev = Device(disk)
  d = SmartDevice(dev)
  devices[dev] = {
  	"model" : d.model,
  	"smart_status" : d.smart_status,
  	"hours": d.hours
  }
else:
  devlist = DeviceList()
  for dev in devlist.devices:
    try:
      d = SmartDevice(dev)
    except Exception as e:
        print("Failed to convert to device",e)
        continue
    if d.dev == exclude :
      continue
    devices[dev] = {
        "model" : d.model,
        "smart_status" : d.smart_status,
        "hours": d.hours
    }

if action == "list" :

  print (jsonpickle.encode(devices,indent=2))

elif action == "long" :
  for dev in devices:
  	dev.run_selftest("long")
  
elif action == "short" :
  for dev in devices:
  	dev.run_selftest("short")
  
elif action == "status" :
  code = 0
  msgList = []
  for dev in devices:
  	r = dev.get_selftest_result()
  	d = SmartDevice(dev)
  	msg = "{} / {} : {}".format(d.dev,d.serial, r[1])
  	if r[0] == 1 or r[0] == 3:
  		code = 1
  		if r[2] :
  			msg += " {}% done".format(r[2])
  	msgList.append(msg)
  print( "\n".join(msgList) )
  sys.exit(code)

elif action == "json":
  print( jsonpickle.encode(devices, indent=2) )