#pragma once

#include <vector>
#include <atomic>
#include <thread>

#include "shm_helper.h"


class ShmWorker {
 public:
  ShmWorker();
  ~ShmWorker();

  bool register_worker(int32_t tgid, int gpu_id, uint64_t *out_shareable_handle,
                       uint64_t *out_vmem_size);

  bool try_lock_gpu(bool &out_self_hold);
  bool lock_gpu(bool &out_self_hold);
  void unlock_gpu(bool keep_wait = false);

  bool has_higher_priority_waiter();
  void start_wait();
  void cancel_wait();

 private:
  int32_t tgid;
  int gpu_id;
  int shm_slot;
  uint16_t priority;
  uint32_t waiting_timestamp;
  bool is_waiting;
  bool is_holding_lock;

  ShmHelper *shm_helper;
  std::thread heart_beat_thread;
  std::atomic<bool> stop_heart_beat;

  // request
  uint64_t make_request(uint32_t type, uint64_t parameter);
  int register_worker_shm();

  // heart beat
  void heart_beat_loop();
};