#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <signal.h>
#include <sys/stat.h>

#ifdef _WIN32
#include <conio.h>
#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <io.h>
#include <direct.h>
#define STDIN_FILENO 0
#define mkdir_p(path) _mkdir(path)
#else
#include <sys/ioctl.h>
#include <termios.h>
#include <unistd.h>
#define mkdir_p(path) mkdir(path, 0755)
#endif

#include "ant.h"
#include "repl.h"
#include "reactor.h"
#include "runtime.h"
#include "internal.h"

#include "silver/ast.h"
#include "silver/engine.h"

#include <crprintf.h>
#include "modules/io.h"
#include "highlight.h"
#include "highlight/regex.h"

#define MAX_HISTORY          512
#define MAX_LINE_LENGTH      4096
#define MAX_MULTILINE_LENGTH 65536

#define INPUT \
  char *line, int *pos, int *len, key_event_t *key, history_t *hist, const char *prompt

static volatile sig_atomic_t ctrl_c_pressed = 0;

typedef struct {
  char **lines;
  int count;
  int capacity;
  int current;
} history_t;

typedef enum {
  CMD_OK,
  CMD_EXIT,
  CMD_NOT_FOUND
} cmd_result_t;

typedef struct {
  const char *name;
  const char *description;
  bool has_arg;
  cmd_result_t (*handler)(ant_t *js, history_t *history, const char *arg);
} repl_command_t;

typedef struct {
  char *name;
  size_t len;
} repl_decl_name_t;

typedef struct {
  repl_decl_name_t *items;
  size_t count;
  size_t cap;
} repl_decl_registry_t;

typedef struct {
  const char **names;
  uint32_t *lens;
  size_t count;
  size_t cap;
} repl_decl_pending_t;

static repl_decl_registry_t *g_repl_decl_registry = NULL;
static void sigint_handler(int sig) { ctrl_c_pressed++; }

static inline void repl_clear_exception_state(ant_t *js) {
  js->thrown_exists = false;
  js->thrown_value = js_mkundef();
}

static void repl_decl_registry_free(repl_decl_registry_t *reg) {
  if (!reg) return;
  for (size_t i = 0; i < reg->count; i++) 
    free(reg->items[i].name);
  free(reg->items);
  reg->items = NULL;
  reg->count = 0;
  reg->cap = 0;
}

static bool repl_decl_registry_contains(
  const repl_decl_registry_t *reg,
  const char *name, uint32_t len
) {
  if (!reg || !name) return false;
  for (size_t i = 0; i < reg->count; i++) {
    if (
      reg->items[i].len == (size_t)len 
      && memcmp(reg->items[i].name, name, (size_t)len) == 0
    ) return true;
  }
  return false;
}

static bool repl_decl_registry_add(
  ant_t *js, repl_decl_registry_t *reg,
  const char *name, uint32_t len
) {
  if (!reg || !name) return true;
  if (repl_decl_registry_contains(reg, name, len)) return true;
  
  if (reg->count >= reg->cap) {
    size_t new_cap = reg->cap ? reg->cap * 2 : 32;
    repl_decl_name_t *ni = realloc(reg->items, new_cap * sizeof(*ni));
    if (!ni) {
      js_mkerr_typed(js, JS_ERR_INTERNAL | JS_ERR_NO_STACK, "out of memory");
      return false;
    }
    reg->items = ni;
    reg->cap = new_cap;
  }
  
  char *copy = malloc((size_t)len + 1);
  if (!copy) {
    js_mkerr_typed(js, JS_ERR_INTERNAL | JS_ERR_NO_STACK, "out of memory");
    return false;
  }
  
  memcpy(copy, name, (size_t)len);
  copy[len] = '\0';
  reg->items[reg->count++] = (repl_decl_name_t){ .name = copy, .len = (size_t)len };
  
  return true;
}

static void repl_decl_pending_free(repl_decl_pending_t *p) {
  if (!p) return;
  free(p->names);
  free(p->lens);
  p->names = NULL;
  p->lens = NULL;
  p->count = 0;
  p->cap = 0;
}

static bool repl_decl_pending_contains(
  const repl_decl_pending_t *p,
  const char *name, uint32_t len
) {
  if (!p || !name) return false;
  for (size_t i = 0; i < p->count; i++) 
    if (p->lens[i] == len && memcmp(p->names[i], name, (size_t)len) == 0) return true;
  return false;
}

static bool repl_decl_pending_push(
  ant_t *js, repl_decl_pending_t *p,
  const char *name, uint32_t len
) {
  if (!p || !name || len == 0) return true;
  if (repl_decl_pending_contains(p, name, len)) return true;
  if (p->count >= p->cap) {
    size_t new_cap = p->cap ? p->cap * 2 : 16;
    const char **nn = realloc(p->names, new_cap * sizeof(*nn));
    if (!nn) {
      js_mkerr_typed(js, JS_ERR_INTERNAL | JS_ERR_NO_STACK, "out of memory");
      return false;
    }
    uint32_t *nl = realloc(p->lens, new_cap * sizeof(*nl));
    if (!nl) {
      p->names = nn;
      js_mkerr_typed(js, JS_ERR_INTERNAL | JS_ERR_NO_STACK, "out of memory");
      return false;
    }
    p->names = nn;
    p->lens = nl;
    p->cap = new_cap;
  }
  p->names[p->count] = name;
  p->lens[p->count] = len;
  p->count++;
  return true;
}

static bool repl_collect_pattern_names(ant_t *js, sv_ast_t *pat, repl_decl_pending_t *p) {
  if (!pat) return true;
  switch (pat->type) {
    case N_IDENT:
      return repl_decl_pending_push(js, p, pat->str, pat->len);
    case N_ASSIGN_PAT:
    case N_ASSIGN:
      return repl_collect_pattern_names(js, pat->left, p);
    case N_REST:
    case N_SPREAD:
      return repl_collect_pattern_names(js, pat->right, p);
    case N_ARRAY:
    case N_ARRAY_PAT:
      for (int i = 0; i < pat->args.count; i++) {
        if (!repl_collect_pattern_names(js, pat->args.items[i], p)) return false;
      }
      return true;
    case N_OBJECT:
    case N_OBJECT_PAT:
      for (int i = 0; i < pat->args.count; i++) {
        sv_ast_t *prop = pat->args.items[i];
        if (!prop) continue;
        if (prop->type == N_PROPERTY) {
          if (!repl_collect_pattern_names(js, prop->right, p)) return false;
        } else if (prop->type == N_REST || prop->type == N_SPREAD) {
          if (!repl_collect_pattern_names(js, prop->right, p)) return false;
        }
      }
      return true;
    default: return true;
  }
}

static bool repl_collect_top_level_decls(ant_t *js, sv_ast_t *stmt, repl_decl_pending_t *p) {
  if (!stmt) return true;
  sv_ast_t *node = (stmt->type == N_EXPORT) ? stmt->left : stmt;
  if (!node) return true;

  if (node->type == N_VAR && node->var_kind != SV_VAR_VAR) {
    for (int i = 0; i < node->args.count; i++) {
      sv_ast_t *decl = node->args.items[i];
      if (!decl || decl->type != N_VARDECL || !decl->left) continue;
      if (!repl_collect_pattern_names(js, decl->left, p)) return false;
    }
    return true;
  }

  if (node->type == N_CLASS && node->str && node->len > 0)
    return repl_decl_pending_push(js, p, node->str, node->len);

  if (node->type == N_IMPORT_DECL) {
    for (int i = 0; i < node->args.count; i++) {
      sv_ast_t *spec = node->args.items[i];
      if (!spec || spec->type != N_IMPORT_SPEC || !spec->right) continue;
      if (spec->right->type != N_IDENT) continue;
      if (!repl_decl_pending_push(js, p, spec->right->str, spec->right->len)) return false;
    }
  }

  return true;
}

static bool repl_precheck_and_commit_lexicals(
  ant_t *js, repl_decl_registry_t *reg,
  const char *code, size_t len
) {
  if (!js || !reg || !code || len == 0) return true;

  code_arena_mark_t mark = code_arena_mark();
  repl_decl_pending_t pending = {0};
  bool ok = true;

  repl_clear_exception_state(js);
  sv_ast_t *program = sv_parse(js, code, (jsoff_t)len, false);
  
  if (!program || js->thrown_exists) {
    ok = true;
    goto done;
  }

  for (int i = 0; i < program->args.count; i++) {
    if (
      !repl_collect_top_level_decls(
      js, program->args.items[i], &pending)
    ) { ok = false; goto done; }
  }

  for (size_t i = 0; i < pending.count; i++) {
    if (repl_decl_registry_contains(reg, pending.names[i], pending.lens[i])) {
      js_mkerr_typed(
        js, JS_ERR_SYNTAX, "Identifier '%.*s' has already been declared",
        (int)pending.lens[i], pending.names[i]
      );
      ok = false; goto done;
    }
  }

  for (size_t i = 0; i < pending.count; i++) {
    if (
      !repl_decl_registry_add(
      js, reg, pending.names[i], pending.lens[i])
    ) { ok = false; goto done; }
  }

done:
  code_arena_rewind(mark);
  repl_decl_pending_free(&pending);

  if (ok && js->thrown_exists)
    repl_clear_exception_state(js);

  return ok;
}

typedef enum {
  REPL_PRINT_INTERACTIVE,
  REPL_PRINT_LOAD,
} repl_print_mode_t;

static void repl_eval_chunk(
  ant_t *js, repl_decl_registry_t *decl_registry,
  const char *code, size_t len,
  repl_print_mode_t print_mode
) {
  if (!repl_precheck_and_commit_lexicals(js, decl_registry, code, len)) {
    print_uncaught_throw(js);
    return;
  }

  repl_clear_exception_state(js);
  jsval_t result = js_eval_bytecode_repl(js, code, len);
  js_run_event_loop(js);

  if (print_uncaught_throw(js)) return;
  if (print_mode == REPL_PRINT_INTERACTIVE) {
    print_repl_value(js, result, stdout);
    return;
  }

  if (vtype(result) == T_ERR) fprintf(stderr, "%s\n", js_str(js, result));
  else if (vtype(result) != T_UNDEF) printf("%s\n", js_str(js, result));
}

static cmd_result_t cmd_help(ant_t *js, history_t *history, const char *arg);
static cmd_result_t cmd_exit(ant_t *js, history_t *history, const char *arg);
static cmd_result_t cmd_load(ant_t *js, history_t *history, const char *arg);
static cmd_result_t cmd_save(ant_t *js, history_t *history, const char *arg);
static cmd_result_t cmd_stats(ant_t *js, history_t *history, const char *arg);
static cmd_result_t cmd_copy(ant_t *js, history_t *history, const char *arg);

static const repl_command_t commands[] = {
  { "help",  "Show this help message", false, cmd_help },
  { "exit",  "Exit the REPL", false, cmd_exit },
  { "load",  "Load JS from a file into the REPL session", true, cmd_load },
  { "save",  "Save all evaluated commands in this REPL session to a file", true, cmd_save },
  { "stats", "Show memory statistics", false, cmd_stats },
  { "copy",  "Evaluate expression and copy its value", true, cmd_copy },
  { NULL, NULL, false, NULL }
};

static cmd_result_t cmd_help(ant_t *js, history_t *history, const char *arg) {
  for (const repl_command_t *cmd = commands; cmd->name; cmd++) {
    printf("  .%-7s - %s\n", cmd->name, cmd->description);
  }
  printf("\nPress Ctrl+C to abort current expression.\n");
  return CMD_OK;
}

static cmd_result_t cmd_exit(ant_t *js, history_t *history, const char *arg) {
  return CMD_EXIT;
}

static cmd_result_t cmd_load(ant_t *js, history_t *history, const char *arg) {
  (void)history;
  if (!arg || *arg == '\0') {
    fprintf(stderr, "Usage: .load <filename>\n");
    return CMD_OK;
  }
  
  FILE *fp = fopen(arg, "r");
  if (fp == NULL) {
    fprintf(stderr, "Failed to open file: %s\n", arg);
    return CMD_OK;
  }
  
  fseek(fp, 0, SEEK_END);
  long file_size = ftell(fp);
  fseek(fp, 0, SEEK_SET);
  
  char *file_buffer = malloc(file_size + 1);
  if (file_buffer) {
    size_t len = fread(file_buffer, 1, file_size, fp);
    file_buffer[len] = '\0';
    repl_eval_chunk(
      js, g_repl_decl_registry, 
      file_buffer, len, REPL_PRINT_LOAD
    );
    free(file_buffer);
  }
  fclose(fp);
  return CMD_OK;
}

static cmd_result_t cmd_save(ant_t *js, history_t *history, const char *arg) {
  (void)js;
  if (!arg || *arg == '\0') {
    fprintf(stderr, "Usage: .save <filename>\n");
    return CMD_OK;
  }
  
  FILE *fp = fopen(arg, "w");
  if (fp == NULL) {
    fprintf(stderr, "Failed to open file for writing: %s\n", arg);
    return CMD_OK;
  }
  
  for (int i = 0; i < history->count; i++) {
    fprintf(fp, "%s\n", history->lines[i]);
  }
  fclose(fp);
  printf("Session saved to %s\n", arg);
  return CMD_OK;
}

static cmd_result_t cmd_stats(ant_t *js, history_t *history, const char *arg) {
  jsval_t stats_fn = js_get(js, rt->ant_obj, "stats");
  jsval_t result = sv_vm_call(js->vm, js, stats_fn, js_mkundef(), NULL, 0, NULL, false);
  console_print(js, &result, 1, NULL, stdout);
  return CMD_OK;
}

#ifdef _WIN32
static bool repl_copy_with_command(const char *data, size_t len) {
  FILE *pipe = _popen("clip", "wb");
  if (!pipe) return false;

  size_t written = fwrite(data, 1, len, pipe);
  int close_rc = _pclose(pipe);
  return written == len && close_rc == 0;
}
#else
static bool repl_copy_with_single_command(const char *cmd, const char *data, size_t len) {
  FILE *pipe = popen(cmd, "w");
  if (!pipe) return false;

  size_t written = fwrite(data, 1, len, pipe);
  int close_rc = pclose(pipe);
  return written == len && close_rc == 0;
}

static bool repl_copy_with_command(const char *data, size_t len) {
  static const char *cmds[] = {
    "pbcopy",
    "wl-copy",
    "xclip -selection clipboard",
    "xsel --clipboard --input",
  };
  for (size_t i = 0; i < sizeof(cmds) / sizeof(cmds[0]); i++) {
    if (repl_copy_with_single_command(cmds[i], data, len)) return true;
  }
  return false;
}
#endif

static cmd_result_t cmd_copy(ant_t *js, history_t *history, const char *arg) {
  (void)history;
  if (!arg || *arg == '\0') {
    fprintf(stderr, "Usage: .copy <expression>\n");
    return CMD_OK;
  }

  repl_clear_exception_state(js);
  jsval_t result = js_eval_bytecode_repl(js, arg, strlen(arg));
  
  js_run_event_loop(js);
  if (print_uncaught_throw(js)) return CMD_OK;

  char cbuf[512];
  js_cstr_t cstr = js_to_cstr(js, result, cbuf, sizeof(cbuf));
  
  bool copied_command = repl_copy_with_command(cstr.ptr, cstr.len);
  if (cstr.needs_free) free((void *)cstr.ptr);

  if (!copied_command) {
    fprintf(stderr, "Failed to copy to clipboard (no clipboard command available).\n");
    return CMD_OK;
  }

  printf("Copied to clipboard.\n");
  return CMD_OK;
}

static cmd_result_t execute_command(ant_t *js, history_t *history, const char *line) {
  const char *cmd_start = line + 1;
  
  for (const repl_command_t *cmd = commands; cmd->name; cmd++) {
    size_t n = strlen(cmd->name);
    if (strncmp(cmd_start, cmd->name, n) != 0) continue;
    
    char next = cmd_start[n];
    if (cmd->has_arg && (next == ' ' || next == '\0')) {
      const char *arg = cmd_start + n;
      while (*arg == ' ') arg++;
      return cmd->handler(js, history, arg);
    }
    if (!cmd->has_arg && next == '\0') return cmd->handler(js, history, NULL);
  }
  
  return CMD_NOT_FOUND;
}

static void history_init(history_t *hist) {
  hist->capacity = MAX_HISTORY;
  hist->lines = malloc(sizeof(char*) * hist->capacity);
  hist->count = 0;
  hist->current = -1;
}

static void history_add(history_t *hist, const char *line) {
  if (strlen(line) == 0) return;
  if (hist->count > 0 && strcmp(hist->lines[hist->count - 1], line) == 0) return;
  
  if (hist->count >= hist->capacity) {
    free(hist->lines[0]);
    memmove(hist->lines, hist->lines + 1, sizeof(char*) * (hist->capacity - 1));
    hist->count--;
  }
  
  hist->lines[hist->count++] = strdup(line);
  hist->current = hist->count;
}

static const char* history_prev(history_t *hist) {
  if (hist->count == 0) return NULL;
  if (hist->current > 0) hist->current--;
  return hist->lines[hist->current];
}

static const char* history_next(history_t *hist) {
  if (hist->count == 0) return NULL;
  if (hist->current < hist->count - 1) {
    hist->current++;
    return hist->lines[hist->current];
  }
  hist->current = hist->count;
  return "";
}

static void history_free(history_t *hist) {
  for (int i = 0; i < hist->count; i++) free(hist->lines[i]);
  free(hist->lines);
}

static char* get_history_path(void) {
  const char *home = getenv("HOME");
  if (!home) home = getenv("USERPROFILE");
  if (!home) return NULL;
  
  size_t len = strlen(home) + 32;
  char *path = malloc(len);
  snprintf(path, len, "%s/.ant", home);
  mkdir_p(path);
  snprintf(path, len, "%s/.ant/repl_history", home);
  return path;
}

static void history_load(history_t *hist) {
  char *path = get_history_path();
  if (!path) return;
  
  FILE *fp = fopen(path, "r");
  free(path);
  if (!fp) return;
  
  char line[MAX_LINE_LENGTH];
  while (fgets(line, sizeof(line), fp)) {
    size_t len = strlen(line);
    if (len > 0 && line[len - 1] == '\n') line[len - 1] = '\0';
    if (line[0]) history_add(hist, line);
  }
  fclose(fp);
}

static void history_save(history_t *hist) {
  char *path = get_history_path();
  if (!path) return;
  
  FILE *fp = fopen(path, "w");
  free(path);
  if (!fp) return;
  
  for (int i = 0; i < hist->count; i++) {
    fprintf(fp, "%s\n", hist->lines[i]);
  }
  fclose(fp);
}

typedef enum { 
  KEY_NONE, KEY_UP, KEY_DOWN, KEY_LEFT, KEY_RIGHT, 
  KEY_BACKSPACE, KEY_ENTER, KEY_EOF, KEY_CHAR 
} key_type_t;

typedef struct { key_type_t type; int ch; } key_event_t;
typedef void (*key_handler_t)(INPUT);

static crprintf_compiled *hl_prog = NULL;
static highlight_state hl_line_state = HL_STATE_INIT;
static int repl_last_render_rows = 1;

static int repl_terminal_cols(void) {
  int cols = 80;
#ifdef _WIN32
  CONSOLE_SCREEN_BUFFER_INFO csbi;
  if (GetConsoleScreenBufferInfo(GetStdHandle(STD_OUTPUT_HANDLE), &csbi)) {
    cols = csbi.srWindow.Right - csbi.srWindow.Left + 1;
  }
#else
  struct winsize ws;
  if (ioctl(STDOUT_FILENO, TIOCGWINSZ, &ws) == 0 && ws.ws_col > 0) {
    cols = ws.ws_col;
  }
#endif
  return cols > 0 ? cols : 80;
}

static void repl_move_to_line_start(const char *prompt, int pos, int cols) {
  int prompt_len = (int)strlen(prompt);
  int cursor_cols = prompt_len + pos;
  int cursor_row = cursor_cols / cols;
  if (cursor_cols > 0 && cursor_cols % cols == 0) cursor_row--;

  if (cursor_row > 0) {
    char move_buf[32];
    snprintf(move_buf, sizeof(move_buf), "\033[%dA", cursor_row);
    fputs(move_buf, stdout);
  }
  fputs("\r", stdout);
}

static void refresh_line(const char *line, int len, int pos, const char *prompt) {
  int cols = repl_terminal_cols();
  int prompt_len = (int)strlen(prompt);
  int line_cols = prompt_len + len;
  int current_rows = line_cols > 0 ? (line_cols - 1) / cols + 1 : 1;
  int rows = repl_last_render_rows > current_rows ? repl_last_render_rows : current_rows;

  repl_move_to_line_start(prompt, pos, cols);
  for (int i = 0; i < rows; i++) {
    fputs("\033[K", stdout);
    if (i < rows - 1) fputs("\033[B\r", stdout);
  }
  for (int i = 0; i < rows - 1; i++) fputs("\033[A", stdout);
  fputs("\r", stdout);

  fputs(prompt, stdout);

  if (crprintf_get_color() && len > 0 && len <= 2048) {
    char tagged[8192];
    char rendered[8192];
    
    highlight_state state = hl_line_state;
    ant_highlight_stateful(line, (size_t)len, tagged, sizeof(tagged), &state);
    
    hl_prog = crprintf_recompile(hl_prog, tagged);
    crprintf_state *rs = crprintf_state_new();
    
    crsprintf_compiled(rendered, sizeof(rendered), rs, hl_prog);
    crprintf_state_free(rs);
    
    fputs(rendered, stdout);
  } else if (len > 0) fwrite(line, 1, (size_t)len, stdout);

  int end_cols = prompt_len + len;
  int end_rows = end_cols > 0 ? (end_cols - 1) / cols + 1 : 1;
  int end_row = end_rows - 1;
  
  int cursor_cols = prompt_len + pos;
  int cursor_row = cursor_cols > 0 ? cursor_cols / cols : 0;
  int cursor_col = cursor_cols > 0 ? cursor_cols % cols : 0;
  int up_rows = end_row - cursor_row;

  if (up_rows > 0) {
    char move_buf[32];
    snprintf(move_buf, sizeof(move_buf), "\033[%dA", up_rows);
    fputs(move_buf, stdout);
  }
  
  fputs("\r", stdout);
  if (cursor_col > 0) {
    char move_buf[32];
    snprintf(move_buf, sizeof(move_buf), "\033[%dC", cursor_col);
    fputs(move_buf, stdout);
  }

  repl_last_render_rows = end_rows;

  fflush(stdout);
}

static void cursor_move(int *pos, int len, int dir) {
  if (dir < 0 && *pos > 0) { printf("\033[D"); fflush(stdout); (*pos)--; }
  else if (dir > 0 && *pos < len) { printf("\033[C"); fflush(stdout); (*pos)++; }
}

static void line_set(char *line, int *pos, int *len, const char *str, const char *prompt) {
  strcpy(line, str);
  *len = (int)strlen(line); *pos = *len;
  refresh_line(line, *len, *pos, prompt);
}

static void line_backspace(char *line, int *pos, int *len, const char *prompt) {
  if (*pos <= 0) return;
  
  memmove(line + *pos - 1, line + *pos, *len - *pos + 1);
  (*pos)--; (*len)--;
  refresh_line(line, *len, *pos, prompt);
}

static void line_insert(char *line, int *pos, int *len, int c, const char *prompt) {
  if (*len >= MAX_LINE_LENGTH - 1) return;
  
  memmove(line + *pos + 1, line + *pos, *len - *pos + 1);
  line[*pos] = (char)c;
  (*pos)++; (*len)++;
  refresh_line(line, *len, *pos, prompt);
}

#ifdef _WIN32
static key_event_t read_key(void) {
  if (ctrl_c_pressed > 0) return (key_event_t){ KEY_EOF, 0 };
  int c = _getch();
  if (c == 0 || c == 0xE0) {
    int ext = _getch();
    switch (ext) {
      case 72: return (key_event_t){ KEY_UP, 0 };
      case 80: return (key_event_t){ KEY_DOWN, 0 };
      case 77: return (key_event_t){ KEY_RIGHT, 0 };
      case 75: return (key_event_t){ KEY_LEFT, 0 };
    }
    return (key_event_t){ KEY_NONE, 0 };
  }
  if (c == 8) return (key_event_t){ KEY_BACKSPACE, 0 };
  if (c == '\r' || c == '\n') return (key_event_t){ KEY_ENTER, 0 };
  if (c == 3) { ctrl_c_pressed++; return (key_event_t){ KEY_EOF, 0 }; }
  if (c == 4 || c == 26) return (key_event_t){ KEY_EOF, 0 };
  if (isprint(c) || (unsigned char)c >= 0x80) return (key_event_t){ KEY_CHAR, c };
  return (key_event_t){ KEY_NONE, 0 };
}
#else
static struct termios saved_tio;
static key_event_t read_key(void) {
  if (ctrl_c_pressed > 0) return (key_event_t){ KEY_EOF, 0 };
  int c = getchar();
  if (c == EOF && !feof(stdin)) { clearerr(stdin); return (key_event_t){ KEY_EOF, 0 }; }
  if (c == EOF) return (key_event_t){ KEY_EOF, 0 };
  if (c == 27) {
    int seq1 = getchar();
    if (seq1 == EOF) return (key_event_t){ KEY_NONE, 0 };
    if (seq1 != '[') return (key_event_t){ KEY_NONE, 0 };
    int seq2 = getchar();
    if (seq2 == EOF) return (key_event_t){ KEY_NONE, 0 };
    switch (seq2) {
      case 'A': return (key_event_t){ KEY_UP, 0 };
      case 'B': return (key_event_t){ KEY_DOWN, 0 };
      case 'C': return (key_event_t){ KEY_RIGHT, 0 };
      case 'D': return (key_event_t){ KEY_LEFT, 0 };
    }
    if (seq2 >= '0' && seq2 <= '9') {
      int seq3 = getchar();
      (void)seq3;
    }
    return (key_event_t){ KEY_NONE, 0 };
  }
  if (c == 127 || c == 8) return (key_event_t){ KEY_BACKSPACE, 0 };
  if (c == '\n' || c == '\r') return (key_event_t){ KEY_ENTER, 0 };
  if (isprint(c) || (unsigned char)c >= 0x80) return (key_event_t){ KEY_CHAR, c };
  return (key_event_t){ KEY_NONE, 0 };
}
#endif

static void handle_up(INPUT) {
  const char *h = history_prev(hist);
  if (h) line_set(line, pos, len, h, prompt);
}

static void handle_down(INPUT) {
  const char *h = history_next(hist);
  if (h) line_set(line, pos, len, h, prompt);
}

static void handle_left(INPUT) { cursor_move(pos, *len, -1); }
static void handle_right(INPUT) { cursor_move(pos, *len, 1); }

static void handle_backspace(INPUT) { line_backspace(line, pos, len, prompt); }
static void handle_char(INPUT) { line_insert(line, pos, len, key->ch, prompt); }

static key_handler_t handlers[] = {
  [KEY_UP] = handle_up,
  [KEY_DOWN] = handle_down,
  [KEY_LEFT] = handle_left,
  [KEY_RIGHT] = handle_right,
  [KEY_BACKSPACE] = handle_backspace,
  [KEY_CHAR] = handle_char,
};

static inline void term_restore(void) {
#ifndef _WIN32
  tcsetattr(STDIN_FILENO, TCSANOW, &saved_tio);
#endif
}

static char *read_line_with_history(history_t *hist, ant_t *js, const char *prompt) {
  char *line = malloc(MAX_LINE_LENGTH);
  int pos = 0, len = 0; line[0] = '\0';
  repl_last_render_rows = 1;
  
  #ifndef _WIN32
  struct termios new_tio;
  tcgetattr(STDIN_FILENO, &saved_tio);
  new_tio = saved_tio;
  new_tio.c_lflag &= ~(ICANON | ECHO);
  tcsetattr(STDIN_FILENO, TCSANOW, &new_tio);
  #endif

  again:
  key_event_t key = read_key();
  
  switch (key.type) {
  case KEY_ENTER:
    putchar('\n');
    term_restore();
    return line;
  
  case KEY_EOF:
    putchar('\n');
    term_restore();
    free(line);
    return NULL;
  
  default:
    if (handlers[key.type]) handlers[key.type](
      line, &pos, &len, &key, hist, prompt
    );
    break;
  }
  
  goto again;
}

typedef struct {
  int paren, bracket, brace;
  int *templates;
  int template_count, template_cap;
  char string_char;
  bool in_string, escaped;
} parse_state_t;

static void push_template(parse_state_t *s) {
  if (s->template_count >= s->template_cap) {
    s->template_cap = s->template_cap ? s->template_cap * 2 : 8;
    int *new_templates = realloc(s->templates, s->template_cap * sizeof(int));
    if (!new_templates) { return; } s->templates = new_templates;
  }
  s->templates[s->template_count++] = s->brace;
}

static bool in_template_text(parse_state_t *s) {
  return s->template_count > 0 && s->brace == s->templates[s->template_count - 1];
}

static bool is_incomplete_input(const char *code, size_t len) {
  parse_state_t s = {0};
  
  for (size_t i = 0; i < len; i++) {
    char c = code[i];
    
    if (s.escaped) { s.escaped = false; continue; }
    if (c == '\\' && (s.in_string || s.template_count > 0)) { s.escaped = true; continue; }
    if (s.in_string) { if (c == s.string_char) s.in_string = false; continue; }
    
    if (in_template_text(&s)) {
      if (c == '`') s.template_count--;
      else if (c == '$' && i + 1 < len && code[i + 1] == '{') { s.brace++; i++; }
      continue;
    }
    
    if (c == '/' && i + 1 < len) {
      if (code[i + 1] == '/') { while (i < len && code[i] != '\n') i++; continue; }
      if (code[i + 1] == '*') {
        for (i += 2; i + 1 < len && !(code[i] == '*' && code[i + 1] == '/'); i++);
        if (i + 1 >= len) { free(s.templates); return true; }
        i++; continue;
      }
      if (js_regex_can_start(code, i)) {
        size_t regex_end = 0;
        if (!js_scan_regex_literal(code, len, i, &regex_end)) { free(s.templates); return true; }
        i = regex_end - 1;
        continue;
      }
    }
    
    switch (c) {
      case '"': case '\'': s.in_string = true; s.string_char = c; break;
      case '`': push_template(&s); break;
      case '(': s.paren++; break;   case ')': s.paren--; break;
      case '[': s.bracket++; break; case ']': s.bracket--; break;
      case '{': s.brace++; break;   case '}': s.brace--; break;
    }
  }
  
  bool incomplete = s.in_string || s.template_count > 0 || s.paren > 0 || s.bracket > 0 || s.brace > 0;
  free(s.templates);
  return incomplete;
}

void ant_repl_run() {
  ant_t *js = rt->js;
  
  js_set_filename(js, "[repl]");
  js_setup_import_meta(js, "[repl]");
  
  printf("Welcome to Ant JavaScript v%s\n", ANT_VERSION);
  printf("Type \".help\" for more information.\n");
  
#ifdef _WIN32
  signal(SIGINT, sigint_handler);
#else
  struct sigaction sa;
  sa.sa_handler = sigint_handler;
  sigemptyset(&sa.sa_mask);
  sa.sa_flags = 0;
  sigaction(SIGINT, &sa, NULL);
#endif
  
  history_t history;
  history_init(&history);
  history_load(&history);
  
  repl_decl_registry_t decl_registry = {0};
  g_repl_decl_registry = &decl_registry;
  
  js_set(js, js_glob(js), "__dirname", js_mkstr(js, ".", 1));
  js_set(js, js_glob(js), "__filename", js_mkstr(js, "[repl]", 6));

  int prev_ctrl_c_count = 0;
  char *multiline_buf = NULL;
  size_t multiline_len = 0;
  size_t multiline_cap = 0;
  
  while (1) {
    const char *prompt = multiline_buf ? "| " : "> ";

    if (multiline_buf && multiline_len > 0) {
      char scratch[8192];
      hl_line_state = HL_STATE_INIT;
      ant_highlight_stateful(multiline_buf, multiline_len, scratch, sizeof(scratch), &hl_line_state);
    } else hl_line_state = HL_STATE_INIT;

    fputs(prompt, stdout);
    fflush(stdout);
    
    ctrl_c_pressed = 0;
    char *line = read_line_with_history(&history, js, prompt);
    
    if (ctrl_c_pressed > 0) {
      if (multiline_buf) {
        free(multiline_buf);
        multiline_buf = NULL;
        multiline_len = 0;
        multiline_cap = 0;
        prev_ctrl_c_count = 0;
        if (line) free(line);
        continue;
      }
      if (prev_ctrl_c_count > 0) {
        if (line) free(line);
        break;
      }
      printf("(To exit, press Ctrl+C again or type .exit)\n");
      prev_ctrl_c_count++;
      if (line) free(line);
      continue;
    }
    
    if (line == NULL) {
      if (multiline_buf) {
        free(multiline_buf);
        multiline_buf = NULL;
        multiline_len = 0;
        multiline_cap = 0;
        continue;
      }
      break;
    }
    
    prev_ctrl_c_count = 0;
    size_t line_len = strlen(line);
    
    if (line_len == 0 && multiline_buf) {
      if (multiline_len + 1 >= multiline_cap) {
        multiline_cap = multiline_cap ? multiline_cap * 2 : 256;
        multiline_buf = realloc(multiline_buf, multiline_cap);
      }
      multiline_buf[multiline_len++] = '\n';
      multiline_buf[multiline_len] = '\0';
      free(line);
      continue;
    }
    
    if (line_len == 0) {
      free(line);
      continue;
    }
    
    if (!multiline_buf && line[0] == '.') {
      cmd_result_t result = execute_command(js, &history, line);
      if (result == CMD_EXIT) {
        free(line);
        break;
      } else if (result == CMD_NOT_FOUND) {
        printf("Unknown command: %s\n", line);
        printf("Type \".help\" for more information.\n");
      }
      free(line);
      continue;
    }
    
    size_t new_len = multiline_len + line_len + 1;
    if (new_len >= multiline_cap || !multiline_buf) {
      multiline_cap = multiline_cap ? multiline_cap * 2 : 256;
      if (multiline_cap < new_len + 1) multiline_cap = new_len + 1;
      multiline_buf = realloc(multiline_buf, multiline_cap);
    }
    
    if (multiline_len > 0) {
      multiline_buf[multiline_len++] = '\n';
    }
    memcpy(multiline_buf + multiline_len, line, line_len);
    multiline_len += line_len;
    multiline_buf[multiline_len] = '\0';
    free(line);
    
    if (is_incomplete_input(multiline_buf, multiline_len)) continue;
    history_add(&history, multiline_buf);
    
    repl_eval_chunk(
      js, &decl_registry, multiline_buf, 
      multiline_len, REPL_PRINT_INTERACTIVE
    );
    
    free(multiline_buf);
    multiline_buf = NULL;
    multiline_len = 0;
    multiline_cap = 0;
  }
  
  if (multiline_buf) free(multiline_buf);
  if (hl_prog) { crprintf_compiled_free(hl_prog); hl_prog = NULL; }
  
  repl_decl_registry_free(&decl_registry);
  g_repl_decl_registry = NULL;
  
  history_save(&history);
  history_free(&history);
}