#include <pthread.h>
#include <usb.h>

#include <csignal>
#include <cstdio>
#include <cstring>
#include <format>
#include <iostream>
#include <sstream>
#include <stdexcept>
#include <string>
#include <vector>

class USBDevice;
class EndpointListenerThread;

void print_raw_data(std::ostream& out, uint8_t* data, int len);
bool global_interrupt = false;

struct usb_device* find_usb_device(uint16_t idVendor, uint16_t idProduct) {
  struct usb_bus* busses = usb_get_busses();

  for (struct usb_bus* bus = busses; bus; bus = bus->next) {
    for (struct usb_device* dev = bus->devices; dev; dev = dev->next) {
      if (dev->descriptor.idVendor == idVendor &&
          dev->descriptor.idProduct == idProduct) {
        return dev;
      }
    }
  }
  return 0;
}

class USBDevice {
 private:
  static USBDevice* current_;

 public:
  static USBDevice* current() { return current_; }

 private:
  struct usb_device* dev;
  struct usb_dev_handle* handle;

  std::vector<EndpointListenerThread*> threads;

 public:
  USBDevice(struct usb_device* dev_) : dev(dev_) {
    USBDevice::current_ = this;

    handle = usb_open(dev);
    if (!handle) {
      throw std::runtime_error("Error opening usb device");
    }
  }

  ~USBDevice() { usb_close(handle); }

  void clear_halt(int ep) {
    if (usb_clear_halt(handle, ep) != 0) {
      std::cout << "Failure to reset_ep: " << ep << std::endl;
    }
  }

  void reset() {
    if (usb_reset(handle) != 0) {
      std::cout << "Failure to reset" << std::endl;
    }
  }

  void detach_kernel_driver(int iface) {
    if (usb_detach_kernel_driver_np(handle, iface) < 0) {
      std::ostringstream str;
      str << "Couldn't detcach interface " << iface;
      throw std::runtime_error(str.str());
    }
  }

  void claim_interface(int iface) {
    if (usb_claim_interface(handle, iface) != 0) {
      std::ostringstream str;
      str << "Couldn't claim interface " << iface;
      throw std::runtime_error(str.str());
    }
  }

  void release_interface(int iface) {
    if (usb_release_interface(handle, iface) != 0) {
      std::ostringstream str;
      str << "Couldn't release interface " << iface;
      throw std::runtime_error(str.str());
    }
  }

  void set_configuration(int configuration) {
    if (usb_set_configuration(handle, configuration) != 0) {
      std::ostringstream str;
      str << "Couldn't set configuration " << configuration;
      throw std::runtime_error(str.str());
    }
  }

  void set_altinterface(int interface) {
    if (usb_set_altinterface(handle, interface) != 0) {
      std::ostringstream str;
      str << "Couldn't set alternative interface " << interface;
      throw std::runtime_error(str.str());
    }
  }

  int read(int endpoint, uint8_t* data, int len) {
    return usb_interrupt_read(handle, endpoint, (char*)data, len, 0);
  }

  int write(int endpoint, uint8_t* data, int len) {
    return usb_interrupt_write(handle, endpoint, (char*)data, len, 0);
  }

  /* uint8_t  requesttype
     uint8_t  request
     uint16_t value;
     uint16_t index;
     uint16_t length;
   */
  int ctrl_msg(int requesttype, int request, int value, int index,
               uint8_t* data, int size) {
    return usb_control_msg(handle, requesttype, request, value, index,
                           (char*)data, size, 0 /* timeout */);
  }

  void print_info() {
    for (int i = 0; i < dev->descriptor.bNumConfigurations; ++i) {
      std::cout << "Configuration: " << i << std::endl;
      for (int j = 0; j < dev->config[i].bNumInterfaces; ++j) {
        std::cout << "  Interface " << j << ":" << std::endl;
        for (int k = 0; k < dev->config[i].interface[j].num_altsetting; ++k) {
          for (int l = 0;
               l < dev->config[i].interface[j].altsetting[k].bNumEndpoints;
               ++l) {
            std::cout << "    Endpoint: "
                      << int(dev->config[i]
                                 .interface[j]
                                 .altsetting[k]
                                 .endpoint[l]
                                 .bEndpointAddress &
                             USB_ENDPOINT_ADDRESS_MASK)
                      << ((dev->config[i]
                               .interface[j]
                               .altsetting[k]
                               .endpoint[l]
                               .bEndpointAddress &
                           USB_ENDPOINT_DIR_MASK)
                              ? " (IN)"
                              : " (OUT)")
                      << std::endl;
          }
        }
      }
    }
  }

  void launch_listener_thread(int endpoint);
};

USBDevice* USBDevice::current_ = 0;

class EndpointListenerThread {
 private:
  pthread_t thread;
  int endpoint;
  bool running;

 public:
  EndpointListenerThread(int endpoint_) : endpoint(endpoint_), running(false) {}

  ~EndpointListenerThread() {}

  void start() {
    running = true;
    pthread_create(&thread, NULL, &EndpointListenerThread::callback, this);
  }

  static void* callback(void* userdata) {
    EndpointListenerThread* thread =
        static_cast<EndpointListenerThread*>(userdata);
    thread->run();
    pthread_exit(NULL);
  }

  void run() {
    bool this_quit = false;
    std::cout << "Reading from endpoint " << endpoint << std::endl;
    ;
    while (!this_quit) {
      uint8_t data[8192];
      int ret = USBDevice::current()->read(endpoint, data, sizeof(data));
      if (ret < 0) {
        std::cerr << "USBError: " << ret << "\n" << usb_strerror() << std::endl;
        std::cerr << "Shutting down" << std::endl;
        this_quit = true;
      } else {
        std::cout << ">>> Ep" << endpoint << ": ";
        print_raw_data(std::cout, data, ret);
        std::cout << std::endl;
      }
    }
  }
};

void USBDevice::launch_listener_thread(int endpoint) {
  try {
    EndpointListenerThread* thread = new EndpointListenerThread(endpoint);
    threads.push_back(thread);
    thread->start();
  } catch (std::exception& err) {
    std::cout << "Error: " << err.what() << std::endl;
  }
}

bool has_prefix(const std::string& lhs, const std::string rhs) {
  if (lhs.length() < rhs.length()) {
    return false;
  } else {
    return lhs.compare(0, rhs.length(), rhs) == 0;
  }
}

bool readline(const std::string& prompt, std::string& line) {
  std::cout << prompt << std::flush;
  bool ret = !std::getline(std::cin, line).eof();
  // line = strip(line);
  return ret;
}

std::vector<std::string> tokenize(const std::string& str,
                                  const std::string& delimiters = " ") {
  std::vector<std::string> tokens;

  // Skip delimiters at beginning.
  std::string::size_type lastPos = str.find_first_not_of(delimiters, 0);
  // Find first "non-delimiter".
  std::string::size_type pos = str.find_first_of(delimiters, lastPos);

  while (std::string::npos != pos || std::string::npos != lastPos) {
    // Found a token, add it to the vector.
    tokens.push_back(str.substr(lastPos, pos - lastPos));
    // Skip delimiters.  Note the "not_of"
    lastPos = str.find_first_not_of(delimiters, pos);
    // Find next "non-delimiter"
    pos = str.find_first_of(delimiters, lastPos);
  }

  return tokens;
}

void print_raw_data(std::ostream& out, uint8_t* data, int len) {
  std::cout << "[" << len << "] { ";

  for (int i = 0; i < len; ++i) {
    std::cout << std::format("{:#02x}", int(data[i]));
    if (i != len - 1) {
      std::cout << ", ";
    }
  }

  std::cout << " }";
}

void console_listen_cmd(const std::vector<std::string>& args) {
  if (args.size() < 2) {
    std::cout << "Usage: listen [ENDPOINT]..." << std::endl;
  } else {
    for (size_t i = 1; i < args.size(); ++i) {
      int endpoint = atoi(args[i].c_str());
      USBDevice::current()->launch_listener_thread(endpoint);
    }
  }
}

void console_claim_cmd(const std::vector<std::string>& args) {
  if (args.size() < 2) {
    std::cout << "Usage: claim [INTERFACE]..." << std::endl;
  } else {
    for (size_t i = 1; i < args.size(); ++i) {
      int interface = atoi(args[i].c_str());
      try {
        USBDevice::current()->claim_interface(interface);
      } catch (std::exception& err) {
        std::cout << "Error: " << err.what() << std::endl;
        goto end;
      }
      std::cout << "Interface " << interface << " successfully claimed"
                << std::endl;
    end:;
    }
  }
}

void console_release_cmd(const std::vector<std::string>& args) {
  if (args.size() < 2) {
    std::cout << "Usage: release [INTERFACE]..." << std::endl;
  } else {
    for (size_t i = 1; i < args.size(); ++i) {
      int interface = atoi(args[i].c_str());
      try {
        USBDevice::current()->release_interface(interface);
      } catch (std::exception& err) {
        std::cout << "Error: " << err.what() << std::endl;
        goto end;
      }
      std::cout << "Interface " << interface << " successfully releaseed"
                << std::endl;
    end:;
    }
  }
}

void console_detach_cmd(const std::vector<std::string>& args) {
  if (args.size() < 2) {
    std::cout << "Usage: detach [INTERFACE]..." << std::endl;
  } else {
    for (size_t i = 1; i < args.size(); ++i) {
      int interface = atoi(args[i].c_str());
      try {
        USBDevice::current()->detach_kernel_driver(interface);
      } catch (std::exception& err) {
        std::cout << "Error: " << err.what() << std::endl;
        goto end;
      }
      std::cout << "Interface " << interface << " successfully detached"
                << std::endl;
    end:;
    }
  }
}

void console_info_cmd(const std::vector<std::string>& args) {
  USBDevice::current()->print_info();
}

class Sequence {
 private:
  int start;
  int end;
  int idx;

 public:
  Sequence(int start_, int end_) : start(start_), end(end_), idx(start_) {}

  Sequence(int start_) : start(start_), end(start_), idx(start_) {}

  int get() {
    if (eol()) {
      return end;
    } else {
      return idx;
    }
  }

  void next() {
    if (!eol()) {
      if (start <= end) {
        idx += 1;
      } else {
        idx -= 1;
      }
    }
  }

  bool eol() {
    if (start <= end) {
      return idx > end;
    } else {
      return idx < end;
    }
  }

  void reset() { idx = start; }

  std::string to_string() const {
    std::ostringstream str;
    if (start == end) {
      str << start;
    } else {
      str << start << "-" << end;
    }
    return str.str();
  }
};

class SequenceGenerator {
 private:
  std::vector<Sequence> sequences;
  std::vector<Sequence>::size_type idx;

 public:
  /** spec: NUM or [NUM, NUM, NUM-NUM]*/
  SequenceGenerator(const std::string& spec) {
    std::vector<std::string> tokens = tokenize(spec, ",");

    for (std::vector<std::string>::iterator i = tokens.begin();
         i != tokens.end(); ++i) {
      int start, end;
      if (sscanf(i->c_str(), "%x-%x", &start, &end) == 2) {
        sequences.push_back(Sequence(start, end));
      } else if (sscanf(i->c_str(), "%x", &start) == 1) {
        sequences.push_back(Sequence(start));
      } else {
        throw std::runtime_error(
            "Syntax Error in sequence, couldn't convert: " + *i);
      }
    }

    idx = 0;
  }

  bool eol() { return (idx == sequences.size()); }

  void reset() {
    for (std::vector<Sequence>::iterator i = sequences.begin();
         i != sequences.end(); ++i) {
      i->reset();
    }
    idx = 0;
  }

  int get() {
    if (eol()) {
      return sequences.back().get();
    } else {
      return sequences[idx].get();
    }
  }

  void next() {
    if (!eol()) {
      if (sequences[idx].eol()) {
        ++idx;
        if (idx < sequences.size()) {
          sequences[idx].next();
        }
      } else {
        return sequences[idx].next();
      }
    }
  }

  std::string to_string() const {
    std::ostringstream str;
    std::vector<Sequence>::const_iterator i = sequences.begin();

    str << "{ ";
    while (i != sequences.end()) {
      const Sequence& seq = *i;

      str << seq.to_string();

      ++i;
      if (i != sequences.end()) {
        str << ", ";
      }
    }
    str << " }";

    return str.str();
  }
};

bool eol(std::vector<SequenceGenerator>& sequences) {
  if (sequences.empty()) {
    return true;
  } else {
    return sequences.back().eol();
  }
}

void next(std::vector<SequenceGenerator>& sequences, int idx) {
  if (idx < int(sequences.size())) {
    if (!sequences[idx].eol()) {
      sequences[idx].next();
      if (sequences[idx].eol()) {
        if (idx + 1 < int(sequences.size())) {
          sequences[idx].reset();
          next(sequences, idx + 1);
        }
      }
    }
  }
}

void next(std::vector<SequenceGenerator>& sequences) {
  if (!sequences.empty()) {
    next(sequences, 0);
  }
}

void console_probe_cmd(const std::vector<std::string>& args) {
  if (args.size() < 3) {
    std::cout << "Usage: send [EP] [DATA]..." << std::endl;
  } else {
    std::vector<uint8_t> data;
    std::vector<SequenceGenerator> sequences;

    int endpoint = atoi(args[1].c_str());
    for (int i = 2; i < int(args.size()); ++i) {
      SequenceGenerator gen(args[i]);
      sequences.push_back(gen);
      std::cout << gen.to_string() << std::endl;
    }

    data.resize(sequences.size());
    while (!eol(sequences) && !global_interrupt) {
      for (int i = 0; i < int(sequences.size()); ++i) {
        data[i] = sequences[i].get();
      }

      std::cout << "Data Ep: " << endpoint << ": ";
      print_raw_data(std::cout, &*data.begin(), data.size());
      int ret =
          USBDevice::current()->write(endpoint, &*data.begin(), data.size());
      std::cout << " -> " << ret;
      std::cout << std::endl;

      next(sequences);

      // usleep(100 * 1000);
      usleep(10 * 1000);
    }
  }
}

void console_wait_cmd(const std::vector<std::string>& args) {
  if (args.size() != 2) {
    std::cout << "Usage: wait [SECONDS]" << std::endl;
  } else {
    float seconds = atof(args[1].c_str());
    std::cout << "Waiting for " << seconds << " seconds... " << std::flush;
    usleep(static_cast<useconds_t>(1000 * 1000 * seconds));
    std::cout << "done" << std::endl;
  }
}

void console_send_cmd(const std::vector<std::string>& args) {
  if (args.size() < 3) {
    std::cout << "Usage: send [EP] [DATA]..." << std::endl;
  } else {
    int endpoint = atoi(args[1].c_str());
    std::vector<uint8_t> data;
    for (int i = 2; i < int(args.size()); ++i) {
      data.push_back(strtol(args[i].c_str(), NULL, 16));
    }

    std::cout << "Sending to endpoint " << endpoint << ": ";
    print_raw_data(std::cout, &*data.begin(), data.size());
    int ret =
        USBDevice::current()->write(endpoint, &*data.begin(), data.size());
    std::cout << " -> " << ret << std::endl;
  }
}

void console_ctrlreq_cmd(const std::vector<std::string>& args) {
  if (args.size() < 5) {
    std::cout << "Usage: ctrlreq REQUESTTYPE REQUEST VALUE INDEX LEN"
              << std::endl;
  } else {  // See USB Specification (usb_20.pdf) page 248
    int requesttype = strtol(args[1].c_str(), NULL, 16);
    int request = strtol(args[2].c_str(), NULL, 16);
    int value = strtol(args[3].c_str(), NULL, 16);
    int index = strtol(args[4].c_str(), NULL, 16);
    int len = strtol(args[5].c_str(), NULL, 16);

    std::cout << "Sending to ctrl: " << requesttype << " " << request << " "
              << value << " " << index << " " << len << ": ";

    uint8_t* data = (len == 0) ? NULL : new uint8_t[len];
    memset(data, 0, len * sizeof(uint8_t));

    int ret = USBDevice::current()->ctrl_msg(requesttype, request, value, index,
                                             data, len);
    std::cout << " -> " << ret << " '" << strerror(-ret) << "'" << std::endl;

    if (!data) {
      std::cout << "no data";
    } else {
      print_raw_data(std::cout, data, len);
    }

    std::cout << std::endl;

    delete[] data;
  }
}

void console_setconfiguration_cmd(const std::vector<std::string>& args) {
  if (args.size() != 2) {
    std::cout << "Usage: setconfiguration CFG" << std::endl;
  } else {
    USBDevice::current()->set_configuration(atoi(args[1].c_str()));
  }
}

void console_setaltinterface_cmd(const std::vector<std::string>& args) {
  if (args.size() != 2) {
    std::cout << "Usage: " << args[0] << " INTERFACE" << std::endl;
  } else {
    USBDevice::current()->set_altinterface(atoi(args[1].c_str()));
  }
}

void console_reset_ep_cmd(const std::vector<std::string>& args) {
  if (args.size() != 2) {
    std::cout << "Usage: clear_halt ENDPOINT" << std::endl;
  } else {
    USBDevice::current()->clear_halt(atoi(args[1].c_str()));
  }
}

void console_reset_cmd(const std::vector<std::string>& args) {
  USBDevice::current()->reset();
}

void console_ctrl_cmd(const std::vector<std::string>& args) {
  if (args.size() < 5) {
    std::cout << "Usage: ctrl REQUESTTYPE REQUEST VALUE INDEX [DATA]..."
              << std::endl;
  } else {  // See USB Specification (usb_20.pdf) page 248
    int requesttype = strtol(args[1].c_str(), NULL, 16);
    int request = strtol(args[2].c_str(), NULL, 16);
    int value = strtol(args[3].c_str(), NULL, 16);
    int index = strtol(args[4].c_str(), NULL, 16);

    std::vector<uint8_t> data;
    for (int i = 5; i < int(args.size()); ++i) {
      data.push_back(strtol(args[i].c_str(), NULL, 16));
    }

    std::cout << "Sending to ctrl: " << requesttype << " " << request << " "
              << value << " " << index << ": ";
    if (data.empty()) {
      std::cout << "no data";
    } else {
      print_raw_data(std::cout, &*data.begin(), data.size());
    }

    int ret = USBDevice::current()->ctrl_msg(
        requesttype, request, value, index,
        data.empty() ? NULL : &*data.begin(), data.size());
    std::cout << " -> " << ret << " '" << strerror(-ret) << "'" << std::endl;
  }
}

std::string strip_comment(const std::string& line) {
  std::string::size_type p = line.find_first_of('#');
  if (p != std::string::npos) {
    return line.substr(0, p);
  } else {
    return line;
  }
}

bool quit = false;

void dispatch_command(const std::vector<std::string>& args) {
  if (0) {
    for (int i = 0; i < int(args.size()); ++i) {
      std::cout << i << ":'" << args[i] << "' ";
    }
    std::cout << std::endl;
  }

  if (args[0] == "help") {
    std::cout << "help:\n   Print this help\n" << std::endl;
    std::cout << "quit:\n   Exit usbdebug\n" << std::endl;
    std::cout << "claim [INTERFACE]...\n   Claim the given interfaces\n"
              << std::endl;
    std::cout << "release [INTERFACE]...\n   Release the given interfaces\n"
              << std::endl;
    std::cout
        << "detach [INTERFACE]...\n   Detach kernel driver from interfaces\n"
        << std::endl;
    std::cout << "listen [ENDPOINT]...\n   On the given endpoints\n"
              << std::endl;
    std::cout << "info\n   Print some info on the current device\n"
              << std::endl;
    std::cout << "send [ENDPOINT] [DATA]...\n   Send data to an USB Endpoint\n"
              << std::endl;
  } else if (args[0] == "listen") {
    console_listen_cmd(args);
  } else if (args[0] == "info") {
    console_info_cmd(args);
  } else if (args[0] == "claim") {
    console_claim_cmd(args);
  } else if (args[0] == "release") {
    console_release_cmd(args);
  } else if (args[0] == "probe") {
    console_probe_cmd(args);
  } else if (args[0] == "detach") {
    console_detach_cmd(args);
  } else if (args[0] == "ctrl") {
    console_ctrl_cmd(args);
  } else if (args[0] == "ctrlreq") {
    console_ctrlreq_cmd(args);
  } else if (args[0] == "wait") {
    console_wait_cmd(args);
  } else if (args[0] == "send") {
    console_send_cmd(args);
  } else if (args[0] == "reset") {
    console_reset_cmd(args);
  } else if (args[0] == "clear_halt") {
    console_reset_ep_cmd(args);
  } else if (args[0] == "setconfiguration") {
    console_setconfiguration_cmd(args);
  } else if (args[0] == "setaltinterface") {
    console_setaltinterface_cmd(args);
  } else if (args[0] == "quit") {
    std::cout << "Exiting" << std::endl;
    quit = true;
  } else {
    std::cout << "Unknown command '" << args[0]
              << "', type 'help' to list all available commands" << std::endl;
  }
}

void eval_console_cmd(const std::string& line_) {
  std::string line = strip_comment(line_);

  std::vector<std::string> lines = tokenize(line, ";");
  for (std::vector<std::string>::iterator i = lines.begin(); i != lines.end();
       ++i) {
    std::vector<std::string> args = tokenize(*i);
    if (!args.empty()) {
      dispatch_command(args);
    }
  }
}

void run_console() {
  std::cout << "Type 'help' to list all available commands" << std::endl;

  std::string line;

  while (!quit && readline("\033[32musb\033[0m> ", line)) {
    std::vector<std::string> cmds = tokenize(line, ";");
    for (std::vector<std::string>::iterator i = cmds.begin(); i != cmds.end();
         ++i) {
      eval_console_cmd(*i);
    }
    global_interrupt = false;
  }
  std::cout << std::endl;
}

void signal_callback(int) {
  if (global_interrupt) {
    exit(EXIT_FAILURE);
  } else {
    global_interrupt = true;
    std::cout << "INTERRUPT" << std::endl;
  }
}

int main(int argc, char** argv) {
  if (argc != 2) {
    std::cout << "Usage: " << argv[0] << " [idVendor:idProduct]" << std::endl;
  } else {
    usb_init();
    usb_find_busses();
    usb_find_devices();

    uint16_t idVendor;
    uint16_t idProduct;
    if (sscanf(argv[1], "%hx:%hx", &idVendor, &idProduct) == 2) {
      struct usb_device* dev = find_usb_device(idVendor, idProduct);

      if (dev) {
        std::cout
            << std::format(
                   "Opening device with idVendor: {:#04x}, idProduct: {:#04x}",
                   idVendor, idProduct)
            << std::endl;
        USBDevice* usbdev = new USBDevice(dev);
        signal(SIGINT, signal_callback);
        run_console();
        delete usbdev;
      } else {
        std::cout << "Couldn't device with " << argv[1] << std::endl;
      }
    } else {
      std::cout << "Syntax error, must be idVendor:idProduct (ex: 045e:028e)"
                << std::endl;
    }
  }

  return 0;
}

/* EOF */
