#include "monitor_kernel.h"

#include <linux/sched.h>
#include <linux/stacktrace.h>
// #include <linux/sched/task.h>
// #include <linux/sched/mm.h>

#define __task_contributes_to_load(task)                                       \
  ((READ_ONCE(task->__state) & TASK_UNINTERRUPTIBLE) != 0 &&                   \
   (task->flags & PF_FROZEN) == 0 &&                                           \
   (READ_ONCE(task->__state) & TASK_NOLOAD) == 0)

/**
 * @brief watch_arg to kernel_watch_arg
 *
 * @param ptr: kernel space address
 * @param warg: watch_arg
 * @param k_watch_arg: kernel_watch_arg
 * @return unsigned char
 */
static unsigned char w_arg2k_w_arg(void *kptr, 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->ptr = warg.ptr;
  k_watch_arg->kptr = kptr;
  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;
}

static void k_w_arg2threshold(kernel_watch_arg *k_watch_arg,
                              threshold *threshold) {
  threshold->task_id = k_watch_arg->task_id;
  strncpy(threshold->name, k_watch_arg->name, MAX_NAME_LEN + 1);
  threshold->name[MAX_NAME_LEN + 1] = '\0';
  threshold->ptr = k_watch_arg->ptr;
  threshold->threshold = k_watch_arg->threshold;
}

static void init_mm_tree(mm_tree *mm_tree) {
  INIT_RADIX_TREE(&mm_tree->mm_tree, GFP_ATOMIC);
  spin_lock_init(&mm_tree->mm_tree_lock);
}

static int init_buffer(unsigned int buf_size) {
  init_mm_tree(&mm_tree_struct); // init mm_tree
  init_diag_variant_buffer(&load_monitor_variant_buffer, buf_size);
  int ret = 0;
  ret = alloc_diag_variant_buffer(&load_monitor_variant_buffer);
  return ret;
}

static void diag_tsk(struct task_struct *p, variable_monitor_task *tsk_info) {
  unsigned int nr_bt;
  printk(KERN_INFO "diag_tsk\n");
  diag_task_brief(p, &tsk_info->task); // task brief
  // printk("1\n");
  diag_task_user_stack(p, &tsk_info->user_stack); // user stack
  // printk("2\n");
  nr_bt = diag_task_kern_stack(p, &tsk_info->kern_stack); // kernel stack
  // int i = 0;
  // printk("pid: %d, kernel stack.stack\n", p->pid);
  // for (i = 0; i < nr_bt; i++) {
  //   printk("%lx\n", tsk_info->kern_stack.stack[i]);
  // }
  // printk("pid: %d, stack_trace_print\n", p->pid);
  // stack_trace_print(tsk_info->kern_stack.stack, nr_bt, 0); /* 打印栈 */

  // printk("3\n");
  dump_proc_chains_argv(1, p, &mm_tree_struct,
                        &tsk_info->proc_chains); // proc chains
  diag_task_raw_stack(p, &tsk_info->raw_stack);   // raw stack
}

static void push_tsk_info(variable_monitor_task *tsk_info,unsigned long *flags) {
  printk(KERN_INFO "push_tsk_info\n");
  diag_variant_buffer_spin_lock(&load_monitor_variant_buffer, *flags);
  diag_variant_buffer_reserve(&load_monitor_variant_buffer,
                              sizeof(variable_monitor_task));
  diag_variant_buffer_write_nolock(&load_monitor_variant_buffer, tsk_info,
                                   sizeof(variable_monitor_task));
  diag_variant_buffer_seal(&load_monitor_variant_buffer);
  diag_variant_buffer_spin_unlock(&load_monitor_variant_buffer, *flags);
}

/// @brief clear all watch and reset kernel_wtimer_list/kernel_wtimer_num
/// @param
static 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));
}

static void sample_task_work(struct work_struct *work){
    kernel_watch_timer *k_watch_timer = container_of(work, kernel_watch_timer, wk);

    if (k_watch_timer->threshold_num <= 0) return;

    struct task_struct *g, *p; // g: task group; p: task
    unsigned long flags;
    unsigned long event_id = get_cycles();
  
    static variable_monitor_task tsk_info = {0};
    static variable_monitor_record vm_record = {0};
    kernel_watch_arg *kwarg;

    vm_record.id = event_id;
    vm_record.et_type = 0; //! todo event type
    vm_record.tv = ktime_get_real();
    vm_record.threshold_num = k_watch_timer->threshold_num;

    int i;
    for (i = 0; i < vm_record.threshold_num; i++) {
      kwarg = &k_watch_timer->k_watch_args[k_watch_timer->threshold_buffer[i]];
      k_w_arg2threshold(kwarg, &vm_record.threshold_record[i]);
    }

    rcu_read_lock();
    
    diag_variant_buffer_spin_lock(&load_monitor_variant_buffer, flags);
    diag_variant_buffer_reserve(&load_monitor_variant_buffer,
                                sizeof(variable_monitor_record));
    diag_variant_buffer_write_nolock(&load_monitor_variant_buffer, &vm_record,
                                     sizeof(variable_monitor_record));
    diag_variant_buffer_seal(&load_monitor_variant_buffer);
    diag_variant_buffer_spin_unlock(&load_monitor_variant_buffer, flags);

    rcu_read_unlock();

    do_each_thread(g, p) {
      if (p->__state == TASK_RUNNING || __task_contributes_to_load(p) ||
          p->__state == TASK_IDLE || 1) {

        get_task_struct(p); // count +1

        tsk_info.et_type = 1; //! todo event type
        tsk_info.id = event_id;
        tsk_info.tv = vm_record.tv;
        diag_tsk(p, &tsk_info);
        
        put_task_struct(p); // count -1

        push_tsk_info(&tsk_info, &flags); // push to buffer
      }
    }
    while_each_thread(g, p);
}
/**
 * @brief all module function init. orig_X | buffer | workqueue
 *
 * @return int
 */
int monitor_init(void) {
  int ret = 0;

  ret = init_orig_fun(); // init orig_X
  if (ret)
    return ret;
  ret = init_buffer(50 * 1024 * 1024); // 50M
  if (ret)
    return -1;
  // init workqueue
  int i;
  for (i=0; i < MAX_TIMER_NUM; i++) {
    kernel_watch_timer *kw_timer = &kernel_wtimer_list[i];
    INIT_WORK(&kw_timer->wk, sample_task_work);
  }
  return 0;
}

/**
 * @brief monitor exit: clear all watch and free buffer
 *
 */
void monitor_exit(void) {
  // clear all watch
  clear_all_watch();
  // free buffer
  destroy_diag_variant_buffer(&load_monitor_variant_buffer);
  printk(KERN_INFO "clear all buffer\n");
}

/**
 * @brief start watch variable
 *
 * @param warg: uapi watch_arg
 * @return int 0 is success
 * !todo: adjust printk
 */
int start_watch_variable(watch_arg warg) {
  void *kptr;
  kernel_watch_timer *timer = NULL;
  kernel_watch_arg k_watch_arg;

  // user space address to kernel space address
  kptr = convert_user_space_ptr(warg.task_id, (unsigned long)warg.ptr);
  if (kptr == NULL) {
    printk(KERN_ERR "Cannot access user space\n");
    return -EACCES;
  }
  // check length
  if (warg.length_byte != 1 && warg.length_byte != 2 && warg.length_byte != 4 &&
      warg.length_byte != 8) {
    printk(KERN_ERR "Invalid length %d\n", warg.length_byte);
    return -EINVAL;
  }
  // k_watch_arg init
  w_arg2k_w_arg(kptr, warg, &k_watch_arg);
  timer = get_timer(warg.time_ns); // get a valuable timer

  printk(KERN_INFO "ptr transform kptr: %p\n", kptr);
  printk(KERN_INFO "timer: %p\n", timer);
  printk(KERN_INFO "timer->sentinel: %d, timer->time_ns: %lld\n",
         timer->sentinel, timer->time_ns);
  printk(KERN_INFO "timer->hr_timer: %p\n", &timer->hr_timer);

  TIMER_CANCEL(timer); // just in case
  timer_add_watch(timer, k_watch_arg);
  TIMER_START(timer);

  printk(KERN_INFO "Start watching var: %s\n", warg.name);
  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 main callback function
 *
 * @param timer
 * @return enum hrtimer_restart
 */
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;
  kernel_watch_arg *kwarg;

  // check all watched kernel_watch_arg
  for (i = 0; i < k_watch_timer->sentinel; i++) {
    kwarg = &k_watch_timer->k_watch_args[i];
    if (read_and_compare(kwarg->kptr, kwarg->length_byte, kwarg->greater_flag,
                         kwarg->unsigned_flag, kwarg->threshold)) {
      k_watch_timer->threshold_buffer[j] = i;
      j++;
    }
  }
  if (j > 0) // if any threshold reached
  {
    k_watch_timer->threshold_num = j;
    // restart timer after 5s
    hrtimer_forward(timer, timer->base->get_time(), ktime_set(5, 0)); //! todo
    // highpri_wq
    queue_work(system_highpri_wq, &k_watch_timer->wk);
  } else {
    // keep frequency
    hrtimer_forward(timer, timer->base->get_time(), k_watch_timer->kt);
  }
  return HRTIMER_RESTART; // restart timer
}

// static int diag_test(int nid); // for test
// static void test(struct task_struct *p, variable_monitor_task *tsk_info){
//   // unsigned int nr_bt;
//   printk(KERN_INFO "diag_tsk\n");
//   diag_task_brief(p, &tsk_info->task); // task brief
//   // printk("1\n");
//   diag_task_user_stack(p, &tsk_info->user_stack); // user stack
//   diag_task_kern_stack(p, &tsk_info->kern_stack); // kernel stack
//   dump_proc_chains_argv(1, p, &mm_tree_struct,
//                         &tsk_info->proc_chains); // proc chains
//   diag_task_raw_stack(p, &tsk_info->raw_stack);   // raw stack
//   printk(KERN_INFO "diag_tsk finish\n");
// }

// static void test2(variable_monitor_task *tsk_info, unsigned long flags){
//   printk(KERN_INFO "test2\n");
//   diag_variant_buffer_spin_lock(&load_monitor_variant_buffer, flags);
//   diag_variant_buffer_reserve(&load_monitor_variant_buffer,sizeof(variable_monitor_task));
//   diag_variant_buffer_write_nolock(&load_monitor_variant_buffer, tsk_info,
//                                    sizeof(variable_monitor_task));
//   diag_variant_buffer_seal(&load_monitor_variant_buffer);
//   diag_variant_buffer_spin_unlock(&load_monitor_variant_buffer, flags);
//   printk(KERN_INFO "test2 finish\n");
// }

int diag_test(int nid){
  // static struct task_struct *tsk;
  // static struct task_struct *leader;
  // static variable_monitor_task tsk_info;
  // // unsigned int nr_bt;

  // int ret;
  // unsigned long flags;
	// pid_t id = (pid_t)nid;

  // rcu_read_lock();
  // tsk = NULL;
	// if (orig_find_task_by_vpid)
	// 	tsk = orig_find_task_by_vpid(id);
	// if (!tsk) {
	// 	ret = -EINVAL;
	// 	rcu_read_unlock();
	// 	return ret;
	// }

	// leader = tsk->group_leader;
	// if (leader == NULL || leader->exit_state == EXIT_ZOMBIE){
	// 	ret = -EINVAL;
	// 	rcu_read_unlock();
	// 	return ret;
	// }

  // get_task_struct(tsk);
  // rcu_read_unlock();

  // tsk_info.et_type = 1;
  // test(tsk, &tsk_info);

  // put_task_struct(tsk);
  // printk(KERN_INFO "put_task_struct finish\n");

  // diag_variant_buffer_spin_lock(&load_monitor_variant_buffer, flags);
  // printk(KERN_INFO "1\n");
  // diag_variant_buffer_reserve(&load_monitor_variant_buffer,sizeof(variable_monitor_task));
  // printk(KERN_INFO "2\n");
  // diag_variant_buffer_write_nolock(&load_monitor_variant_buffer, &tsk_info,
  //                                  sizeof(variable_monitor_task));
  // printk(KERN_INFO "3\n");
  // diag_variant_buffer_seal(&load_monitor_variant_buffer);
  // printk(KERN_INFO "4\n");
  // diag_variant_buffer_spin_unlock(&load_monitor_variant_buffer, flags);
  // printk(KERN_INFO "5\n");

    struct task_struct *g, *p; // g: task group; p: task
    unsigned long flags;
    unsigned long event_id = get_cycles();

    static variable_monitor_task tsk_info = {0};
    static variable_monitor_record vm_record = {0};

    // vm_record.id = event_id;
    // vm_record.et_type = 0; //! todo event type
    vm_record.tv = ktime_get_real();
    // vm_record.threshold_num = j;

    // printk("-------------------------------------\n");
    // printk("-------------watch monitor-----------\n");
    // printk("Threshold reached:\n");
    // for (i = 0; i < j; i++) {
    //   kwarg = &k_watch_timer->k_watch_args[buffer[i]];
    //   k_w_arg2threshold(kwarg, &vm_record.threshold_record[i]);
    // }
    // rcu_read_lock();
    
    // diag_variant_buffer_spin_lock(&load_monitor_variant_buffer, flags);
    // diag_variant_buffer_reserve(&load_monitor_variant_buffer,
    //                             sizeof(variable_monitor_record));
    // diag_variant_buffer_write_nolock(&load_monitor_variant_buffer, &vm_record,
    //                                  sizeof(variable_monitor_record));
    // diag_variant_buffer_seal(&load_monitor_variant_buffer);
    // diag_variant_buffer_spin_unlock(&load_monitor_variant_buffer, flags);

    rcu_read_unlock();

    do_each_thread(g, p) {
      if (p->__state == TASK_RUNNING || __task_contributes_to_load(p) ||
          p->__state == TASK_IDLE || 1) {

        get_task_struct(p);

        tsk_info.et_type = 1; //! todo event type
        tsk_info.id = event_id;
        tsk_info.tv = vm_record.tv;

        diag_tsk(p, &tsk_info);
        
        put_task_struct(p);

        push_tsk_info(&tsk_info, &flags);
      }
    }
    while_each_thread(g, p);
    printk("-------------------------------------\n");
  
  return 0;
}