#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>
#include <dirent.h>
#include <errno.h>

#include <sys/utsname.h>
#ifdef __linux__
#include <sys/sysinfo.h>
#elif __FreeBSD__
#elif __OpenBSD__
#endif
#include <sys/statvfs.h>

#ifdef __linux__
#include <pci/pci.h>
#elif __FreeBSD__
#elif __OpenBSD__
#endif

#include <X11/Xlib.h>
#include <X11/Xatom.h>
#ifdef __linux__
#include <wayland-client.h>
#include <wayland-client-protocol.h>
#endif

#include "paleofetch.h"
#include "config.h"

#define BUF_SIZE 150
#define COUNT(x) (int)(sizeof x / sizeof *x)

#define halt_and_catch_fire(fmt, ...) \
  do { \
    if (status != 0) { \
      fprintf(stderr, "paleofetch: " fmt "\n", ##__VA_ARGS__); \
      exit(status); \
    } \
  } while(0)

struct conf {
  char *label, *(*function)();
  bool cached;
} config[] = CONFIG;

struct {
  char *substring;
  char *repl_str;
  size_t length;
  size_t repl_len;
} cpu_config[] = CPU_CONFIG, gpu_config[] = GPU_CONFIG;

struct statvfs file_stats;
struct utsname uname_info;
struct sysinfo my_sysinfo;
int title_length, status;

int remove_newline(char *s) {
  int i;
  for (i = 0; s[i] != '\0' && s[i] != '\n'; ++i);
  s[i] = '\0';
  return i;
}

void truncate_spaces(char *str) {
  char *src = str, *dst = str;

  while (*dst == ' ') ++dst;

  while ((*src = *dst) != '\0') {
    if (*(dst++) == ' ') {
      while (*dst == ' ') ++dst;
    }
    ++src;
  }
}

void remove_substring(char *str, const char* substring, size_t len) {
  char *sub = strstr(str, substring);
  if (sub) {
    memmove(sub, sub + len, strlen(sub + len) + 1);
  }
}

void replace_substring(char *str, const char *sub_str, const char *repl_str, size_t sub_len, size_t repl_len) {
  char buffer[BUF_SIZE / 2];
  char *start = strstr(str, sub_str);
  if (start) {
    size_t tail_len = strlen(start + sub_len);
    if (tail_len + repl_len < BUF_SIZE / 2) {
      strcpy(buffer, start + sub_len);
      strncpy(start, repl_str, repl_len);
      strcpy(start + repl_len, buffer);
    } else {
      status = -1;
      halt_and_catch_fire("new substring too long to replace");
    }
  }
}

static char *get_title() {
  char hostname[BUF_SIZE / 3];
  status = gethostname(hostname, BUF_SIZE / 3);
  halt_and_catch_fire("unable to retrieve host name");

  char *username = getenv("USER");
  if (username == NULL) {
    halt_and_catch_fire("unable to retrieve login name");
  }

  title_length = strlen(hostname) + strlen(username) + 1;

  char *title = malloc(BUF_SIZE);
  snprintf(title, BUF_SIZE, TITLECOLOR"%s\e[0m@"TITLECOLOR"%s", username, hostname);

  return title;
}

static char *get_bar() {
  char *bar = malloc(BUF_SIZE);
  char *s = bar;
  for (int i = 0; i < title_length; i++) {
    *(s++) = '-';
  }
  *s = '\0';
  return bar;
}

static char *get_os() {
  char *os = malloc(BUF_SIZE), *name = malloc(BUF_SIZE), *line = NULL;
  size_t len;
  char *version = malloc(BUF_SIZE);
  int getname = 0, getver = 0;

  FILE *os_release = fopen("/etc/os-release", "r");
  if (os_release == NULL) {
    status = -1;
    halt_and_catch_fire("unable to open /etc/os-release");
  }

  while (getline(&line, &len, os_release) != -1) {
    if (!getname && strstr(line, "NAME=") == line) {
      if (line[5] == '"') {
        sscanf(line, "NAME=\"%[^\"]\"", name);
      } else {
        sscanf(line, "NAME=%[^\n]", name);
      }
      getname = 1;
    }

    if (!getver && strstr(line, "VERSION=") == line) {
      if (line[8] == '"') {
        sscanf(line, "VERSION=\"%[^\"]\"", version);
      } else {
        sscanf(line, "VERSION=%[^\n]", version);
      }
      getver = 1;
    }

    if (getname && getver) {
      break;
    }
  }

  if (getver) {
    snprintf(os, BUF_SIZE, "%s %s %s", name, version, uname_info.machine);
  } else {
    snprintf(os, BUF_SIZE, "%s %s", name, uname_info.machine);
  }

  free(line);
  fclose(os_release);
  free(name);
  free(version);

  return os;
}

static char *get_kernel() {
  struct utsname uname_info;

  if (uname(&uname_info) == -1) {
    return "不明";
  }

  char *kernel = malloc(BUF_SIZE);
  if (kernel == NULL) {
    return "不明";
  }

  strncpy(kernel, uname_info.release, BUF_SIZE);
  kernel[BUF_SIZE -1] = '\0';

  return kernel;
}

static char *get_host() {
  FILE *product_name, *product_version, *model;
  char buffer[BUF_SIZE / 2];
  char *host = NULL;

  product_name = fopen("/sys/devices/virtual/dmi/id/product_name", "r");
  if (product_name != NULL) {
    host = malloc(BUF_SIZE);
    size_t read_len = fread(host, 1, BUF_SIZE / 2, product_name);
    if (read_len > 0) {
      remove_newline(host);

      product_version = fopen("/sys/devices/virtual/dmi/id/product_version", "r");
      if (product_version != NULL) {
        read_len = fread(buffer, 1, BUF_SIZE / 2, product_version);
        if (read_len > 0) {
          remove_newline(buffer);
          if (strcmp(buffer, "To Be Filled By O.E.M.") != 0) {
            size_t host_len = strlen(host);
            size_t remain_size = (BUF_SIZE - host_len) - 1;
            if (remain_size > 0) {
              snprintf(host + host_len, remain_size, " %s", buffer);
            }
          }
        }
        fclose(product_version);
      }
      fclose(product_name);
      return host;
    }
    fclose(product_name);
  }

  model = fopen("/sys/firmware/devicetree/base/model", "r");
  if (model != NULL) {
    host = malloc(BUF_SIZE);
    size_t read_len = fread(host, 1, BUF_SIZE, model);
    if (read_len > 0) {
      remove_newline(host);
      return host;
    }
    fclose(model);
  }

  status = -1;
  halt_and_catch_fire("unable to get host");
  return NULL;
}

static char *get_uptime() {
  long seconds = my_sysinfo.uptime;
  struct { char *name; int secs; } units[] = {
    { "day",  60 * 60 * 24 },
    { "hour", 60 * 60 },
    { "min",  60 },
  };

  int n, len = 0;
  char *uptime = malloc(BUF_SIZE);
  for (int i = 0; i < 3; ++i ) {
    if ((n = seconds / units[i].secs) || i == 2) {
      len += snprintf(uptime + len, BUF_SIZE - len, "%d %s%s, ", n, units[i].name, n != 1 ? "s": "");
    }
    seconds %= units[i].secs;
  }

  uptime[len - 2] = '\0';
  return uptime;
}

static char *get_packages() {
  char *packages = malloc(BUF_SIZE);
  snprintf(packages, BUF_SIZE, "%s", PACKAGES);
  return packages;
}

static char *get_shell() {
  char *shell = malloc(BUF_SIZE);
  char *shell_path = getenv("SHELL");
  char *shell_name = strrchr(getenv("SHELL"), '/');

  if (shell_name == NULL) {
    strncpy(shell, shell_path, BUF_SIZE - 1);
  } else {
    strncpy(shell, shell_name + 1, BUF_SIZE - 1);
  }

  return shell;
}

static char *get_resolution() {
  int screen, width, height;
  char *resolution = malloc(BUF_SIZE);
  char *wayland_display = getenv("WAYLAND_DISPLAY");

  if (wayland_display != NULL) { // TODO
    struct wl_display *dp = wl_display_connect(NULL);
    if (dp == NULL) {
      exit(EXIT_FAILURE);
    }

    struct wl_registry *rg = wl_display_get_registry(dp);
    if (rg == NULL) {
      wl_display_disconnect(dp);
      exit(EXIT_FAILURE);
    }

    wl_display_disconnect(dp);
  } else {
    Display *display = XOpenDisplay(NULL);
    if (display != NULL) {
      screen = DefaultScreen(display);
    
      width = DisplayWidth(display, screen);
      height = DisplayHeight(display, screen);

      snprintf(resolution, BUF_SIZE, "%dx%d", width, height);
    } else {
      DIR *dir;
      struct dirent *entry;
      char dir_name[] = "/sys/class/drm";
      char modes_file_name[BUF_SIZE * 2];
      FILE *modes;
      char *line = NULL;
      size_t len;

      strncpy(resolution, "", BUF_SIZE);

      dir = opendir(dir_name);
      if (dir == NULL) {
        status = -1;
        halt_and_catch_fire("Could not open /sys/class/drm to determine resolution in tty mode.");
      }

      while ((entry = readdir(dir)) != NULL) {
        if (entry->d_type == DT_LNK) {
          snprintf(modes_file_name, BUF_SIZE * 2, "%s/%s/modes", dir_name, entry->d_name);

          modes = fopen(modes_file_name, "r");
          if (modes != NULL) {
            if (getline(&line, &len, modes) != -1) {
              strncpy(resolution, line, BUF_SIZE - 1);
              remove_newline(resolution);

              free(line);
              fclose(modes);

              break;
            }

            fclose(modes);
          }
        }
      }

      closedir(dir);
    }
  }

  return resolution;
}

static char *get_wm() {
  char *wayland_display = getenv("WAYLAND_DISPLAY");
  if (wayland_display != NULL) {
    char *wayland_wms[] = {
      "[a]rcan", "[a]sc", "[c]layland", "[d]wc", "[d]wl", "[f]ireplace",
      "[g]nome-shell", "[g]reenfield", "[g]refsen", "[k]win",
      "[l]ipstick", "[m]aynard", "[m]azecompositor", "[m]otorcar",
      "[o]rbital", "[o]rbment", "[p]erceptia", "[r]ustland", "[s]way",
      "[u]lubis", "[v]elox", "[w]avy", "[w]ay-cooler", "[w]ayfire",
      "[w]ayhouse", "[w]esteros", "[w]estford", "[w]eston"
    };

    char command[1024] = "ps aux | grep -m 1 -E -o '";
    size_t num_wms = sizeof(wayland_wms) / sizeof(wayland_wms[0]);
    for (size_t i = 0; i < num_wms - 1; i++) {
      strcat(command, wayland_wms[i]);
      strcat(command, "|");
    }
    strcat(command, wayland_wms[sizeof(wayland_wms) / sizeof(wayland_wms[0]) - 1]);
    strcat(command, "'");

    FILE *fp = popen(command, "r");
    if (fp == NULL) {
      perror("popen");
      return strdup("不明");
    }

    char wm_name[256];
    if (fgets(wm_name, sizeof(wm_name), fp) != NULL) {
      char *newline = strchr(wm_name, '\n');
      if (newline) {
        *newline = '\0';
      }
    } else {
      strcpy(wm_name, "不明");
    }

    pclose(fp);
    return strdup(wm_name);
  }

  Display *dpy = XOpenDisplay(NULL);
  if (dpy == NULL) {
    //fprintf(stderr, "Unable to connect to display\n");
    return strdup("不明");
  }

  int screen = DefaultScreen(dpy);
  Window root = RootWindow(dpy, screen);

  Atom atom = XInternAtom(dpy, "_NET_SUPPORTING_WM_CHECK", False);
  Atom type;
  int format;
  unsigned long nitems, bytes_after;
  unsigned char *prop;
  int status = XGetWindowProperty(dpy, root, atom, 0, (~0L), False, AnyPropertyType, &type, &format, &nitems, &bytes_after, &prop);
  if (status != Success || !prop) {
    fprintf(stderr, "Unable to get window manager property\n");
    XCloseDisplay(dpy);
    return strdup("不明");
  }

  Window wm_window = *(Window *)prop;
  XFree(prop);

  status = XGetWindowProperty(dpy, wm_window, XInternAtom(dpy, "_NET_WM_NAME", False), 0, (~0L), False, AnyPropertyType, &type, &format, &nitems, &bytes_after, &prop);
  if (status != Success || !prop) {
    fprintf(stderr, "Unable to get window manager name property\n");
    XCloseDisplay(dpy);
    return strdup("不明");
  }

  char *wm_name = strdup((char *)prop);
  XFree(prop);
  XCloseDisplay(dpy);

  return wm_name;
}

static void get_terminal_wayland(char *terminal) {
  struct wl_display *display = wl_display_connect(NULL);
  if (!display) {
    return;
  }

  printf("%s\n", terminal);

  wl_display_disconnect(display);
}

static void get_terminal_x11(char *terminal) {
  Display *display = XOpenDisplay(NULL);
  if (!display) {
    return;
  }

  Atom a, active, class;
  unsigned long _, window;
  unsigned char *prop;

  window = RootWindow(display, XDefaultScreen(display));
  active = XInternAtom(display, "_NET_ACTIVE_WINDOW", True);
  class = XInternAtom(display, "WM_CLASS", True);

  if (XGetWindowProperty(display, window, active, 0, 64, 0, 0, &a, (int *)&_, &_, &_, &prop) == Success) {
    window = (prop[3] << 24) + (prop[2] << 16) + (prop[1] << 8) + prop[0];
    free(prop);

    if (window && XGetWindowProperty(display, window, class, 0, 64, 0, 0, &a, (int *)&_, &_, &_, &prop) == Success) {
      snprintf(terminal, BUF_SIZE, "%s", prop);
      free(prop);
    }
  }

  XCloseDisplay(display);
}

static char *get_terminal() {
  //unsigned char *prop;
  char *terminal = calloc(1, BUF_SIZE);
  if (!terminal) {
    return "不明";
  }
  char *wayland_display = getenv("WAYLAND_DISPLAY");

  if (wayland_display != NULL) { // TODO
    get_terminal_wayland(terminal);
  } else {
    get_terminal_x11(terminal);

    if (!(*terminal)) {
      strncpy(terminal, getenv("TERM"), BUF_SIZE - 1);
      if (strcmp(terminal, "linux") == 0) {
        strncpy(terminal, ttyname(STDIN_FILENO), BUF_SIZE -1);
      }
    }
  }

  return terminal;
}

static char *get_cpu() {
  FILE *cpuinfo = fopen("/proc/cpuinfo", "r");
  if (cpuinfo == NULL) {
    status = -1;
    halt_and_catch_fire("Unable to open cpuinfo");
  }

  char *cpu_model = malloc(BUF_SIZE / 2);
  char *line = malloc(BUF_SIZE);
  size_t len = 0;
  int num_cores = 0, cpu_freq;
  double freq;
  char freq_unit[] = "GHz";
  int prec = 3;

  while (getline(&line, &len, cpuinfo) != -1) {
    num_cores += sscanf(line, "model name : %[^\n@]", cpu_model);
  }

  fseek(cpuinfo, 0, SEEK_SET);

  while (getline(&line, &len, cpuinfo) != -1) {
    if (sscanf(line, "cpu MHz : %lf", &freq) > 0) {
      break;
    }
  }

  fclose(cpuinfo);
  free(line);

  cpu_freq = (int) freq;

  if (cpu_freq < 1000) {
    freq = (double) cpu_freq;
    freq_unit[0] = 'M';
    prec = 0;
  } else {
    freq = cpu_freq / 1000.0;

    while (cpu_freq % 10 == 0) {
      --prec;
      cpu_freq /= 10;
    }

    if (prec == 0) {
      prec = 1;
    }
  }

  for (int i = 0; i < COUNT(cpu_config); ++i) {
    if (cpu_config[i].repl_str == NULL) {
      remove_substring(cpu_model, cpu_config[i].substring, cpu_config[i].length);
    } else {
      replace_substring(cpu_model, cpu_config[i].substring, cpu_config[i].repl_str, cpu_config[i].length, cpu_config[i].repl_len);
    }
  }

  char *cpu = malloc(BUF_SIZE);
  snprintf(cpu, BUF_SIZE, "%s (%d) @ %.*f%s", cpu_model, num_cores, prec, freq, freq_unit);
  free(cpu_model);

  truncate_spaces(cpu);

  if (num_cores == 0) {
    *cpu = '\0';
  }

  return cpu;
}

static char *find_gpu(int index) {
  char buffer[BUF_SIZE], *device_class;
  struct pci_access *pacc;
  struct pci_dev *dev;
  int gpu_index = 0;
  bool found = false;
  char *gpu = malloc(BUF_SIZE);

  pacc = pci_alloc();
  pci_init(pacc);
  pci_scan_bus(pacc);

  for (dev = pacc->devices; dev != NULL; dev = dev->next) {
    pci_fill_info(dev, PCI_FILL_IDENT | PCI_FILL_BASES | PCI_FILL_CLASS);
    device_class = pci_lookup_name(pacc, buffer, sizeof(buffer), PCI_LOOKUP_CLASS, dev->device_class);

    if (!strcmp("VGA compatible controller", device_class) || !strcmp("3D controller", device_class)) {
      snprintf(gpu, BUF_SIZE, "%s", pci_lookup_name(pacc, buffer, sizeof(buffer), PCI_LOOKUP_DEVICE | PCI_LOOKUP_VENDOR, dev->vendor_id, dev->device_id));

      if (gpu_index == index) {
        found = true;
        break;
      }
      gpu_index++;
    }
  }

  pci_cleanup(pacc);

  if (!found) {
    *gpu = '\0';
  } else {
    for (int i = 0; i < COUNT(gpu_config); ++i) {
      if (gpu_config[i].repl_str == NULL) {
        remove_substring(gpu, gpu_config[i].substring, gpu_config[i].length);
      } else {
        replace_substring(gpu, gpu_config[i].substring, gpu_config[i].repl_str, gpu_config[i].length, gpu_config[i].repl_len);
      }
    }
    truncate_spaces(gpu);
  }

  return gpu;
}

static char *get_gpu1() {
  return find_gpu(0);
}

static char *get_memory() {
  int total_memory, used_memory;
  int total, shared, memfree, buffers, cached, reclaimable;

  FILE *meminfo = fopen("/proc/meminfo", "r");
  if (meminfo == NULL) {
    status = -1;
    halt_and_catch_fire("Unable to open meminfo");
  }

  char *line = NULL;
  size_t len;

  while (getline(&line, &len, meminfo) != -1) {
    sscanf(line, "MemTotal: %d", &total);
    sscanf(line, "Shmem: %d", &shared);
    sscanf(line, "MemFree: %d", &memfree);
    sscanf(line, "Buffers: %d", &buffers);
    sscanf(line, "Cached: %d", &cached);
    sscanf(line, "SReclaimable: %d", &reclaimable);
  }

  free(line);
  fclose(meminfo);

  used_memory = (total + shared - memfree - buffers - cached - reclaimable) / 1024;
  total_memory = total / 1024;
  int percentage = (int) (100 * (used_memory / (double) total_memory));

  char *memory = malloc(BUF_SIZE);
  snprintf(memory, BUF_SIZE, "%dMiB / %dMiB (%d%%)", used_memory, total_memory, percentage);

  return memory;
}

static char *get_colors1() {
  char *colors1 = malloc(BUF_SIZE);
  char *s = colors1;

  for (int i = 0; i < 8; i++) {
    sprintf(s, "\e[4%dm   ", i);
    s += 8;
  }
  snprintf(s, 5, "\e[0m");

  return colors1;
}

static char *get_colors2() {
  char *colors2 = malloc(BUF_SIZE);
  char *s = colors2;

  for (int i = 8; i < 16; i++) {
    sprintf(s, "\e[48;5;%dm   ", i);
    s += 12 + (i >= 10 ? 1 : 0);
  }
  snprintf(s, 5, "\e[0m");

  return colors2;
}

static char *spacer() {
  return calloc(1, 1);
}

char *get_cache_file() {
  char *cache_file = malloc(BUF_SIZE);
  char *env = getenv("XDG_CACHE_HOME");
  if (env == NULL) {
    snprintf(cache_file, BUF_SIZE, "%s/.cache/paleofetch", getenv("HOME"));
  } else {
    snprintf(cache_file, BUF_SIZE, "%s/paleofetch", env);
  }

  return cache_file;
}

char *search_cache(char *cache_data, char *label) {
  char *start = strstr(cache_data, label);
  if (start == NULL) {
    status = ENODATA;
    halt_and_catch_fire("cache miss on key '%s'; need to --recache?", label);
  }
  start += strlen(label);
  char *end = strchr(start, ';');
  char *buf = calloc(1, BUF_SIZE);
  strncpy(buf, start + 1, end - start - 1);

  return buf;
}

char *get_value(struct conf c, int read_cache, char *cache_data) {
  char *value;

  if (c.cached && read_cache) {
    value = search_cache(cache_data, c.label);
  } else {
    value = c.function();
    if (c.cached) {
      char *buf = malloc(BUF_SIZE);
      sprintf(buf, "%s=%s;", c.label, value);
      strcat(cache_data, buf);
      free(buf);
    }
  }

  return value;
}

int main(int argc, char *argv[]) {
  char *cache, *cache_data = NULL;
  FILE *cache_file;
  int read_cache;
  char *wayland_display = getenv("WAYLAND_DISPLAY");

  if (wayland_display != NULL) { // TODO
  } else {
    Display *display;
    status = uname(&uname_info);
    halt_and_catch_fire("uname failed");
    status = sysinfo(&my_sysinfo);
    halt_and_catch_fire("sysinfo failed");
    display = XOpenDisplay(NULL);

    cache = get_cache_file();
    if (argc == 2 && strcmp(argv[1], "--recache") == 0) {
      read_cache = 0;
    } else {
      cache_file = fopen(cache, "r");
      read_cache = cache_file != NULL;
    }

    if (!read_cache) {
      cache_data = calloc(4, BUF_SIZE);
    } else {
      size_t len;
      getline(&cache_data, &len, cache_file);
      fclose(cache_file);
    }

    int offset = 0;

    for (int i = 0; i < COUNT(LOGO); i++) {
      if (i >= COUNT(config) - offset) {
        printf("%s\n", LOGO[i]);
      } else {
        char *label = config[i+offset].label, *value = get_value(config[i+offset], read_cache, cache_data);
        if (strcmp(value, "") != 0) {
          printf("%s"COLOR"%s\e[0m%s\n", LOGO[i], label, value);
        } else {
          if (strcmp(label, "") != 0) {
            ++offset;
            free(value);
            label = config[i+offset].label;
            value = get_value(config[i+offset], read_cache, cache_data);
          }
          printf("%s"COLOR"%s\e[0m%s\n", LOGO[i], label, value);
        }
        free(value);

      }
    }
    puts("\e[0m");

    if (!read_cache && *cache_data) {
      cache_file = fopen(cache, "w");
      fprintf(cache_file, "%s", cache_data);
      fclose(cache_file);
    }

    free(cache);
    free(cache_data);
    if (display != NULL) {
      XCloseDisplay(display);
    }
  }

  return 0;
}