#include "monitor_kernel.h"

unsigned char w_arg2k_w_arg(void *ptr, watch_arg warg,
                            kernel_watch_arg *k_watch_arg) {
  // k_watch_arg init
  k_watch_arg->task_id = warg.task_id;
  strncpy(k_watch_arg->name, warg.name, MAX_NAME_LEN + 1); // name
  k_watch_arg->name[MAX_NAME_LEN + 1] = '\0';              // just in case
  k_watch_arg->kptr = ptr;
  k_watch_arg->length_byte = warg.length_byte;
  k_watch_arg->threshold = warg.threshold;
  k_watch_arg->unsigned_flag = warg.unsigned_flag;
  k_watch_arg->greater_flag = warg.greater_flag;
  return 0;
}

/// @brief get a valuable timer
/// @param time_ns
/// @return kernel_watch_timer *, NULL means fail
kernel_watch_timer *get_timer(unsigned long long time_ns) {
  int i = 0;
  kernel_watch_timer *timer = NULL;
  // chose a timer
  for (i = 0; i < kernel_wtimer_num; i++) {
    timer = &kernel_wtimer_list[i];

    if (TIMER_EMPTY(timer)) {
      break;
    }
    if ((timer->time_ns == time_ns) && (!TIMER_FILLED(timer))) {
      break;
    }
  }
  // if all timer is full
  if (i >= MAX_TIMER_NUM) {
    printk(KERN_ERR "No timer available\n");
    return NULL;
  }
  // if a new timer, init it
  if (i > kernel_wtimer_num - 1) {
    printk(KERN_INFO "New timer\n");

    kernel_wtimer_list[i].time_ns = time_ns;
    kernel_wtimer_list[i].sentinel = 0;

    kernel_wtimer_list[i].kt = ktime_set(0, (unsigned long)time_ns); // ns
    // CLOCK_MONOTONIC: time since boot | HRTIMER_MODE_REL : relative time
    hrtimer_init(&(kernel_wtimer_list[i].hr_timer), CLOCK_MONOTONIC,
                 HRTIMER_MODE_REL);
    kernel_wtimer_list[i].hr_timer.function =
        check_variable_cb; // callback function

    kernel_wtimer_num = i + 1;
  }
  printk(KERN_INFO "now, we have %d timers\n", kernel_wtimer_num);
  return &kernel_wtimer_list[i];
}

/// @brief hrTimer add watch
/// @param timer
/// @param k_watch_arg
/// @return 0 is success
unsigned char timer_add_watch(kernel_watch_timer *timer,
                              kernel_watch_arg k_watch_arg) {
  if (TIMER_FILLED(timer)) {
    printk(KERN_ERR "Timer is full\n");
    return -1;
  }
  memcpy(&timer->k_watch_args[timer->sentinel], &k_watch_arg,
         sizeof(k_watch_arg));
  // timer->k_watch_args[timer->sentinel] = k_watch_arg;
  timer->sentinel++;
  return 0;
}

unsigned char timer_del_watch_by_pid(kernel_watch_timer *timer, pid_t pid) {
  int i = 0;
  for (i = 0; i < timer->sentinel; i++) {
    // if pid match, delete it and move the last one to this position, check
    // again
    if (timer->k_watch_args[i].task_id == pid) {
      if (i != timer->sentinel - 1) {
        memcpy(&timer->k_watch_args[i],
               &timer->k_watch_args[timer->sentinel - 1],
               sizeof(kernel_watch_arg));
      }
      timer->sentinel--;
      i--;
    }
  }
  return 0;
}

/// @brief transfer user space address to kernel space address
///        change static global "kaddr" and "page" value
/// @param pid: process id
/// @param kaddr: user space address
/// @return kernel space address + offset
void *convert_user_space_ptr(pid_t pid, unsigned long addr) {
  struct task_struct *task;
  struct mm_struct *mm;
  int ret;

  // unsigned long aligned_addr = 0;
  // unsigned long offset = 0;

  watch_local_memory *node;

  // if (addr < TASK_SIZE || addr > -PAGE_SIZE)
  // {
  //     printk(KERN_ERR "Invalid address\n");
  //     return NULL;
  // }

  // for get_user_pages_remote
  unsigned long aligned_addr = addr & PAGE_MASK;
  unsigned long offset = addr & ~PAGE_MASK;

  printk(KERN_INFO "%s\n", __FUNCTION__);

  node = kmalloc(sizeof(watch_local_memory), GFP_KERNEL);
  node->task_id = pid;

  // Find the task with pid
  rcu_read_lock();
  task = pid_task(find_vpid(pid), PIDTYPE_PID);
  rcu_read_unlock();

  if (!task) {
    printk(KERN_ERR "Cannot find task for PID %d\n", pid);
    kfree(node); // careful there is kfree
    return NULL;
  }
  // Get memory descriptor
  mm = get_task_mm(task);
  if (!mm) {
    printk(KERN_ERR "Cannot get memory descriptor\n");
    kfree(node); // careful there is kfree
    return NULL;
  }
  down_read(&task->mm->mmap_lock);
  ret = get_user_pages_remote(task->mm, aligned_addr, 1, FOLL_FORCE,
                              &(node->page), NULL, NULL);
  up_read(&task->mm->mmap_lock);

  if (ret != 1) {
    printk(KERN_ERR "Cannot get user page\n");
    kfree(node); // careful there is kfree
    return NULL;
  }
  // Map the page to kernel space
  node->kaddr = kmap(node->page);
  list_add_tail(&node->entry, &watch_local_memory_list); // add to list
  // printk(KERN_INFO "node->kaddr: %p, aligned_addr: %ld, offset: %ld\n",
  // node->kaddr, aligned_addr, offset);
  return (void *)((unsigned long)(node->kaddr) + offset);
}

/// @brief free page in watch_local_memory_list with task_id
/// @param task_id
void free_page_list(pid_t task_id) {
  watch_local_memory *node, *next;
  list_for_each_entry_safe(node, next, &watch_local_memory_list, entry) {
    if (node == NULL)
      break;
    if (node->task_id == task_id) {
      // unmap and release the page
      if (node->kaddr)
        kunmap(node->kaddr);
      if (node->page)
        put_page(node->page);
      list_del(&node->entry);
      kfree(node); // careful there is kfree
    }
  }
}

/// @brief free all page in watch_local_memory_list
/// @param
void free_all_page_list(void) {
  watch_local_memory *node, *next;
  list_for_each_entry_safe(node, next, &watch_local_memory_list, entry) {
    if (node == NULL)
      break;
    // unmap and release the page
    if (node->kaddr)
      kunmap(node->kaddr);
    if (node->page)
      put_page(node->page);
    list_del(&node->entry);
    kfree(node); // careful there is kfree
  }
}

/// @brief hrTimer handler
enum hrtimer_restart check_variable_cb(struct hrtimer *timer) {
  kernel_watch_timer *k_watch_timer =
      container_of(timer, kernel_watch_timer, hr_timer);
  int i = 0, j = 0;
  int buffer[TIMER_MAX_WATCH_NUM]; // Buffer to store the messages

  // check all watched kernel_watch_arg
  for (i = 0; i < k_watch_timer->sentinel; i++) {
    if (read_and_compare(&k_watch_timer->k_watch_args[i])) {
      // snprintf(buffer + strlen(buffer), sizeof(buffer) - strlen(buffer), "
      // name: %s, threshold: %lld, pid: %d\n",
      //          k_watch_timer->k_watch_args[i].name,
      //          k_watch_timer->k_watch_args[i].threshold,
      //          k_watch_timer->k_watch_args[i].task_id);
      buffer[j] = i;
      j++;

      // printk(KERN_INFO "j: name %s, threshold: %lld\n",
      // k_watch_timer->k_watch_args[i].name,
      //        k_watch_timer->k_watch_args[i].threshold);
      // printk(KERN_INFO "j: %d\n", j);
    }
  }
  if (j > 0) // if any threshold reached
  {
    printk("-------------------------------------\n");
    printk("-------------watch monitor-----------\n");
    printk("Threshold reached:\n");

    for (i = 0; i < j; i++) {
      printk(" name: %s, threshold: %lld, pid: %d\n",
             k_watch_timer->k_watch_args[buffer[i]].name, //! todo
             k_watch_timer->k_watch_args[buffer[i]].threshold,
             k_watch_timer->k_watch_args[buffer[i]].task_id);
    }
    print_task_stack();
    // restart timer after 1s
    hrtimer_forward(timer, timer->base->get_time(), ktime_set(1, 0)); //! todo
    printk("-------------------------------------\n");
  } else {
    // keep frequency
    hrtimer_forward(timer, timer->base->get_time(), k_watch_timer->kt);
  }
  return HRTIMER_RESTART; // restart timer
}

/// @brief start hrTimer
/// @param timeout: timeout in us
/// @return 0 is success
// int start_hrTimer(unsigned long timeout)
// {
//     printk("HrTimer Start\n");

//     kt = ktime_set(0, (unsigned long)timeout); //  us -> ns
//     // CLOCK_MONOTONIC: time since boot | HRTIMER_MODE_REL : relative time
//     hrtimer_init(&hr_timer, CLOCK_MONOTONIC, HRTIMER_MODE_REL);
//     hr_timer.function = check_variable_cb;
//     // mode the same as hrtimer_init
//     hrtimer_start(&hr_timer, kt, HRTIMER_MODE_REL);
//     return 0;
// }

/// @brief start all hrTimer
/// @param
void start_all_hrTimer(void) {
  int i = 0;
  kernel_watch_timer *timer = NULL;
  for (i = 0; i < kernel_wtimer_num; i++) {
    timer = &(kernel_wtimer_list[i]);
    TIMER_START(timer);
  }
  printk("HrTimer start,module keep %d hrtimer for now\n", kernel_wtimer_num);
}

/// @brief cancel hrTimer
/// @param
void cancel_all_hrTimer(void) {
  int i = 0;
  kernel_watch_timer *timer = NULL;
  for (i = 0; i < kernel_wtimer_num; i++) {
    timer = &(kernel_wtimer_list[i]);
    TIMER_CANCEL(timer);
  }

  printk("HrTimer cancel,module keep %d hrtimer for now\n", kernel_wtimer_num);
}

// for read_and_compare
typedef unsigned char (*compare_func)(void *, long long);

unsigned char compare_1_byte_signed(void *ptr, long long threshold) {
  // printk("compare_1_byte_signed: value %d, biss: %lld\n", *(char *)ptr,
  // threshold);
  return *(char *)ptr > threshold;
}
unsigned char compare_1_byte_unsigned(void *ptr, long long threshold) {
  // printk("compare_1_byte_unsigned: value %d, biss: %lld\n", *(unsigned char
  // *)ptr, threshold);
  return *(unsigned char *)ptr > threshold;
}
unsigned char compare_2_byte_signed(void *ptr, long long threshold) {
  // printk("compare_2_byte_signed: value %d, biss: %lld\n", *(short int *)ptr,
  // threshold);
  return *(short int *)ptr > threshold;
}
unsigned char compare_2_byte_unsigned(void *ptr, long long threshold) {
  // printk("compare_2_byte_unsigned: value %d, biss: %lld\n", *(unsigned short
  // int *)ptr, threshold);
  return *(unsigned short int *)ptr > threshold;
}
unsigned char compare_4_byte_signed(void *ptr, long long threshold) {
  // printk("compare_4_byte_signed: value %d, biss: %lld\n", *(int *)ptr,
  // threshold);
  return *(int *)ptr > threshold;
}
unsigned char compare_4_byte_unsigned(void *ptr, long long threshold) {
  // printk("compare_4_byte_unsigned: value %d, biss: %lld\n", *(unsigned int
  // *)ptr, threshold);
  return *(unsigned int *)ptr > threshold;
}
unsigned char compare_8_byte_signed(void *ptr, long long threshold) {
  // printk("compare_8_byte_signed: value %lld, biss: %lld\n", *(long long
  // *)ptr, threshold);
  return *(long long *)ptr > threshold;
}
unsigned char compare_8_byte_unsigned(void *ptr, long long threshold) {
  // printk("compare_8_byte_unsigned: value %lld, biss: %lld\n", *(unsigned long
  // long *)ptr, threshold);
  return *(unsigned long long *)ptr > threshold;
}
// list of compare functions
static compare_func compare_funcs[8] = {
    compare_1_byte_signed,   compare_2_byte_signed,   compare_4_byte_signed,
    compare_8_byte_signed,   compare_1_byte_unsigned, compare_2_byte_unsigned,
    compare_4_byte_unsigned, compare_8_byte_unsigned};

static int func_indices[2][9] = {{0, 0, 1, 0, 2, 0, 0, 0, 3},
                                 {0, 4, 5, 0, 6, 0, 0, 0, 7}};

/// @brief read k_arg->kptr and compare with threshold
/// @param k_arg
/// @return result of compare
unsigned char read_and_compare(kernel_watch_arg *k_arg) {
  void *ptr = k_arg->kptr;
  int len = k_arg->length_byte;
  unsigned char is_unsigned = k_arg->unsigned_flag;
  long long threshold = k_arg->threshold;

  unsigned char result = 0;

  // if (len != 1 && len != 2 && len != 4 && len != 8)
  // {
  //     printk(KERN_ERR "Invalid length\n");
  //     return 0;
  // }

  result = compare_funcs[func_indices[is_unsigned][len]](ptr, threshold);

  // printk(KERN_INFO "read_and_compare: name %s, value %d, biss: %lld, result:
  // %d \n", k_arg->name, *(int *)ptr,
  //        threshold, result);

  if (k_arg->greater_flag)
    return result;
  else
    return !result;
}

/// @brief init kallsyms_lookup_name
/// @param
/// @return 0 is success
int fn_kallsyms_lookup_name_init(void) {
  register_kprobe(&kprobe_kallsyms_lookup_name);
  diag_kallsyms_lookup_name = (void *)kprobe_kallsyms_lookup_name.addr;
  unregister_kprobe(&kprobe_kallsyms_lookup_name);

  printk("xby-debug, diag_kallsyms_lookup_name is %p\n",
         diag_kallsyms_lookup_name);

  if (!diag_kallsyms_lookup_name) {
    return -EINVAL;
  }
  return 0;
}

unsigned char del_all_kwarg_by_pid(pid_t pid) {
  int i = 0;
  kernel_watch_timer *timer = NULL;

  printk(KERN_INFO "del kwarg...");

  for (i = 0; i < kernel_wtimer_num; i++) {
    timer = &(kernel_wtimer_list[i]);
    timer_del_watch_by_pid(timer, pid);
  }
  for (i = 0; i < kernel_wtimer_num; i++) {
    timer = &(kernel_wtimer_list[i]);
    if (TIMER_NO_KWARG(timer)) // no available kwarg
    {
      if (i != kernel_wtimer_num - 1) {
        memcpy(timer, &kernel_wtimer_list[kernel_wtimer_num - 1],
               sizeof(kernel_watch_timer));
      }
      kernel_wtimer_num--;
      i--;
    }
  }
  return 0;
}

/// @brief clear watch with pid
/// @param pid
void clear_watch(pid_t pid) {
  printk(KERN_INFO "clear pid %d 's watch variable\n", pid);
  cancel_all_hrTimer();      // just in case
  del_all_kwarg_by_pid(pid); // delete all kwarg with pid
  free_page_list(pid);       // free page with pid
  start_all_hrTimer();       // restart timer
}

/// @brief clear all watch and reset kernel_wtimer_list/kernel_wtimer_num
/// @param
void clear_all_watch(void) {
  printk(KERN_INFO "clear all watch variable\n");
  // unmap and release the page
  free_all_page_list();
  // cancel timer
  cancel_all_hrTimer();
  // clear timer
  kernel_wtimer_num = 0;
  memset(kernel_wtimer_list, 0, sizeof(kernel_wtimer_list));
}