main.cpp

#include <algorithm>
#include <cctype>
#include <chrono>
#include <cstdlib>
#include <iomanip>
#include <iostream>
#include <libproc.h>
#include <map>
#include <print>
#include <sstream>
#include <string>
#include <sys/ioctl.h>
#include <sys/stat.h>
#include <thread>
#include <unistd.h>
#include <vector>

// ANSI color codes
const std::string COLOR_RESET = "\033[0m";
const std::string COLOR_BLUE = "\033[34m";
const std::string COLOR_GREEN = "\033[32m";
const std::string COLOR_YELLOW = "\033[33m";
const std::string COLOR_CYAN = "\033[36m";
const std::string COLOR_MAGENTA = "\033[35m";
const std::string COLOR_BOLD = "\033[1m";
const std::string COLOR_DIM = "\033[2m";

struct FileInfo {
  std::string path;
  off_t size;
  off_t position;
  int fd;
  std::chrono::steady_clock::time_point lastUpdate;
  off_t lastPosition;
  double speed; // bytes per second
};

struct ProcessInfo {
  pid_t pid;
  std::string name;
  std::string path;
};

bool containsAlpha(const std::string &str) {
  return std::any_of(str.begin(), str.end(),
                     [](char c) { return std::isalpha(c); });
}

std::vector<ProcessInfo> findProcessesByName(const std::string &searchTerm) {
  std::vector<ProcessInfo> matchingProcesses;

  // Get list of all processes
  int numPids = proc_listpids(PROC_ALL_PIDS, 0, nullptr, 0);
  if (numPids <= 0)
    return matchingProcesses;

  std::vector<pid_t> pids(numPids);
  int actualSize =
      proc_listpids(PROC_ALL_PIDS, 0, pids.data(), numPids * sizeof(pid_t));

  if (actualSize <= 0)
    return matchingProcesses;

  int actualNumPids = actualSize / sizeof(pid_t);

  for (int i = 0; i < actualNumPids; ++i) {
    if (pids[i] <= 0)
      continue;

    char pathBuffer[PROC_PIDPATHINFO_MAXSIZE];
    int pathLen = proc_pidpath(pids[i], pathBuffer, sizeof(pathBuffer));

    if (pathLen > 0) {
      std::string fullPath(pathBuffer);
      std::string processName = fullPath.substr(fullPath.find_last_of('/') + 1);

      // Case-insensitive search
      std::string lowerProcessName = processName;
      std::string lowerSearchTerm = searchTerm;
      std::transform(lowerProcessName.begin(), lowerProcessName.end(),
                     lowerProcessName.begin(), ::tolower);
      std::transform(lowerSearchTerm.begin(), lowerSearchTerm.end(),
                     lowerSearchTerm.begin(), ::tolower);

      if (lowerProcessName.find(lowerSearchTerm) != std::string::npos) {
        ProcessInfo procInfo;
        procInfo.pid = pids[i];
        procInfo.name = processName;
        procInfo.path = fullPath;
        matchingProcesses.push_back(procInfo);
      }
    }
  }

  return matchingProcesses;
}

std::string formatBytes(uint64_t bytes) {
  const char *units[] = {"B", "KB", "MB", "GB", "TB", "PB"};
  const int numUnits = sizeof(units) / sizeof(units[0]);

  if (bytes == 0)
    return "0 B";

  int unitIndex = 0;
  double size = static_cast<double>(bytes);

  while (size >= 1024.0 && unitIndex < numUnits - 1) {
    size /= 1024.0;
    unitIndex++;
  }

  std::ostringstream oss;
  if (unitIndex == 0) {
    oss << static_cast<uint64_t>(size) << " " << units[unitIndex];
  } else {
    oss << std::fixed << std::setprecision(1) << size << " "
        << units[unitIndex];
  }

  return oss.str();
}

std::string formatTime(double seconds) {
  if (seconds < 0 || std::isinf(seconds) || std::isnan(seconds)) {
    return "--:--";
  }

  int hours = static_cast<int>(seconds) / 3600;
  int minutes = (static_cast<int>(seconds) % 3600) / 60;
  int secs = static_cast<int>(seconds) % 60;

  if (hours > 0) {
    return std::to_string(hours) + "h " + std::to_string(minutes) + "m";
  } else if (minutes > 0) {
    return std::to_string(minutes) + "m " + std::to_string(secs) + "s";
  } else {
    return std::to_string(secs) + "s";
  }
}

int getTerminalWidth() {
  struct winsize w;
  if (ioctl(STDOUT_FILENO, TIOCGWINSZ, &w) == 0) {
    return w.ws_col;
  }
  return 80; // Default fallback width
}

std::string getProgressBarColor(double progress) {
  if (progress < 0.3)
    return COLOR_YELLOW;
  if (progress < 0.7)
    return COLOR_BLUE;
  return COLOR_GREEN;
}

void printProgressBar(const FileInfo &fileInfo) {
  int terminalWidth = getTerminalWidth();

  // Calculate space needed for non-progress-bar elements
  std::string filename =
      fileInfo.path.substr(fileInfo.path.find_last_of('/') + 1);
  if (filename.length() > 25) {
    filename = filename.substr(0, 22) + "...";
  }

  // Estimate space for: filename(25) + " [" + "] " + percentage(7) + " (" +
  // sizes + ") | " + speed + " | ETA: " + eta
  int fixedWidth = 25 + 2 + 2 + 7 + 3 + 20 + 3 + 15 + 8 + 10; // Rough estimate
  int progressBarWidth = std::max(
      20, terminalWidth - fixedWidth); // Minimum 20 chars for progress bar
  progressBarWidth =
      std::min(progressBarWidth, 60); // Maximum 60 chars to keep it reasonable
  if (fileInfo.size <= 0) {
    std::cout << filename << ": [Unknown size]" << std::endl;
    return;
  }

  double progress = static_cast<double>(fileInfo.position) / fileInfo.size;
  int filled = static_cast<int>(progress * progressBarWidth);

  std::print("{}{:<25}{} [", COLOR_CYAN, filename, COLOR_RESET);
  std::string barColor = getProgressBarColor(progress);
  std::print("{}", barColor);
  for (int i = 0; i < progressBarWidth; ++i) {
    std::print("{}", (i < filled) ? "█" : "░");
  }
  std::print("{}", COLOR_RESET);

  std::print("] {}{:>6.1f}%{} ({}{}/{}{})", COLOR_BOLD, progress * 100,
             COLOR_RESET, COLOR_MAGENTA, formatBytes(fileInfo.position),
             formatBytes(fileInfo.size), COLOR_RESET);

  // Show speed and ETA
  if (fileInfo.speed > 0) {
    std::print("{} │ {}{}/s{}", COLOR_DIM, COLOR_GREEN,
               formatBytes(static_cast<uint64_t>(fileInfo.speed)), COLOR_RESET);

    double remainingBytes = fileInfo.size - fileInfo.position;
    double eta = remainingBytes / fileInfo.speed;
    std::print("{} │ {}ETA: {}{}{}", COLOR_DIM, COLOR_RESET, COLOR_YELLOW,
               formatTime(eta), COLOR_RESET);
  } else {
    std::print("{} │ {}--/s{} │ {}ETA: --:--", COLOR_DIM, COLOR_RESET,
               COLOR_DIM, COLOR_RESET);
  }

  std::println("");
}

std::string getProcessCommand(pid_t pid) {
#ifdef __APPLE__
  char pathbuf[PROC_PIDPATHINFO_MAXSIZE];
  char command[1024];
  int ret = proc_pidpath(pid, pathbuf, sizeof(pathbuf));
  if (ret > 0) {
    FILE *pipe =
        popen(("ps -p " + std::to_string(pid) + " -o command=").c_str(), "r");
    if (pipe) {
      if (fgets(command, sizeof(command), pipe) != nullptr) {
        command[strcspn(command, "\n")] = 0;
        pclose(pipe);
        return std::string(command);
      }
      pclose(pipe);
    }
  }
  return "";
#else
  std::string cmdPath = "/proc/" + std::to_string(pid) + "/cmdline";
  std::ifstream cmdFile(cmdPath);
  std::string cmd;
  if (cmdFile.is_open()) {
    std::getline(cmdFile, cmd);
    cmdFile.close();
    return cmd;
  }
  return "";
#endif
}

std::vector<FileInfo> getOpenFiles(pid_t pid) {
  std::vector<FileInfo> files;

  // Get file descriptor info
  int bufferSize = proc_pidinfo(pid, PROC_PIDLISTFDS, 0, nullptr, 0);
  if (bufferSize <= 0)
    return files;

  std::vector<proc_fdinfo> fdInfos(bufferSize / sizeof(proc_fdinfo));
  int actualSize =
      proc_pidinfo(pid, PROC_PIDLISTFDS, 0, fdInfos.data(), bufferSize);

  if (actualSize <= 0)
    return files;

  int numFds = actualSize / sizeof(proc_fdinfo);

  for (int i = 0; i < numFds; ++i) {
    if (fdInfos[i].proc_fdtype == PROX_FDTYPE_VNODE) {
      vnode_fdinfowithpath vnodeInfo;
      int result =
          proc_pidfdinfo(pid, fdInfos[i].proc_fd, PROC_PIDFDVNODEPATHINFO,
                         &vnodeInfo, sizeof(vnodeInfo));

      if (result == sizeof(vnodeInfo)) {
        FileInfo fileInfo;
        fileInfo.path = vnodeInfo.pvip.vip_path;
        fileInfo.fd = fdInfos[i].proc_fd;
        fileInfo.position = vnodeInfo.pfi.fi_offset;

        // Get file size using stat
        struct stat fileStat;
        if (stat(fileInfo.path.c_str(), &fileStat) == 0) {
          fileInfo.size = fileStat.st_size;
        } else {
          fileInfo.size = 0;
        }

        // Only include regular files that are readable
        if (!fileInfo.path.empty() && fileInfo.size > 0 &&
            S_ISREG(fileStat.st_mode)) {
          // Initialize timing info
          fileInfo.lastUpdate = std::chrono::steady_clock::now();
          fileInfo.lastPosition = fileInfo.position;
          fileInfo.speed = 0.0;
          files.push_back(fileInfo);
        }
      }
    }
  }

  return files;
}
bool hasShownFiles = false;

void showProgress(const std::vector<pid_t> &pids, int intervalMs = 1000) {
  std::print("\033[2J\033[H"); // Clear screen and move cursor to top
  std::println("{}{}{} for {}{} process(es){}", COLOR_BOLD,
               "File Progress Monitor", COLOR_RESET, COLOR_CYAN, pids.size(),
               COLOR_RESET);
  std::println("{}{}{}", COLOR_DIM, "Press Ctrl+C to exit\n", COLOR_RESET);

  std::map<pid_t, std::map<std::string, FileInfo>> previousFiles;

  while (true) {
    std::map<pid_t, std::vector<FileInfo>> allFiles;
    auto currentTime = std::chrono::steady_clock::now();

    // Collect files from all processes
    for (pid_t pid : pids) {
      allFiles[pid] = getOpenFiles(pid);
    }

    // Calculate speed and display progress for each process
    std::print("\033[3;1H"); // Move cursor to line 3
    bool hasFiles = false;

    for (const auto &[pid, files] : allFiles) {
      if (!files.empty()) {
        hasFiles = true;
        hasShownFiles = true;
        std::string cmd = getProcessCommand(pid);
        std::println("[PID {}{}{}{}] {}{}{}", COLOR_BOLD, COLOR_CYAN, pid,
                     COLOR_RESET, COLOR_DIM, cmd, COLOR_RESET);
        // Top border
        std::println("+{}{}+{}", COLOR_DIM,
                     std::string(std::min(118, getTerminalWidth() - 2), '-'),
                     COLOR_RESET);

        for (auto file : files) {
          auto &prevFiles = previousFiles[pid];
          auto it = prevFiles.find(file.path);
          if (it != prevFiles.end()) {
            auto timeDiff =
                std::chrono::duration_cast<std::chrono::milliseconds>(
                    currentTime - it->second.lastUpdate)
                    .count();
            if (timeDiff > 0) {
              off_t positionDiff = file.position - it->second.lastPosition;
              file.speed =
                  (static_cast<double>(positionDiff) * 1000.0) / timeDiff;

              if (it->second.speed > 0) {
                file.speed = 0.7 * it->second.speed + 0.3 * file.speed;
              }
            }
          }
          file.lastUpdate = currentTime;
          file.lastPosition = file.position;
          printProgressBar(file);
        }
        // Bottom border
        std::println("+{}{}+{}", COLOR_DIM,
                     std::string(std::min(118, getTerminalWidth() - 2), '-'),
                     COLOR_RESET);
        std::cout << std::endl;
      }

      // Update previous files for this process
      previousFiles[pid].clear();
      for (const auto &file : files) {
        previousFiles[pid][file.path] = file;
      }
    }
    if (hasShownFiles && !hasFiles) {
      // clear screen:
      std::print("\033[2J\033[H");

      // exit process:
      exit(0);
    }
    if (!hasFiles) {
      std::println("No regular files currently open by any monitored process.");
    }

    std::this_thread::sleep_for(std::chrono::milliseconds(intervalMs));
  }
}

void showUsage(const char *progName) {
  std::println(
      "Usage: {} [--progress|-p] [--interval|-i <ms>] <pid|process_name>",
      progName);
  std::println("\nOptions:");
  std::println("  --progress, -p       Show real-time progress of open files");
  std::println("  --interval, -i <ms>  Update interval in milliseconds "
               "(default: 2000, min: 100)");
  std::println("\nExamples:");
  std::println("  {} 1234          # Show disk I/O statistics for PID 1234",
               progName);
  std::println("  {} chrome        # Find and show stats for processes "
               "containing 'chrome'",
               progName);
  std::println("  {} --progress 1234  # Monitor file progress for PID 1234",
               progName);
  std::println("  {} --progress zstd  # Monitor file progress for processes "
               "containing 'zstd'",
               progName);
}

int main(int argc, char *argv[]) {
  bool progressMode = false;
  int intervalMs = 1000; // Default 2 seconds
  std::string input;

  // Parse command line arguments
  for (int i = 1; i < argc; ++i) {
    std::string arg = argv[i];
    if (arg == "--progress" || arg == "-p") {
      progressMode = true;
    } else if (arg == "--interval" || arg == "-i") {
      if (i + 1 < argc) {
        intervalMs = std::atoi(argv[++i]);
        if (intervalMs < 100) {
          intervalMs = 100; // Clamp to minimum 100ms
        }
      } else {
        std::cerr << "Error: --interval requires a value in milliseconds"
                  << std::endl;
        return 1;
      }
    } else {
      input = arg;
    }
  }

  if (input.empty()) {
    showUsage(argv[0]);
    return 1;
  }

  std::vector<pid_t> targetPids;

  // Check if input contains alphabetic characters (process name search)
  if (containsAlpha(input)) {
    std::vector<ProcessInfo> processes = findProcessesByName(input);

    if (processes.empty()) {
      std::println("No processes found containing '{}'", input);
      return 1;
    }

    std::println("Found {} process(es) containing '{}':", processes.size(),
                 input);
    std::println("{}", std::string(80, '-'));

    for (const auto &proc : processes) {
      targetPids.push_back(proc.pid);
      std::println("{:>8}  {}", proc.pid, proc.name);
    }
    std::println("{}", std::string(80, '-'));

  } else {
    // Numeric PID
    pid_t pid = std::atoi(input.c_str());
    if (pid <= 0) {
      std::println("Error: Invalid PID. Please provide a positive integer or "
                   "process name.");
      return 1;
    }
    targetPids.push_back(pid);
  }

  if (progressMode) {
    showProgress(targetPids, intervalMs);
    return 0;
  }

  // Show disk I/O statistics for all matching processes
  for (size_t i = 0; i < targetPids.size(); ++i) {
    pid_t pid = targetPids[i];

    rusage_info_current rusage;
    int result = proc_pid_rusage(pid, RUSAGE_INFO_CURRENT, (void **)&rusage);

    if (result != 0) {
      std::println("Error: Failed to get resource usage for PID {}", pid);
      std::println(
          "Make sure the process exists and you have permission to access it.");
      continue;
    }

    if (targetPids.size() > 1) {
      std::println("\nDisk I/O Statistics for PID {}:", pid);
    } else {
      std::println("Disk I/O Statistics for PID {}:", pid);
    }

    std::println("Bytes Read:    {} ({} bytes)",
                 formatBytes(rusage.ri_diskio_bytesread),
                 rusage.ri_diskio_bytesread);
    std::println("Bytes Written: {} ({} bytes)",
                 formatBytes(rusage.ri_diskio_byteswritten),
                 rusage.ri_diskio_byteswritten);
  }

  return 0;
}