mainwindow.cpp

#include "mainwindow.h"
#include "./ui_mainwindow.h"
#include <QHeaderView>
#include <QInputDialog>
#include <QMessageBox>

MainWindow::MainWindow(QWidget *parent)
    : QMainWindow(parent), ui(new Ui::MainWindow),
      m_driveManager(new DriveManager(this)), m_shredder(new Shredder(this)),
      m_formatter(new Formatter(this)) {
  ui->setupUi(this);

  // Add new erasure methods
  ui->methodCombo->addItem("Secure Erase (ATA/NVMe)");
  ui->methodCombo->addItem("Cryptographic Erasure");

  // Configure table
  ui->driveTable->horizontalHeader()->setSectionResizeMode(
      QHeaderView::Stretch);
  ui->driveTable->setSelectionBehavior(QAbstractItemView::SelectRows);
  ui->driveTable->setEditTriggers(QAbstractItemView::NoEditTriggers);

  // Connect signals
  connect(m_driveManager, &DriveManager::drivesUpdated, this,
          &MainWindow::updateDrives);
  connect(ui->shredButton, &QPushButton::clicked, this,
          &MainWindow::onShredClicked);
  connect(m_shredder, &Shredder::progressUpdated, this,
          &MainWindow::onShredProgress);
  connect(m_shredder, &Shredder::timerUpdated, this,
          &MainWindow::onTimerUpdated);
  connect(m_shredder, &Shredder::finished, this, &MainWindow::onShredFinished);
  connect(m_formatter, &Formatter::finished, this,
          &MainWindow::onFormatFinished);

  // Initial update
  updateDrives(m_driveManager->getDrives());
}

MainWindow::~MainWindow() { delete ui; }

void MainWindow::updateDrives(const QList<QStorageInfo> &drives) {
  ui->driveTable->setRowCount(0);

  for (const QStorageInfo &drive : drives) {
    int row = ui->driveTable->rowCount();
    ui->driveTable->insertRow(row);

    ui->driveTable->setItem(row, 0, new QTableWidgetItem(drive.name()));
    ui->driveTable->setItem(row, 1, new QTableWidgetItem(drive.device()));
    ui->driveTable->setItem(row, 2,
                            new QTableWidgetItem(drive.fileSystemType()));

    QString totalSize =
        QString::number(drive.bytesTotal() / (1024.0 * 1024.0 * 1024.0), 'f',
                        2) +
        " GB";
    ui->driveTable->setItem(row, 3, new QTableWidgetItem(totalSize));

    QString freeSpace =
        QString::number(drive.bytesAvailable() / (1024.0 * 1024.0 * 1024.0),
                        'f', 2) +
        " GB";
    ui->driveTable->setItem(row, 4, new QTableWidgetItem(freeSpace));
  }
}

void MainWindow::onShredClicked() {
  QList<QTableWidgetItem *> selectedItems = ui->driveTable->selectedItems();
  if (selectedItems.isEmpty()) {
    QMessageBox::warning(this, "No Drive Selected",
                         "Please select a drive to shred.");
    return;
  }

  int row = selectedItems.first()->row();
  QString drivePath =
      ui->driveTable->item(row, 1)->text(); // Path is in column 1
  QString driveName = ui->driveTable->item(row, 0)->text();

  QMessageBox::StandardButton reply;
  reply = QMessageBox::question(
      this, "Confirm Shredding",
      QString("Are you sure you want to PERMANENTLY ERASE all data on %1 "
              "(%2)?\n\nThis action cannot be undone!")
          .arg(driveName, drivePath),
      QMessageBox::Yes | QMessageBox::No);

  if (reply == QMessageBox::Yes) {
    // Unmount the drive first
    if (!unmountDrive(drivePath)) {
      QMessageBox::warning(this, "Unmount Failed",
                           "Failed to unmount the drive. It might be in use.");
      return;
    }

    // Store path for later use (formatting)
    m_currentDrivePath = drivePath;

    // Disconnect drive updates to prevent UI from clearing the row
    disconnect(m_driveManager, &DriveManager::drivesUpdated, this,
               &MainWindow::updateDrives);

    ui->shredButton->setEnabled(false);
    ui->driveTable->setEnabled(false);
    ui->progressBar->setValue(0);
    ui->elapsedLabel->setText("Elapsed: 00:00:00");
    ui->remainingLabel->setText("Remaining: Calculating...");

    Shredder::Method method =
        static_cast<Shredder::Method>(ui->methodCombo->currentIndex());
    m_shredder->startShredding(drivePath, method);
  }
}

void MainWindow::onShredProgress(int percent, const QString &status) {
  ui->progressBar->setValue(percent);
  if (!status.isEmpty()) {
    ui->statusbar->showMessage(status);
  }
}

void MainWindow::onTimerUpdated(const QString &elapsed,
                                const QString &remaining) {
  ui->elapsedLabel->setText("Elapsed: " + elapsed);
  ui->remainingLabel->setText("Remaining: " + remaining);
}

void MainWindow::onShredFinished(bool success, const QString &message) {
  ui->shredButton->setEnabled(true);
  ui->driveTable->setEnabled(true);
  ui->progressBar->setValue(success ? 100 : 0);
  ui->statusbar->showMessage(success ? "Shredding Complete"
                                     : "Shredding Failed");

  if (success) {
    QMessageBox::StandardButton reply;
    reply = QMessageBox::question(
        this, "Shredding Complete",
        "Shredding completed successfully.\n\nDo you want to FORMAT the drive "
        "now to make it reusable?",
        QMessageBox::Yes | QMessageBox::No);
    if (reply == QMessageBox::Yes) {
      QStringList filesystems;
      filesystems << "ext4" << "ntfs" << "fat32" << "exfat";

      bool ok;
      QString fs = QInputDialog::getItem(this, "Select Filesystem",
                                         "Choose a filesystem:", filesystems, 0,
                                         false, &ok);
      if (ok && !fs.isEmpty()) {
        ui->shredButton->setEnabled(false);
        ui->driveTable->setEnabled(false);
        ui->statusbar->showMessage("Formatting drive...");
        // Use the stored path since the UI might not have it selected or
        // updated
        m_formatter->formatDrive(m_currentDrivePath, fs);
      } else {
        // User cancelled format, reconnect signals
        connect(m_driveManager, &DriveManager::drivesUpdated, this,
                &MainWindow::updateDrives);
        updateDrives(m_driveManager->getDrives());
      }
    } else {
      QMessageBox::information(this, "Success", message);
      // Reconnect signals
      connect(m_driveManager, &DriveManager::drivesUpdated, this,
              &MainWindow::updateDrives);
      updateDrives(m_driveManager->getDrives());
    }
  } else {
    QMessageBox::critical(this, "Error", message);
    // Reconnect signals
    connect(m_driveManager, &DriveManager::drivesUpdated, this,
            &MainWindow::updateDrives);
    updateDrives(m_driveManager->getDrives());
  }
}

void MainWindow::onFormatFinished(bool success, const QString &message) {
  ui->shredButton->setEnabled(true);
  ui->driveTable->setEnabled(true);
  ui->statusbar->showMessage(success ? "Formatting Complete"
                                     : "Formatting Failed");

  // Reconnect signals
  connect(m_driveManager, &DriveManager::drivesUpdated, this,
          &MainWindow::updateDrives);
  updateDrives(m_driveManager->getDrives());

  if (success) {
    QMessageBox::information(this, "Success", message);
  } else {
    QMessageBox::critical(this, "Error", message);
  }
}

bool MainWindow::unmountDrive(const QString &path) {
  QProcess process;
  process.start("umount", QStringList() << path);
  if (!process.waitForStarted())
    return false;
  process.waitForFinished();
  return process.exitCode() == 0;
}