#include <linux/hrtimer.h>
#include <linux/kprobes.h>
#include <linux/ktime.h>
#include <linux/list.h>
#include <linux/slab.h> /*  for kmalloc */
#include <linux/string.h>

#include <asm/uaccess.h>
#include <linux/cdev.h>
#include <linux/highmem.h>
#include <linux/sched.h>
#include <linux/sched/loadavg.h> /* for avenrun, LOAD_* */
#include <linux/sched/mm.h>
#include <linux/sched/signal.h>
#include <linux/stacktrace.h> /* for stack_trace_print */

#define MAX_TIMER_NUM (128)      // max timer number
#define TIMER_MAX_WATCH_NUM (32) // A timer max watch number at once time
#define MAX_NAME_LEN (15)        // max name length
typedef struct {
  pid_t task_id;               // current process id
  char name[MAX_NAME_LEN + 1]; // name
  void *ptr;                   // virtual address
  int length_byte;             // byte
  long long threshold;         // threshold value
  unsigned char unsigned_flag; // unsigned flag (true: unsigned, false: signed)
  unsigned char greater_flag;  // reverse flag (true: >, false: <)
  unsigned long time_ns;       // timer interval (ns)
} watch_arg;

typedef struct {
  pid_t task_id;               // current process id
  char name[MAX_NAME_LEN + 2]; // name, last char automatically add '\0'
  void *kptr;                  // kernel address + offset
  int length_byte;             // byte
  long long threshold;         // threshold value
  unsigned char unsigned_flag; // unsigned flag (true: unsigned, false: signed)
  unsigned char greater_flag;  // reverse flag (true: >, false: <)
} kernel_watch_arg;

typedef struct {
  unsigned long long time_ns; // hrTimer time interval (ns)
  struct hrtimer hr_timer;    // hrTimer
  ktime_t kt;                 // hrTimer time
  unsigned sentinel;          // sentinel
  kernel_watch_arg
      k_watch_args[TIMER_MAX_WATCH_NUM]; // all watched kernel_watch_arg
} kernel_watch_timer;

#define TIMER_FILLED(timer) ((timer)->sentinel >= TIMER_MAX_WATCH_NUM)
#define TIMER_EMPTY(timer) (!((timer)->time_ns | (timer)->sentinel))
#define TIMER_NO_KWARG(timer) ((timer)->sentinel == 0)

#define TIMER_START(timer)                                                     \
  (hrtimer_start(&timer->hr_timer, timer->kt, HRTIMER_MODE_REL))
#define TIMER_CANCEL(timer) (hrtimer_cancel(&timer->hr_timer))

kernel_watch_timer kernel_wtimer_list[MAX_TIMER_NUM] = {
    0};                    // all kernel_watch_timer
int kernel_wtimer_num = 0; // current kernel_watch_timer number

EXPORT_SYMBOL(kernel_wtimer_list); // export kernel_watch_timer_list
EXPORT_SYMBOL(kernel_wtimer_num);  // export kernel_watch_timer_num

// Helper function
unsigned char w_arg2k_w_arg(void *ptr, watch_arg warg,
                            kernel_watch_arg *k_watch_arg);

// for timer
kernel_watch_timer *get_timer(unsigned long long time_ns);
unsigned char timer_add_watch(kernel_watch_timer *timer,
                              kernel_watch_arg k_watch_arg);
unsigned char timer_del_watch_by_pid(kernel_watch_timer *timer, pid_t pid);

// for memory access
typedef struct {
  pid_t task_id; // current process id
  struct page *page;
  void *kaddr;
  struct list_head entry;
} watch_local_memory;

static LIST_HEAD(watch_local_memory_list);

void free_page_list(pid_t task_id);
void free_all_page_list(void);

// static struct page *page = NULL;
// static void *kaddr = NULL;

void *convert_user_space_ptr(pid_t pid, unsigned long kaddr);

// for timer
// #define US2NS (1000) // Interval in microseconds
// static struct hrtimer hr_timer;
// static ktime_t kt;

// hrTimer
enum hrtimer_restart check_variable_cb(struct hrtimer *timer);
void start_all_hrTimer(void);
void cancel_all_hrTimer(void);

unsigned char read_and_compare(kernel_watch_arg *k_arg);

// for diag_kallsyms_lookup_name
unsigned long (*diag_kallsyms_lookup_name)(const char *name);
static struct kprobe kprobe_kallsyms_lookup_name = {.symbol_name =
                                                        "kallsyms_lookup_name"};

int fn_kallsyms_lookup_name_init(void); // init kallsyms_lookup_name

// form
// https://github.com/alibaba/diagnose-tools/blob/8cd905a1c17f2201e460a2d607413a1303757a32/SOURCE/module/internal.h#L65
// look for current function address, all the function with prefix "orig_" are
#define LOOKUP_SYMS(name)                                                      \
  do {                                                                         \
    orig_##name = (void *)diag_kallsyms_lookup_name(#name);                    \
    if (!orig_##name) {                                                        \
      printk(KERN_ERR "kallsyms_lookup_name: %s\n", #name);                    \
      return -EINVAL;                                                          \
    }                                                                          \
  } while (0)

#define LOOKUP_SYMS_NORET(name)                                                \
  do {                                                                         \
    orig_##name = (void *)diag_kallsyms_lookup_name(#name);                    \
    if (!orig_##name)                                                          \
      pr_err("kallsyms_lookup_name: %s\n", #name);                             \
  } while (0)

#define BACKTRACE_DEPTH 20 // max stack depth

// LOOKUP_SYMS(stack_trace_save_tsk);
unsigned int (*orig_stack_trace_save_tsk)(struct task_struct *task,
                                          unsigned long *store,
                                          unsigned int size,
                                          unsigned int skipnr);
// LOOKUP_SYMS(show_stack);
void (*orig_show_stack)(struct task_struct *task, unsigned long *sp,
                        const char *loglvl);

// https://www.spinics.net/lists/kernel/msg3582022.html
// remove from 5.8.rc3,but it still work
// whether the task contributes to the load
#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 print all task stack
/// @param
static void print_task_stack(void) {
  struct task_struct *g, *p;                // g: task group; p: task
  unsigned long backtrace[BACKTRACE_DEPTH]; // save stack
  unsigned int nr_bt;                       // stack depth
  unsigned long long current_time;          // last time
  current_time = ktime_get_real();
  printk("Timestamp (ns): %lld\n", current_time);
  printk("Recent Load: %lu.%02lu, %lu.%02lu, %lu.%02lu\n", // recent load
         LOAD_INT(avenrun[0]), LOAD_FRAC(avenrun[0]), LOAD_INT(avenrun[1]),
         LOAD_FRAC(avenrun[1]), LOAD_INT(avenrun[2]), LOAD_FRAC(avenrun[2]));
  rcu_read_lock(); // lock run queue
  // printk("Running task\n");
  do_each_thread(g, p) {
    if (p->__state == TASK_RUNNING || __task_contributes_to_load(p) ||
        p->__state == TASK_IDLE) {
      printk("task: %s, pid %d, state %d\n", p->comm, p->pid,
             p->__state); //! todo
      nr_bt = orig_stack_trace_save_tsk(p, backtrace, BACKTRACE_DEPTH, 0);
      stack_trace_print(backtrace, nr_bt, 0); // print
    }
  }
  while_each_thread(g, p);
  rcu_read_unlock(); // unlock run queue
}

unsigned char del_all_kwarg_by_pid(pid_t pid);
void clear_watch(pid_t pid);
void clear_all_watch(void);

// one more thing
int monitor_init(void);