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

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

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

  // Find nth col
  size_t pos = 0;
  int  cnt = 0;

  while(cnt != col) {
    pos = line.find(" ", pos+1);
    if (pos == std::string::npos) {
      throw std::runtime_error("Column out of range");
    }
    cnt++;
  }
  if(pos != 0) {
    ++pos;
  }

  size_t end_pos = line.find(" ", pos);
  if(end_pos != std::string::npos) {
     return stoi(line.substr(pos, line.find(" ", pos)));
  }
  return stoi(line.substr(pos));
}

int getSmartAttrID(string line) {
  size_t first = line.find_first_not_of(' ');
  size_t last = line.find_first_of(' ',first);
  int id = -1;
  try {
    id = stoi(line.substr(first,last-first));
  }
  catch(...){
    return -1;
  }
  return id;
}

int checkDriveType(const char* disk) {
  string output = "";
  string RRLine = "";
  string line= "";

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

  stringstream sstream;
  sstream.str(output);
  while(getline(sstream,line) && RRLine == "") {
    size_t found = line.find(ROTATION_INFO_STR);
    if(found != string::npos) {
      RRLine = line;
    }
  }
  if(RRLine != "") {
    size_t found = RRLine.find(SSD_DEV_STR);
    if(found != string::npos) {
      return SSD;
    }
  }
  return HDD;
}

map<int,SMARTAttr> prepareAttrMap(int driveType) {
  map<int,SMARTAttr> map;
  map[REALLOC_SEC_COUNT_ID] = reallocated;
  map[REP_UNCORRECT_ID] = rep_uncorrect;

  switch(driveType) {
    case HDD:
      map[CURRENT_PENDING_SEC_ID] = pending;
      map[OFFLINE_UNCORRECT_ID] = off_uncorrect;
    break;
    case SSD:
      map[WEAR_COUNT_ID] = wear;
      map[MEDIA_WEAROUT_ID] = wearout;
      map[RUNTIME_BAD_BLOCKS_ID] = badblocks;
    break;
  }
  return map;
}

int evalStatus(const char* disk, int driveType, 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<int,SMARTAttr> attrMap = prepareAttrMap(driveType);

  stringstream sstream;
  sstream.str(output);
  while(getline(sstream,line)) {
    for(map<int,SMARTAttr>::iterator it = attrMap.begin(); it != attrMap.end(); ++it) {
      int id = it->first;
      int attrID = getSmartAttrID(line);
      if(attrID == -1) {
        continue;
      }
      if(getSmartAttrID(line) == id) {
        attrMap[id].value = getSmartAttrValue(line, attrMap[id].col);
      }
    }
  }
  int ret = OK;
  *status = string(disk);
  for(map<int,SMARTAttr>::iterator it = attrMap.begin(); it != attrMap.end(); ++it) {
    int id = it->first;
    SMARTAttr 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_smart v" << VERSION << endl << endl;
}

void printHelp(bool longVersion) {
  if(longVersion) {
    printVersion();
    cout << "Checks for pending, reallocated or uncorrectable sectors in disks using SMART" << 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++) {
    int driveType = checkDriveType(argv[i]);
    switch(driveType) {
      case HDD:
      case SSD:
        break;
      default:
        cout << "Unknown disk type" << endl;
        return UNKN;
        break;
    }

    switch(evalStatus(argv[i],driveType,&(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;
}