#include "check_nvme.h"
#include <regex>

const char *servicename = (const char*)"NVME";

int getNVMeAttrValue(string line) {
  line = std::regex_replace(line, std::regex("\\s+"), " ");
  line = std::regex_replace(line, std::regex("^ "), "");

  return stoi(line.substr(line.find(":")+1));
}

std::string getNVMeAttrName(string line) {
  size_t first = line.find_first_not_of(' ');
  size_t last = line.find_first_of(':',first);
  return line.substr(first,last-first);
}

int evalStatus(const char* disk, string *status) {
  string output = "";
  string line = "";

  int rc = run_smartctl_cmd(string(SMARTCTL_CMD_ATTRS), disk,&output);
  if(rc) {
    cout << "Error reading SMART data from disk " << disk << endl;
    exit(UNKN);
  }

  map<std::string,NVMeAttr> attrMap;
  attrMap[AVAILABLE_SPARE] = spare;
  attrMap[MEDIA_ERRORS] = errors;
  attrMap[PERCENTAGE_USED] = used;

  stringstream sstream;
  sstream.str(output);
  while(getline(sstream,line)) {
    if(line == "") {
      continue;
    }
    for(map<std::string,NVMeAttr>::iterator it = attrMap.begin(); it != attrMap.end(); ++it) {
      std::string name = it->first;
      if(getNVMeAttrName(line) == name) {
        attrMap[name].value = getNVMeAttrValue(line);
      }
    }
  }

  int ret = OK;
  *status = string(disk);
  for(map<std::string, NVMeAttr>::iterator it = attrMap.begin(); it != attrMap.end(); ++it) {
    NVMeAttr attr = it->second;

    if(attr.value == -1) {
      if(attr.optional) {
        continue;
      }
      *status = string(disk) + " status UNKNOWN";
      return UNKN;
    }

    int veredict = 0;

    if(attr.lower_than) {
      if(attr.value < attr.threshold_warn) {
        veredict = WARN;
      }
      if(attr.threshold_crit != -1 && attr.value < attr.threshold_crit) {
        veredict = CRIT;
      }
    } else {
      if(attr.value > attr.threshold_warn) {
        veredict = WARN;
      }
      if(attr.threshold_crit != -1 && attr.value > attr.threshold_crit) {
        veredict = CRIT;
      }
    }

    switch(veredict) {
      case OK:
        break;
      case WARN:
        if(ret == OK) {
          ret = WARN;
        }
        break;
      case CRIT:
        ret = CRIT;
        break;
    }
    *status += " " + attr.name + ":" + to_string(attr.value);
  }
  return ret;
}

int run_smartctl_cmd(const string command, const char* disk, string* output) {
  int uid = getuid();
  setreuid(0,0);

  int rc = exec(command + disk,output);

  setreuid(uid,0);
  return rc;
}

void printVersion() {
  cout << "check_nvme v" << VERSION << endl << endl;
}

void printHelp(bool longVersion) {
  if(longVersion) {
    printVersion();
    cout << "Checks for pending, reallocated or uncorrectable sectors in NVMe disks" << endl << "WARNING: Requires the setuid bit to be set to run as root" << endl << endl;
    printHelp(false);
    cout << "Options:" << endl;
    cout << " -h" << endl;
    cout << "    Print detailed help screen" << endl;
    cout << " -V" << endl;
    cout << "    Print version information" << endl;
    cout << " DISKS" << endl;
    cout << "    Disks for which to retrieve S.M.A.R.T data" << endl << endl;
    return;
  }
  cout << "Usage: " << endl << "check_smart [-hV] DISKS..." << endl << endl;
}

void set_timeout(unsigned int sec) {
  struct itimerval timer;
  timer.it_value.tv_sec = sec;
  timer.it_value.tv_usec = 0;
  timer.it_interval.tv_sec = 0;
  timer.it_interval.tv_usec = 0;
  setitimer (ITIMER_VIRTUAL, &timer, 0);

  struct sigaction sa;
  memset (&sa, 0, sizeof (sa));
  sa.sa_handler = &timer_handler;
  sigaction (SIGVTALRM, &sa, 0);
}

int main(int argc, char **argv) {
  set_timeout(10);
  int c;
  while ((c = getopt (argc, argv, "Vh")) != -1) {
    switch(c) {
      case 'h':
        printHelp(true);
        return OK;
      case 'V':
        printVersion();
        return OK;
      case '?':
        printHelp(false);
        return UNKN;
    }
  }

  if(argc == 1) {
    cout << "No disks checked" << endl;
    return 3;
  }

  string *results = new string[argc-1];
  int returnCode = OK;

  for(int i = 1; i < argc; i++) {
    switch(evalStatus(argv[i],&(results[i-1]))) {
      case OK:
        break;
      case WARN:
        if(returnCode == OK) {returnCode = WARN;}
        break;
      case CRIT:
        returnCode = CRIT;
        break;
      case UNKN:
        returnCode = UNKN;
        break;
    }
  }

  cout << servicename;
  if(returnCode == OK) {cout << " OK - disk status: ";}
  else if(returnCode == WARN) {cout << " WARNING - disk status: ";}
  else if(returnCode == CRIT) {cout << " CRITICAL - disk status: ";}
  else if(returnCode == UNKN) {cout << " UNKNOWN - disk status: ";}
  for(int i = 0; i < (argc-1); i++) {
    cout << results[i];
    if(i != (argc-2)) {cout << ", ";}
  }
  cout << endl;
  return returnCode;
}