#include "BlkDevInfo.h"
#include <QtGui>
#include <sys/types.h>
#include <dirent.h>
#include <libudev.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <algorithm>
#include <unistd.h>

std::vector<BlkDevInfo*> BlkDevInfo::allblkdev;

QStringList BlkDevInfo::allejectable;

BlkDevInfo::BlkDevInfo(QString init_devname)
 : devname(init_devname),
   size(0),
   parent(NULL),
   ejectable(false),
   invisible(false)
 {
   return;
}

BlkDevInfo::~BlkDevInfo() {
   return;
}

QString
BlkDevInfo::PVendor() {

   // Search parent if necessary to find a vendor string.

   BlkDevInfo * bdi = this;
   while (bdi != NULL) {
      if (! bdi->pvendor.isNull()) {
         return bdi->pvendor;
      }
      bdi = bdi->parent;
   }
   return QString();
}

QString
BlkDevInfo::PModel() {

   // Search parent if necessary to find a model string.

   BlkDevInfo * bdi = this;
   while (bdi != NULL) {
      if (! bdi->pmodel.isNull()) {
         return bdi->pmodel;
      }
      bdi = bdi->parent;
   }
   return QString();
}

QString
BlkDevInfo::Description() {

   // If we can find a vendor and model string, stick them together
   // to form a description.

   QString pv = PVendor();
   QString pm = PModel();
   if (pv.isNull() || pm.isNull()) {
      return devmodel;
   } else {
      pv.append(" ");
      pv.append(pm);
      return pv;
   }
}

void
BlkDevInfo::ReadMountInfo() {

   // Most of this code stolen from the source for the udevadm
   // utility. Digs up everything that seems useful to display
   // about mounted and unmounted disk devices so you can
   // make an intelligent choice about what to mount where.

   // Zero out the list of known entries

   while (allblkdev.size() > 0) {
      delete allblkdev.back();
      allblkdev.pop_back();
   }
   allejectable.clear();

   // Find all the names of all the disk devices on the system and
   // create the BlkDevInfo objects for them.

   QStringList dlist;
   FindDiskDevices(dlist);
   for (int d = 0; d < dlist.size(); ++d) {
      QString dname = dlist.at(d);
      FindBlkDevInfo(dname);
   }

   // Query the properties and attributes of those devices with libudev

   for (int d = 0; d < dlist.size(); ++d) {
      QString dname = dlist.at(d);
      BlkDevInfo * bdi = SoftFindBlkDevInfo(dname);
      std::string cdevname(dname.toAscii().constData());
      struct udev *udev;
      udev = udev_new();
      if (udev != NULL) {
         struct stat statbuf;
         struct udev_device *device = NULL;
         if (stat(cdevname.c_str(), &statbuf) == 0) {
            char type;
            if (S_ISBLK(statbuf.st_mode)) {
               type = 'b';
            } else if (S_ISCHR(statbuf.st_mode)) {
               type = 'c';
            } else {
               type = '\0';
            }
            if (type != '\0') {
               device = udev_device_new_from_devnum(udev, type, statbuf.st_rdev);
            }
            if (device != NULL) {
               struct udev_list_entry *list_entry;
               const char *value;
               udev_list_entry_foreach(list_entry,
                  udev_device_get_properties_list_entry(device)) {
                  const char * name = udev_list_entry_get_name(list_entry);
                  const char * value = udev_list_entry_get_value(list_entry);
                  if (value != NULL) {
                     if (strcmp(name, "ID_MODEL") == 0) {
                        bdi->devmodel = QString(value);
                     } else if (strcmp(name, "ID_FS_LABEL") == 0) {
                        bdi->devlabel = QString(value);
                     } else if (strcmp(name, "ID_FS_UUID") == 0) {
                        bdi->devuuid = QString(value);
                     } else if (strcmp(name, "ID_FS_TYPE") == 0) {
                        bdi->devtype = QString(value);
                     } else if (strcmp(name, "UDISKS_PARTITION_SIZE") == 0) {
                        bdi->size = strtoull(value, NULL, 10);
                        bdi->devsize = FormatSize(bdi->size);
                     }
                  }
               }
               value = udev_device_get_sysattr_value(device, "events");
               if (value != NULL) {
                  if (strstr(value, "eject_request") != NULL) {

                     // Hey! This device supports ejectable media!

                     if (bdi != NULL) {
                        bdi->ejectable = true;
                        allejectable << bdi->devname;
                     }
                  }
               }

               // See if this disk is a partition in a parent disk device

               device = udev_device_get_parent(device);
               if (device != NULL) {
                  QString pdname("/dev/");
                  pdname.append(udev_device_get_sysname(device));
                  BlkDevInfo * pbdi = SoftFindBlkDevInfo(pdname);
                  if (pbdi != NULL) {
                     BlkDevInfo * bdi = SoftFindBlkDevInfo(dname);
                     bdi->parent = pbdi;
                     pbdi->invisible = true;
                  } else {
                     value = udev_device_get_sysattr_value(device, "vendor");
                     if (value != NULL) {
                        bdi->pvendor = QString(value).trimmed();
                     }
                     value = udev_device_get_sysattr_value(device, "model");
                     if (value != NULL) {
                        bdi->pmodel = QString(value).trimmed();
                     }
                  }
                  udev_device_unref(device);
               }
            }
         }
      }
      udev_unref(udev);
   }

   // Use /proc/self/mountinfo file to determine the best mountpoint
   // to mention (picking the one with the shortest root name, where
   // that short name is usually "/" and indicates the one we
   // really want that isn't a bind mount).

   QFile mounts(QString("/proc/self/mountinfo"));
   QRegExp spaces("[ \t\r\n]+");
   if (mounts.open(QIODevice::ReadOnly | QIODevice::Text)) {
      for ( ; ; ) {
         QString m(mounts.readLine());
         QStringList parts = m.split(spaces, QString::SkipEmptyParts);
         QString rootbit;
         QString mntbit;
         QString devbit;
         if (parts.size() >= 5) {
            rootbit = parts.at(3);
            mntbit = parts.at(4);
         }
         for (int i = 6; i < parts.size(); ++i) {
            if (parts.at(i) == QString("-")) {
               if ((i + 2) < parts.size()) {
                  devbit = parts.at(i+2);
               }
               break;
            }
         }
         if ((! devbit.isNull()) &&
             (! mntbit.isNull()) &&
             (! rootbit.isNull())) {
            BlkDevInfo * bdi = SoftFindBlkDevInfo(devbit);
            if (bdi != NULL) {
               if (bdi->devmountroot.isNull() ||
                   (bdi->devmountroot.size() > rootbit.size())) {
                  bdi->devmountroot = rootbit;
                  bdi->devmountpoint = mntbit;
               }
            }
         }
         if (mounts.atEnd()) {
            break;
         }
      }
      mounts.close();
   }
   std::sort(allblkdev.begin(), allblkdev.end(), BlkDevGreater);
}

void
BlkDevInfo::PopulateUnmounted(QTableWidget * qtw) {
   qtw->clearSelection();
   qtw->clearContents();
   qtw->setRowCount(0);
   QStringList headers;
   headers << "Device" << "Type" << "Label" << "Size" << "Model" << "UUID";
   qtw->setHorizontalHeaderLabels(headers);
   qtw->horizontalHeader()->setStretchLastSection(true);
   qtw->verticalHeader()->hide();
   int len = allblkdev.size();
   int row_count = 0;
   for (int i = 0; i < len; ++i) {
      BlkDevInfo * bdi = allblkdev[i];
      if ((! bdi->invisible) &&
          (bdi->devtype.size() > 0) &&
          (bdi->devtype != QString("swap")) &&
          (bdi->devmountpoint.size() == 0)) {
         qtw->setRowCount(row_count + 1);
         qtw->setItem(row_count, 0, new QTableWidgetItem(bdi->devname));
         qtw->setItem(row_count, 1, new QTableWidgetItem(bdi->devtype));
         qtw->setItem(row_count, 2, new QTableWidgetItem(bdi->devlabel));
         qtw->setItem(row_count, 3, new QTableWidgetItem(bdi->devsize));
         qtw->setItem(row_count, 4, new QTableWidgetItem(bdi->Description()));
         qtw->setItem(row_count, 5, new QTableWidgetItem(bdi->devuuid));
         ++row_count;
      }
   }
   for (int c = 0; c < qtw->columnCount() - 1; ++c) {
      qtw->resizeColumnToContents(c);
   }
}

void
BlkDevInfo::PopulateMounted(QTableWidget * qtw) {
   qtw->clearSelection();
   qtw->clearContents();
   qtw->setRowCount(0);
   QStringList headers;
   headers << "Device" << "Mountpoint" << "Type" << "Label" << "Size"
           << "Model" << "UUID";
   qtw->setHorizontalHeaderLabels(headers);
   qtw->horizontalHeader()->setStretchLastSection(true);
   qtw->verticalHeader()->hide();
   int len = allblkdev.size();
   int row_count = 0;
   for (int i = 0; i < len; ++i) {
      BlkDevInfo * bdi = allblkdev[i];
      if ((! bdi->invisible) &&
          (bdi->devtype.size() > 0) &&
          (bdi->devtype != QString("swap")) &&
          (bdi->devmountpoint.size() > 0)) {
         qtw->setRowCount(row_count + 1);
         qtw->setItem(row_count, 0, new QTableWidgetItem(bdi->devname));
         qtw->setItem(row_count, 1, new QTableWidgetItem(bdi->devmountpoint));
         qtw->setItem(row_count, 2, new QTableWidgetItem(bdi->devtype));
         qtw->setItem(row_count, 3, new QTableWidgetItem(bdi->devlabel));
         qtw->setItem(row_count, 4, new QTableWidgetItem(bdi->devsize));
         qtw->setItem(row_count, 5, new QTableWidgetItem(bdi->Description()));
         qtw->setItem(row_count, 6, new QTableWidgetItem(bdi->devuuid));
         ++row_count;
      }
   }
   for (int c = 0; c < qtw->columnCount() - 1; ++c) {
      qtw->resizeColumnToContents(c);
   }
}

QStringList &
BlkDevInfo::GetEjectableDevices() {

   // Return all the names of devices we found that support an eject
   // request.

   return allejectable;
}

QString
BlkDevInfo::GetDeviceDescription(QString devname) {

   // Get description of device which includes more than just devname if
   // available.

   BlkDevInfo * bdi = SoftFindBlkDevInfo(devname);
   if (bdi != NULL) {
      QString desc = bdi->Description();
      QString lab = bdi->devlabel;
      if (! lab.isNull()) {
         devname.append(": ");
         devname.append(lab);
         if (! desc.isNull()) {
            devname.append(" (");
            devname.append(desc);
            devname.append(")");
         }
      } else if (! desc.isNull()) {
         devname.append(": ");
         devname.append(desc);
      }
   }
   return devname;
}

BlkDevInfo *
BlkDevInfo::FindBlkDevInfo(QString devname) {

   // Find (or create if not found) the BlkDevInfo* object for the
   // given device name.

   BlkDevInfo * rval = SoftFindBlkDevInfo(devname);
   if (rval == NULL) {
      rval = new BlkDevInfo(devname);
      allblkdev.push_back(rval);
   }
   return rval;
}

BlkDevInfo *
BlkDevInfo::SoftFindBlkDevInfo(QString devname) {

   // Dig up the BlkDevInfo* for the given device name or NULL if we
   // don't already have such a device name in the list.

   int len = allblkdev.size();
   for (int i = 0; i < len; ++i) {
      if (allblkdev[i]->devname == devname) {
         return allblkdev[i];
      }
   }
   return NULL;
}

QString
BlkDevInfo::FormatSize(unsigned long long size_data) {

   // I tried these calculations with both 1000 and 1024 and the 1000
   // version seemed to more often result in the same number as the
   // product claims, so I stuck with 1000

   QString devsize;
   double dev_blocks = (double)size_data;
   if (dev_blocks > 0.0) {
      QString units;
      if (dev_blocks <= 1000.0) {
         // Print size in bytes
         units = "Bytes";
      } else {
         dev_blocks /= 1000.0;
         if (dev_blocks <= 1000.0) {
            // Print size in KB
            units = "KB";
         } else {
            dev_blocks /= 1000.0;
            if (dev_blocks <= 1000.0) {
               // Print size in MB
               units = "MB";
            } else {
               dev_blocks /= 1000.0;
               // Print size in GB
               units = "GB";
            }
         }
      }
      devsize.setNum(dev_blocks, 'f', 0);
      devsize.append(" ");
      devsize.append(units);
   }
   return devsize;
}

void
BlkDevInfo::FindDiskDevices(QStringList & dlist) {

   // Read the /dev/disk/by-id symlinks to find all the names of all
   // the /dev/<whatever> devices on the system which udev has classified
   // as disk drives or partitions.

   std::string by_id("/dev/disk/by-id");
   DIR *dir;
   struct dirent *dent;
   dir = opendir(by_id.c_str());
   if (dir != NULL) {
      for (dent = readdir(dir); dent != NULL; dent = readdir(dir)) {
         if (strcmp(dent->d_name, ".") == 0) continue;
         if (strcmp(dent->d_name, "..") == 0) continue;
         std::string ln = by_id;
         ln.append("/");
         ln.append(dent->d_name);
         char buf[1025];
         ssize_t len;
         if ((len = readlink(ln.c_str(), buf, sizeof(buf)-1)) != -1) {
            buf[len] = '\0';
            const char * p = strrchr(buf, '/');
            if (p != NULL) {
               QString dname("/dev/");
               dname.append(p+1);
               if (dlist.indexOf(dname) == -1) {
                  dlist << dname;
               }
            }
         }
      }
      closedir(dir);
   }
   dlist.sort();
}

bool
BlkDevInfo::BlkDevGreater(const BlkDevInfo * a, const BlkDevInfo * b) {

   // Comapre BlkDevInfo * contents to get them sorted in a readable
   // fashion - first by mountpoint, then by partition label, then
   // by the /dev/<devname> strings.

   if (a->devmountpoint > b->devmountpoint) {
      return true;
   } else if (a->devmountpoint == b->devmountpoint) {
      if (a->devlabel > b->devlabel) {
         return true;
      } else if (a->devlabel == b->devlabel) {
         if (a->devname > b->devname) {
            return true;
         }
      }
   }
   return false;
}
