#include <X11/Xlib.h>
#include <X11/Xutil.h>
#include <X11/keysym.h>
#include <errno.h>
#include <fcntl.h>
#include <pty.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <termios.h>
#include <unistd.h>

#include "config.h"

// There isn't a need for MAX_ROWS or MAX_COLS, but ¯\_(ツ)_/¯
#define MAX_ROWS 1000
#define MAX_COLS 1000
#define CHAR_WIDTH 8
#define CHAR_HEIGHT 16

static pid_t child_pid;
static char **terminal_buffer = NULL;
static unsigned long **color_buffer = NULL;
static int **attr_buffer = NULL;
static int term_rows;
static int term_cols;
static int cursor_x = 0;
static int cursor_y = 0;
static unsigned long current_color;
static int current_attr = 0;
static GC gc = NULL;
static GC gc_bold = NULL;
static Display *display = NULL;
static int cursor_visible = 1;
static struct timeval last_blink;
static struct timeval last_refresh;
static Pixmap buffer_pixmap = None;
static int master_fd = -1;
static int needs_refresh = 0;
static XFontStruct *regular_font = NULL;
static XFontStruct *bold_font = NULL;
static XFontStruct *italic_font = NULL;
static int window_focused = 0;

#define ATTR_BOLD 1
#define ATTR_ITALIC 2

// Pesky static smth smth error, so we "declare" it early
static void draw_terminal(Display *display, Window window);
static void draw_char(Display *display, Drawable d, int x, int y);

// Cursor gameplay <== ==>
static void draw_cursor(Display *display, Window window) {
  if (!gc || !terminal_buffer || !color_buffer || cursor_y >= term_rows ||
      cursor_x >= term_cols)
    return;

  // First restore the character that was under the cursor
  if (cursor_x < term_cols && cursor_y < term_rows) {
    draw_char(display, buffer_pixmap, cursor_x, cursor_y);
  }

  if (cursor_visible) {
    // Save original color
    unsigned long original_color = color_buffer[cursor_y][cursor_x];

    if (!window_focused) {
      XSetForeground(display, gc, original_color);
      XDrawRectangle(display, buffer_pixmap, gc, cursor_x * CHAR_WIDTH,
                     cursor_y * CHAR_HEIGHT, CHAR_WIDTH - 1, CHAR_HEIGHT - 1);
    } else {
      XSetForeground(display, gc, original_color);
      XFillRectangle(display, buffer_pixmap, gc, cursor_x * CHAR_WIDTH,
                     cursor_y * CHAR_HEIGHT, CHAR_WIDTH, CHAR_HEIGHT);

      if (cursor_x < term_cols && cursor_y < term_rows) {
        char str[2] = {terminal_buffer[cursor_y][cursor_x], '\0'};
        XSetForeground(display, gc,
                       XBlackPixel(display, DefaultScreen(display)));

        XFontStruct *font = regular_font;
        if (attr_buffer[cursor_y][cursor_x] & ATTR_BOLD) {
          font = bold_font;
        } else if (attr_buffer[cursor_y][cursor_x] & ATTR_ITALIC) {
          font = italic_font;
        }
        XSetFont(display, gc, font->fid);

        XDrawString(display, buffer_pixmap, gc, cursor_x * CHAR_WIDTH,
                    (cursor_y + 1) * CHAR_HEIGHT - 2, str, 1);
      }
    }
  }

  // Copy the relevant area to the window
  XCopyArea(display, buffer_pixmap, window, gc, cursor_x * CHAR_WIDTH,
            cursor_y * CHAR_HEIGHT, CHAR_WIDTH, CHAR_HEIGHT,
            cursor_x * CHAR_WIDTH, cursor_y * CHAR_HEIGHT);
}

static void update_cursor_state(Display *display, Window window) {
  if (window_focused) {
    cursor_visible = 1;
    return;
  }
}

// resize_buffers visual:
//
// Old buffer (3x3):     New buffer (4x4):
// [a][b][c]            [a][b][c][ ]
// [d][e][f]     ->     [d][e][f][ ]
// [g][h][i]            [g][h][i][ ]
//                      [ ][ ][ ][ ]
//
// 1. Allocates new larger buffer
// 2. Copies existing content
// 3. Fills new space with blanks
// 4. Updates cursor position
// 5. Frees old buffer
static void resize_buffers(int new_rows, int new_cols) {
  if (new_rows > MAX_ROWS)
    new_rows = MAX_ROWS;
  if (new_cols > MAX_COLS)
    new_cols = MAX_COLS;
  if (new_rows < 1)
    new_rows = 1;
  if (new_cols < 1)
    new_cols = 1;

  size_t row_size = ((new_cols * sizeof(char) + 15) & ~15);
  size_t color_row_size = ((new_cols * sizeof(unsigned long) + 15) & ~15);
  size_t attr_row_size = ((new_cols * sizeof(int) + 15) & ~15);

  char **new_term_buffer = malloc(new_rows * sizeof(char *));
  unsigned long **new_color_buffer = malloc(new_rows * sizeof(unsigned long *));
  int **new_attr_buffer = malloc(new_rows * sizeof(int *));

  if (!new_term_buffer || !new_color_buffer || !new_attr_buffer) {
    fprintf(stderr, "Failed to allocate memory for buffers\n");
    return;
  }

  for (int i = 0; i < new_rows; i++) {
    new_term_buffer[i] = malloc(row_size);
    new_color_buffer[i] = malloc(color_row_size);
    new_attr_buffer[i] = malloc(attr_row_size);

    if (!new_term_buffer[i] || !new_color_buffer[i] || !new_attr_buffer[i]) {
      fprintf(stderr, "Failed to allocate memory for buffer row\n");
      // Clean up allocated memory
      for (int j = 0; j < i; j++) {
        free(new_term_buffer[j]);
        free(new_color_buffer[j]);
        free(new_attr_buffer[j]);
      }
      free(new_term_buffer);
      free(new_color_buffer);
      free(new_attr_buffer);
      return;
    }

    if (i < term_rows && i < new_rows && terminal_buffer && color_buffer &&
        attr_buffer) {
      int copy_cols = (new_cols < term_cols) ? new_cols : term_cols;
      memcpy(new_term_buffer[i], terminal_buffer[i], copy_cols * sizeof(char));
      memcpy(new_color_buffer[i], color_buffer[i],
             copy_cols * sizeof(unsigned long));
      memcpy(new_attr_buffer[i], attr_buffer[i], copy_cols * sizeof(int));
      if (new_cols > term_cols) {
        memset(new_term_buffer[i] + term_cols, ' ', new_cols - term_cols);
        for (int x = term_cols; x < new_cols; x++) {
          new_color_buffer[i][x] = current_color;
          new_attr_buffer[i][x] = current_attr;
        }
      }
    } else {
      memset(new_term_buffer[i], ' ', new_cols);
      for (int x = 0; x < new_cols; x++) {
        new_color_buffer[i][x] = current_color;
        new_attr_buffer[i][x] = current_attr;
      }
    }
  }

  // Free old buffers after successful allocation
  if (terminal_buffer && color_buffer && attr_buffer) {
    for (int i = 0; i < term_rows; i++) {
      free(terminal_buffer[i]);
      free(color_buffer[i]);
      free(attr_buffer[i]);
    }
    free(terminal_buffer);
    free(color_buffer);
    free(attr_buffer);
  }

  terminal_buffer = new_term_buffer;
  color_buffer = new_color_buffer;
  attr_buffer = new_attr_buffer;
  term_rows = new_rows;
  term_cols = new_cols;

  cursor_x = (cursor_x >= term_cols) ? term_cols - 1 : cursor_x;
  cursor_y = (cursor_y >= term_rows) ? term_rows - 1 : cursor_y;

  // Resize buffer pixmap
  if (buffer_pixmap != None) {
    XFreePixmap(display, buffer_pixmap);
  }
  buffer_pixmap = XCreatePixmap(display, DefaultRootWindow(display),
                                term_cols * CHAR_WIDTH, term_rows * CHAR_HEIGHT,
                                DefaultDepth(display, DefaultScreen(display)));
  needs_refresh = 1;
}

static void destroy_cb(Display *display, Window window) {
  if (child_pid > 0) {
    kill(child_pid, SIGTERM);
    waitpid(child_pid, NULL, 0);
  }
  if (gc)
    XFreeGC(display, gc);
  if (gc_bold)
    XFreeGC(display, gc_bold);
  if (buffer_pixmap != None)
    XFreePixmap(display, buffer_pixmap);
  if (display)
    XCloseDisplay(display);
  if (terminal_buffer) {
    for (int i = 0; i < term_rows; i++) {
      free(terminal_buffer[i]);
    }
    free(terminal_buffer);
  }
  if (color_buffer) {
    for (int i = 0; i < term_rows; i++) {
      free(color_buffer[i]);
    }
    free(color_buffer);
  }
  if (attr_buffer) {
    for (int i = 0; i < term_rows; i++) {
      free(attr_buffer[i]);
    }
    free(attr_buffer);
  }
  if (master_fd >= 0)
    close(master_fd);
  exit(0);
}

static void draw_char(Display *display, Drawable d, int x, int y) {
  if (!gc || !terminal_buffer || !color_buffer || !attr_buffer)
    return;
  if (x < 0 || x >= term_cols || y < 0 || y >= term_rows)
    return;

  // Clear the character position first
  XSetForeground(display, gc, XBlackPixel(display, DefaultScreen(display)));
  XFillRectangle(display, d, gc, x * CHAR_WIDTH, y * CHAR_HEIGHT, CHAR_WIDTH,
                 CHAR_HEIGHT);

  char str[2] = {terminal_buffer[y][x], '\0'};
  XSetForeground(display, gc, color_buffer[y][x]);

  // Select appropriate font based on attributes
  XFontStruct *font = regular_font;
  if (attr_buffer[y][x] & ATTR_BOLD) {
    font = bold_font;
  } else if (attr_buffer[y][x] & ATTR_ITALIC) {
    font = italic_font;
  }
  XSetFont(display, gc, font->fid);

  XDrawString(display, d, gc, x * CHAR_WIDTH, (y + 1) * CHAR_HEIGHT - 2, str,
              1);
}

static void draw_terminal(Display *display, Window window) {
  if (!gc) {
    gc = XCreateGC(display, window, 0, NULL);

    regular_font = XLoadQueryFont(display, REGULAR_FONT);
    bold_font = XLoadQueryFont(display, BOLD_FONT);
    italic_font = XLoadQueryFont(display, ITALICS_FONT);

    if (!regular_font)
      regular_font = XLoadQueryFont(display, "fixed");
    if (!bold_font)
      bold_font = regular_font;
    if (!italic_font)
      italic_font = regular_font;

    XSetFont(display, gc, regular_font->fid);
  }

  if (!terminal_buffer || !color_buffer || !attr_buffer)
    return;

  if (buffer_pixmap == None) {
    buffer_pixmap = XCreatePixmap(
        display, DefaultRootWindow(display), term_cols * CHAR_WIDTH,
        term_rows * CHAR_HEIGHT, DefaultDepth(display, DefaultScreen(display)));
  }

  // Draw to buffer first
  for (int y = 0; y < term_rows; y++) {
    for (int x = 0; x < term_cols; x++) {
      draw_char(display, buffer_pixmap, x, y);
    }
  }

  // Copy buffer to window
  XCopyArea(display, buffer_pixmap, window, gc, 0, 0, term_cols * CHAR_WIDTH,
            term_rows * CHAR_HEIGHT, 0, 0);

  draw_cursor(display, window);
  XFlush(display);
  needs_refresh = 0;
}

static void check_refresh(Display *display, Window window) {
  if (!needs_refresh)
    return;

  struct timeval now;
  gettimeofday(&now, NULL);

  long elapsed = (now.tv_sec - last_refresh.tv_sec) * 1000000 +
                 (now.tv_usec - last_refresh.tv_usec);

  if (elapsed >= REFRESH_INTERVAL) {
    draw_terminal(display, window);
    last_refresh = now;
  }
}

static void clear_terminal(Display *display, Window window) {
  if (!terminal_buffer || !color_buffer || !attr_buffer)
    return;
  for (int y = 0; y < term_rows; y++) {
    memset(terminal_buffer[y], ' ', term_cols);
    for (int x = 0; x < term_cols; x++) {
      color_buffer[y][x] = current_color;
      attr_buffer[y][x] = current_attr;
    }
  }
  cursor_x = 0;
  cursor_y = 0;
  needs_refresh = 1;
}

// Scroll the terminal up one line
static void scroll_up(void) {
  // Move all lines up one position
  for (int y = 0; y < term_rows - 1; y++) {
    memcpy(terminal_buffer[y], terminal_buffer[y + 1], term_cols);
    memcpy(color_buffer[y], color_buffer[y + 1],
           term_cols * sizeof(unsigned long));
    memcpy(attr_buffer[y], attr_buffer[y + 1], term_cols * sizeof(int));
  }

  // Clear the bottom line
  memset(terminal_buffer[term_rows - 1], ' ', term_cols);
  for (int x = 0; x < term_cols; x++) {
    color_buffer[term_rows - 1][x] = current_color;
    attr_buffer[term_rows - 1][x] = current_attr;
  }
  needs_refresh = 1;
}

// Handles those pesky ANSI color escape codes that terminals use.
// When \033[<number>m is sent to the terminal, it changes text color:
//   30 = Black    31 = Red      32 = Green   33 = Yellow
//   34 = Blue     35 = Magenta  36 = Cyan    37 = White
// 0 sets everything back to plain white
// P.S. yes I had to google this up
static void parse_ansi_code(const char *buf, int *idx, int max_len) {
  char num_buf[16] = {0};
  int num_idx = 0;
  (*idx)++; // Skip [

  while (*idx < max_len && buf[*idx] != 'm') {
    if (buf[*idx] >= '0' && buf[*idx] <= '9') {
      num_buf[num_idx++] = buf[*idx];
    }
    (*idx)++;
  }

  if (num_buf[0] != '\0') {
    int code = atoi(num_buf);
    switch (code) {
    case 0:
      current_color = XWhitePixel(display, DefaultScreen(display));
      current_attr = 0;
      break;
    case 1:
      current_attr |= ATTR_BOLD;
      break;
    case 3:
      current_attr |= ATTR_ITALIC;
      break;
    case 22:
      current_attr &= ~ATTR_BOLD;
      break;
    case 23:
      current_attr &= ~ATTR_ITALIC;
      break;
    case 30:
      current_color = XBlackPixel(display, DefaultScreen(display));
      break;
    case 31:
      current_color = RED; // Red
      break;
    case 32:
      current_color = GREEN; // Green
      break;
    case 33:
      current_color = YELLOW; // Yellow
      break;
    case 34:
      current_color = BLUE; // Blue
      break;
    case 35:
      current_color = MAGENTA; // Magenta
      break;
    case 36:
      current_color = CYAN; // Cyan
      break;
    case 37:
      current_color = XWhitePixel(display, DefaultScreen(display));
      break;
    case 90:
      current_color = DARK_GRAY; // Dark gray
      break;
    case 91:
      current_color = LIGHT_RED; // Light red
      break;
    case 92:
      current_color = LIGHT_GREEN; // Light green
      break;
    case 93:
      current_color = LIGHT_YELLOW; // Light yellow
      break;
    case 94:
      current_color = LIGHT_BLUE; // Light blue
      break;
    case 95:
      current_color = LIGHT_MAGENTA; // Light magenta
      break;
    case 96:
      current_color = LIGHT_CYAN; // Light cyan
      break;
    case 97:
      current_color = 0xFFFFFF; // Bright white teehee
      break;
    }
  }
}

static int master_cb(int fd, void *data) {
  if (!terminal_buffer || !color_buffer || !attr_buffer)
    return 0;
  char buf[4096];
  ssize_t bytes_read;
  struct {
    Display *display;
    Window window;
  } *ctx = data;

  bytes_read = read(fd, buf, sizeof(buf));
  if (bytes_read <= 0 && errno != EAGAIN && errno != EINTR) {
    destroy_cb(ctx->display, ctx->window);
    return 0;
  }
  if (bytes_read > 0) {
    for (int i = 0; i < bytes_read; i++) {
      if (buf[i] == '\033' && i + 1 < bytes_read && buf[i + 1] == '[') {
        parse_ansi_code(buf, &i, bytes_read);
        continue;
      }
      if (buf[i] == '\n') {
        cursor_x = 0;
        cursor_y++;
        if (cursor_y >= term_rows) {
          scroll_up();
          cursor_y = term_rows - 1;
          needs_refresh = 1;
        }
      } else if (buf[i] == '\r') {
        cursor_x = 0;
      } else if (buf[i] == '\b') {
        if (cursor_x > 0) {
          cursor_x--;
          terminal_buffer[cursor_y][cursor_x] = ' ';
          needs_refresh = 1;
        }
      } else {
        if (cursor_y < term_rows && cursor_x < term_cols) {
          terminal_buffer[cursor_y][cursor_x] = buf[i];
          color_buffer[cursor_y][cursor_x] = current_color;
          attr_buffer[cursor_y][cursor_x] = current_attr;
          needs_refresh = 1;
          cursor_x++;
          if (cursor_x >= term_cols) {
            cursor_x = 0;
            cursor_y++;
            if (cursor_y >= term_rows) {
              scroll_up();
              cursor_y = term_rows - 1;
              needs_refresh = 1;
            }
          }
        }
      }
    }
    cursor_visible = 1;
    if (needs_refresh) {
      check_refresh(ctx->display, ctx->window);
    } else {
      draw_cursor(ctx->display, ctx->window);
      XFlush(ctx->display);
    }
  }
  return 1;
}

static void key_press_cb(XKeyEvent *event, void *data) {
  int *master = (int *)data;
  char buf[32];
  KeySym keysym;
  int len;

  len = XLookupString(event, buf, sizeof(buf), &keysym, NULL);

  if (event->state & ControlMask) {
    if ((keysym & 0x1f) == ('L' & 0x1f)) {
      write(*master, "\f", 1);
      clear_terminal(event->display, event->window);
      return;
    }
    if ((keysym & 0x1f) == ('C' & 0x1f)) {
      kill(child_pid, SIGINT);
      return;
    }
    if ((keysym & 0x1f) == ('U' & 0x1f)) {
      write(*master, "\025", 1);
      return;
    }
    char ctrl_char = '@' + (keysym & 0x1f);
    write(*master, &ctrl_char, 1);
    return;
  }

  if (keysym == XK_Return || keysym == XK_KP_Enter) {
    write(*master, "\r", 1);
  } else if (keysym == XK_BackSpace) {
    write(*master, "\b", 1);
  } else if (len > 0) {
    write(*master, buf, len);
  }
}

int main(void) {
  int master, slave;
  pid_t pid;

  term_rows = 24;
  term_cols = 80;

  gettimeofday(&last_blink, NULL);
  gettimeofday(&last_refresh, NULL);

  struct winsize ws = {.ws_row = term_rows,
                       .ws_col = term_cols,
                       .ws_xpixel = term_cols * CHAR_WIDTH,
                       .ws_ypixel = term_rows * CHAR_HEIGHT};

  terminal_buffer = malloc(term_rows * sizeof(char *));
  color_buffer = malloc(term_rows * sizeof(unsigned long *));
  attr_buffer = malloc(term_rows * sizeof(int *));

  if (!terminal_buffer || !color_buffer || !attr_buffer) {
    fprintf(stderr, "Failed to allocate memory for buffers\n");
    return 1;
  }

  for (int i = 0; i < term_rows; i++) {
    terminal_buffer[i] = malloc(term_cols * sizeof(char));
    color_buffer[i] = malloc(term_cols * sizeof(unsigned long));
    attr_buffer[i] = malloc(term_cols * sizeof(int));

    if (!terminal_buffer[i] || !color_buffer[i] || !attr_buffer[i]) {
      fprintf(stderr, "Failed to allocate memory for buffer row\n");
      return 1;
    }
    memset(terminal_buffer[i], ' ', term_cols);
    memset(attr_buffer[i], 0, term_cols * sizeof(int));
  }

  display = XOpenDisplay(NULL);
  if (!display) {
    fprintf(stderr, "Cannot open display\n");
    return 1;
  }
  XInitThreads();

  current_color = XWhitePixel(display, DefaultScreen(display));
  current_attr = 0;

  for (int y = 0; y < term_rows; y++) {
    for (int x = 0; x < term_cols; x++) {
      color_buffer[y][x] = current_color;
      attr_buffer[y][x] = current_attr;
    }
  }

  XSizeHints hints;
  hints.flags = PResizeInc | PMinSize;
  hints.width_inc = CHAR_WIDTH;
  hints.height_inc = CHAR_HEIGHT;
  hints.min_width = CHAR_WIDTH * 4;
  hints.min_height = CHAR_HEIGHT * 1;

  Window window = XCreateSimpleWindow(
      display, DefaultRootWindow(display), 0, 0, term_cols * CHAR_WIDTH,
      term_rows * CHAR_HEIGHT, 0, XBlackPixel(display, DefaultScreen(display)),
      XBlackPixel(display, DefaultScreen(display)));

  buffer_pixmap = XCreatePixmap(display, DefaultRootWindow(display),
                                term_cols * CHAR_WIDTH, term_rows * CHAR_HEIGHT,
                                DefaultDepth(display, DefaultScreen(display)));

  XSetWMNormalHints(display, window, &hints);
  XStoreName(display, window, window_name);
  XSelectInput(display, window,
               KeyPressMask | ExposureMask | StructureNotifyMask |
                   FocusChangeMask);
  XMapWindow(display, window);

  if (openpty(&master, &slave, NULL, NULL, &ws) == -1) {
    perror("openpty");
    return 1;
  }

  master_fd = master;

  pid = fork();
  if (pid < 0) {
    perror("fork");
    return 1;
  }

  if (pid == 0) {
    close(master);
    close(STDIN_FILENO);
    close(STDOUT_FILENO);
    close(STDERR_FILENO);

    setsid();
    dup2(slave, 0);
    dup2(slave, 1);
    dup2(slave, 2);

    if (slave > 2)
      close(slave);

    char *shell = getenv("SHELL");
    if (!shell)
      shell = "/bin/sh";

    setenv("TERM", "xterm-256color", 1);
    execlp(shell, shell, NULL);
    perror("execlp");
    exit(1);
  }

  child_pid = pid;
  close(slave);

  struct {
    Display *display;
    Window window;
  } ctx = {display, window};

  XEvent event;
  while (1) {
    int status;
    pid_t wpid = waitpid(child_pid, &status, WNOHANG);
    if (wpid == child_pid) {
      destroy_cb(display, window);
      return 0;
    }

    while (XPending(display)) {
      XNextEvent(display, &event);
      switch (event.type) {
      case ConfigureNotify: {
        XConfigureEvent *ce = (XConfigureEvent *)&event;
        int new_rows = ce->height / CHAR_HEIGHT;
        int new_cols = ce->width / CHAR_WIDTH;
        if (new_rows != term_rows || new_cols != term_cols) {
          resize_buffers(new_rows, new_cols);
          ws.ws_row = term_rows;
          ws.ws_col = term_cols;
          ws.ws_xpixel = term_cols * CHAR_WIDTH;
          ws.ws_ypixel = term_rows * CHAR_HEIGHT;
          ioctl(master, TIOCSWINSZ, &ws);
        }
        break;
      }
      case FocusIn:
        window_focused = 1;
        needs_refresh = 1;
        break;
      case FocusOut:
        window_focused = 0;
        needs_refresh = 1;
        break;
      case KeyPress:
        key_press_cb((XKeyEvent *)&event, &master);
        break;
      case Expose:
        needs_refresh = 1;
        check_refresh(display, window);
        break;
      }
    }

    fd_set fds;
    struct timeval tv = {0, 1000}; // 1ms timeout for faster refresh
    FD_ZERO(&fds);
    FD_SET(master, &fds);
    if (select(master + 1, &fds, NULL, NULL, &tv) > 0) {
      master_cb(master, &ctx);
    }

    check_refresh(display, window);
    update_cursor_state(display, window);
  }

  return 0;
}