/* Copyright (C) 2008  Emmanuel Varoquaux

   This file is part of XOS.

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>. */

#include <getopt.h>
#include <error.h>
#include <string.h>
#include <stdlib.h>
#include <dirent.h>
#include <stdio.h>
#include <unistd.h>
#include <sys/stat.h>
#include <sys/ioctl.h>
#include <errno.h>

#define PROC_ROOT "/proc"

struct procinfo_struct {
  pid_t pid;
  pid_t ppid;
  pid_t pgid;
  unsigned long vsz;
  int state;
  unsigned foreground : 1;
  char *command;

  struct procinfo_struct *next;
};

struct fieldwidth_struct {
  int pid;
  int ppid;
  int pgid;
  int vsz;
  int state;
};

typedef int (*comparator_t)(const struct procinfo_struct *, const struct procinfo_struct *);

struct proc_list_struct {
  comparator_t comparator;
  struct procinfo_struct *first;
  struct fieldwidth_struct fieldwidth;
};

static const struct {
  const char *pid;
  const char *ppid;
  const char *pgid;
  const char *vsz;
  const char *state;
  const char *command;
} header = {
  .pid = "PID",
  .ppid = "PPID",
  .pgid = "PGID",
  .vsz = "VSZ",
  .state = "STATE",
  .command = "COMMAND"
};

static const char *proc_root = PROC_ROOT;

static int screenwidth;

static void get_stat_from_file(const char *filename, int *pid, int *ppid, int *pgid, unsigned long *vsz, int *state)
{
  FILE *fp;
  unsigned long vsize;
  unsigned char c;

  if (!(fp = fopen(filename, "r")))
    goto error;
  if (fscanf(fp, "%d %c %d %d %lu", pid, &c, ppid, pgid, &vsize) != 5)
    goto error_fclose;
  *vsz = vsize / 1024;
  *state = c;
  fclose(fp);
  return;

 error_fclose:
  fclose(fp);
 error:
  *pid = -1;
  *state = -1;
  *ppid = -1;
  *pgid = -1;
  *vsz = -1;
}

static void get_cmdline_from_file(const char *filename, char **s)
{
  struct stat statbuf;
  size_t size;
  FILE *fp;
  char *s1;
  size_t len;

  if (stat(filename, &statbuf) == -1)
    goto error;
  if ((size = statbuf.st_size)) {
    if (!(*s = malloc(size)))
      error(EXIT_FAILURE, errno, "malloc");
    if (!(fp = fopen(filename, "r")))
      goto error_free;
    if (fread(*s, 1, size, fp) < size)
      goto error_fclose;
    fclose(fp);
    s1 = *s;
    while (size) {
      if ((len = strnlen(s1, size)) == size)
        goto error_free;
      s1 += len;
      *s1 = ' ';
      s1++;
      if (size < len + 1)
        goto error_free;
      size -= len + 1;
    }
  }
  else
    if (!(*s = strdup("(defunct)")))
      error(EXIT_FAILURE, errno, "strdup");
  return;

 error_fclose:
  fclose(fp);
 error_free:
  free(*s);
 error:
  *s = NULL;
}

static void build_procinfo(struct procinfo_struct *procinfo, pid_t pid, const char *path)
{
  char *buf;
  pid_t fpgid;

  if (!(buf = malloc(strlen(path) + 1 + NAME_MAX + 1)))
    error(EXIT_FAILURE, errno, "malloc");
  sprintf(buf, "%s/stat", path);
  get_stat_from_file(buf, &procinfo->pid, &procinfo->ppid, &procinfo->pgid, &procinfo->vsz, &procinfo->state);
  if (procinfo->state != -1)
    procinfo->foreground = procinfo->pgid != -1 && (fpgid = tcgetpgrp(fileno(stderr))) != -1 && procinfo->pgid == fpgid ? 1 : 0;
  sprintf(buf, "%s/cmdline", path);
  get_cmdline_from_file(buf, &procinfo->command);
  free(buf);
}

static void destroy_procinfo(struct procinfo_struct *procinfo)
{
  if (procinfo->command)
    free(procinfo->command);
}

static void set_header_fieldwidth(struct fieldwidth_struct *fieldwidth)
{
  fieldwidth->pid = strlen(header.pid);
  fieldwidth->ppid = strlen(header.ppid);
  fieldwidth->pgid = strlen(header.pgid);
  fieldwidth->vsz = strlen(header.vsz);
  fieldwidth->state = strlen(header.state);
}

static int num_fieldwidth(int n)
{
  return snprintf(NULL, 0, "%d", n);
}

static void update_fieldwidth(struct fieldwidth_struct *fieldwidth, const struct procinfo_struct *procinfo)
{
  int w;

  if ((w = num_fieldwidth(procinfo->pid)) > fieldwidth->pid)
    fieldwidth->pid = w;
  if ((w = procinfo->ppid != -1 ? num_fieldwidth(procinfo->ppid) : 1) > fieldwidth->ppid)
    fieldwidth->ppid = w;
  if ((w = procinfo->pgid != -1 ? num_fieldwidth(procinfo->pgid) : 1) > fieldwidth->pgid)
    fieldwidth->pgid = w;
  if ((w = procinfo->vsz != (unsigned long)-1 ? num_fieldwidth(procinfo->vsz) : 1) > fieldwidth->vsz)
    fieldwidth->vsz = w;
  if (fieldwidth->state < 2)
    fieldwidth->state = 2;
}

static void init_proc_list(struct proc_list_struct *proc_list, comparator_t comparator)
{
  proc_list->comparator = comparator;
  proc_list->first = NULL;
  set_header_fieldwidth(&proc_list->fieldwidth);
}

static void destroy_proc_list(struct proc_list_struct *proc_list)
{
  struct procinfo_struct *procinfo, *next;

  procinfo = proc_list->first;
  while (procinfo) {
    next = procinfo->next;
    destroy_procinfo(procinfo);
    free(procinfo);
    procinfo = next;
  }
}

static int pid_comparator(const struct procinfo_struct *procinfo1, const struct procinfo_struct *procinfo2)
{
  return procinfo1->pid < procinfo2->pid;
}

static int compare(const struct procinfo_struct *procinfo1, const struct procinfo_struct *procinfo2, struct proc_list_struct *proc_list)
{
  return proc_list->comparator(procinfo1, procinfo2);
}

static void insert_into_proc_list(struct proc_list_struct *proc_list, struct procinfo_struct *procinfo)
{
  struct procinfo_struct *procinfo1, *procinfo2;

  procinfo1 = NULL;
  procinfo2 = proc_list->first;
  while (procinfo2 && !compare(procinfo, procinfo2, proc_list)) {
    procinfo1 = procinfo2;
    procinfo2 = procinfo2->next;
  }
  procinfo->next = procinfo2;
  if (!procinfo1)
    proc_list->first = procinfo;
  else
    procinfo1->next = procinfo;

  update_fieldwidth(&proc_list->fieldwidth, procinfo);
}

static int populate_proc_list(struct proc_list_struct *proc_list)
{
  int err;
  DIR *dir;
  struct dirent *dent;
  int num;
  char *buf;
  struct procinfo_struct *procinfo;

  err = 0;
  if (!(dir = opendir(proc_root))) {
    error(0, errno, "cannot open %s", proc_root);
    err = 1;
    goto end;
  }
  while (1) {
    errno = 0;
    if (!(dent = readdir(dir))) {
      if (errno) {
        error(0, errno, "reading directory %s", proc_root);
        err = 1;
        goto closedir;
      }
      break;
    }
    if (sscanf(dent->d_name, "%d", &num) != 1)
      continue;
    if (!(buf = malloc(strlen(proc_root) + 1 + strlen(dent->d_name) + 1)))
      error(EXIT_FAILURE, errno, "malloc");
    if (!(procinfo = malloc(sizeof (struct procinfo_struct))))
      error(EXIT_FAILURE, errno, "malloc");
    sprintf(buf, "%s/%s", proc_root, dent->d_name);
    build_procinfo(procinfo, num, buf);
    free(buf);
    insert_into_proc_list(proc_list, procinfo);
  }
 closedir:
  if (closedir(dir) == -1) {
    error(0, errno, "closing directory %s", proc_root);
    err = 1;
  }
 end:
  return !err;
}

static void print_header(const struct fieldwidth_struct *fieldwidth)
{
  printf("%-*s", fieldwidth->pid, header.pid);
  putchar(' ');
  printf("%-*s", fieldwidth->ppid, header.ppid);
  putchar(' ');
  printf("%-*s", fieldwidth->pgid, header.pgid);
  putchar(' ');
  printf("%-*s", fieldwidth->vsz, header.vsz);
  putchar(' ');
  printf("%-*s", fieldwidth->state, header.state);
  putchar(' ');
  printf("%s", header.command);
  putchar('\n');
}

static char *make_state_str(const struct procinfo_struct *procinfo, char *buf)
{
  char *p;

  p = buf;
  *p++ = procinfo->state;
  if (procinfo->foreground)
    *p++ = '+';
  *p = '\0';
  return buf;
}

static void print_procinfo(const struct procinfo_struct *procinfo, const struct fieldwidth_struct *fieldwidth)
{
  int rem_width;
  char state_buf[3];

  rem_width = screenwidth;
  if (procinfo->pid != -1)
    printf("%*d", fieldwidth->pid, procinfo->pid);
  else
    printf("%*s", fieldwidth->pid, "?");
  putchar(' ');
  rem_width -= fieldwidth->pid + 1;
  if (procinfo->ppid != -1)
    printf("%*d", fieldwidth->ppid, procinfo->ppid);
  else
    printf("%*s", fieldwidth->ppid, "?");
  putchar(' ');
  rem_width -= fieldwidth->ppid + 1;
  if (procinfo->pgid != -1)
    printf("%*d", fieldwidth->pgid, procinfo->pgid);
  else
    printf("%*s", fieldwidth->pgid, "?");
  putchar(' ');
  rem_width -= fieldwidth->pgid + 1;
  if (procinfo->vsz != (unsigned long)-1)
    printf("%*lu", fieldwidth->vsz, procinfo->vsz);
  else
    printf("%*s", fieldwidth->vsz, "?");
  putchar(' ');
  rem_width -= fieldwidth->vsz + 1;
  printf("%-*s", fieldwidth->state, procinfo->state != -1 ? make_state_str(procinfo, state_buf) : "?");
  putchar(' ');
  rem_width -= fieldwidth->state + 1;
  if (screenwidth >= 0)
    printf("%.*s", rem_width, procinfo->command ? : "?");
  else
    printf("%s", procinfo->command ? : "?");
  putchar('\n');
}

static void print_proc_list(const struct proc_list_struct *proc_list)
{
  const struct procinfo_struct *procinfo;

  print_header(&proc_list->fieldwidth);
  for (procinfo = proc_list->first; procinfo; procinfo = procinfo->next)
    print_procinfo(procinfo, &proc_list->fieldwidth);
}

static int list_processes(comparator_t comparator)
{
  int err;
  struct proc_list_struct proc_list;

  err = 0;
  init_proc_list(&proc_list, comparator);
  if (!populate_proc_list(&proc_list))
    err = 1;
  print_proc_list(&proc_list);
  destroy_proc_list(&proc_list);
  return !err;
}

static void __attribute__ ((noreturn)) die_bad_usage()
{
  fprintf(stderr, "Try `%s --help' for more information.\n", program_invocation_name);
  exit(EXIT_FAILURE);
}

static void print_help()
{
  printf("Usage: %s [OPTION]...\n", program_invocation_name);
  puts("Print a snapshot of the current processes.");
  putchar('\n');
  puts("  -h, --help  display this help and exit");
}

int main(int argc, char **argv)
{
  static const struct option longopts[] = {
    {"help", no_argument, NULL, 'h'},
    {NULL, 0, NULL, 0}
  };

  int c;
  struct winsize winsize;

  while ((c = getopt_long(argc, argv, "h", longopts, NULL)) != -1)
    switch (c) {
    case 'h':
      print_help();
      return EXIT_SUCCESS;
    case '?':
      die_bad_usage();
    }

  screenwidth = ioctl(fileno(stdout), TIOCGWINSZ, &winsize) != -1 ? winsize.ws_col : -1;
  return list_processes(pid_comparator) ? EXIT_SUCCESS : EXIT_FAILURE;
}