// linmctool - Command-line tool for Bluetooth motion controllers.
// Copyright (c) 2010 pabr@pabr.org
// See http://www.pabr.org/linmctool/
//
// Modifications by Thomas Perl <thp.io/2010/psmove>:
// - Remove WiiMote + Sixaxis-related code
// - Add support for LEDs and Rumble
// - Add support for reading trigger value and buttons
// - Removed unrelated / unnecessary code
//
// Compile with: gcc --std=gnu99 -Wall linmctool-20101117.c -lusb -o linmctool

#include <poll.h>

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <errno.h>
#include <sys/socket.h>
#include <bluetooth/bluetooth.h>
#include <bluetooth/l2cap.h>
#include <assert.h>

/* List of Buttons, based on MoveOnPC's moveonpc.h */
enum {
    L2,
    R2,
    L1,
    R1,
    TRIANGLE,
    CIRCLE,
    CROSS,
    SQUARE,
    SELECT,
    L3,
    R3,
    START,
    UP,
    RIGHT,
    DOWN,
    LEFT,
    PS,
    UNK1,
    UNK2,
    MOVE,
    T,
    UNK3,
    UNK4,
    UNK5,
    BUTTON_COUNT,
};

#define BUTTON_MASK(x) (1<<x)

static const char* BUTTON_NAMES[] = {
    "L2",
    "R2",
    "L1",
    "R1",
    "TRIANGLE",
    "CIRCLE",
    "CROSS",
    "SQUARE",
    "SELECT",
    "L3",
    "R3",
    "START",
    "UP",
    "RIGHT",
    "DOWN",
    "LEFT",
    "PS",
    "UNK1",
    "UNK2",
    "MOVE",
    "T",
    "UNK3",
    "UNK4",
    "UNK5",
};

// Global flags

int verbose = 0;
enum { TEXT, BINARY };
int output_type = TEXT;
char *usb_force_master = NULL;

void fatal(char *msg) {
  if ( errno ) perror(msg); else fprintf(stderr, "%s\n", msg);
  exit(1);
}

// ----------------------------------------------------------------------
// Replacement for libbluetooth

int mystr2ba(const char *s, bdaddr_t *ba) {
  if ( strlen(s) != 17 ) return 1;
  for ( int i=0; i<6; ++i ) {
    int d = strtol(s+15-3*i, NULL, 16);
    if ( d<0 || d>255 ) return 1;
    ba->b[i] = d;
  }
  return 0;
}

char *myba2str(const bdaddr_t *ba) {
  static char buf[2][18];  // Static buffer valid for two invocations.
  static int index = 0;
  index = (index+1)%2;
  sprintf(buf[index], "%02x:%02x:%02x:%02x:%02x:%02x",
	  ba->b[5], ba->b[4], ba->b[3], ba->b[2], ba->b[1], ba->b[0]);
  return buf[index];
}

// ----------------------------------------------------------------------
// USB support

#ifndef WITHOUT_USB

#include <usb.h>
#define USB_DIR_IN 0x80
#define USB_DIR_OUT 0
#define USB_GET_REPORT 0x01
#define USB_SET_REPORT 0x09
#define VENDOR_SONY 0x054c
#define PRODUCT_PSMOVE 0x03d5

void usb_pair_device(struct usb_device *dev, int itfnum) {

  usb_dev_handle *devh = usb_open(dev);
  if ( ! devh ) fatal("usb_open");
  usb_detach_kernel_driver_np(devh, itfnum);
  int res = usb_claim_interface(devh, itfnum);
  if ( res < 0 ) fatal("usb_claim_interface");

  bdaddr_t current_ba;  // Current pairing address.

  switch ( dev->descriptor.idProduct ) {
  case PRODUCT_PSMOVE: {
    fprintf(stderr, "USB: PS MOVE\n");
    unsigned char msg[16];
    int res = usb_control_msg
      (devh, USB_DIR_IN | USB_TYPE_CLASS | USB_RECIP_INTERFACE,
       USB_GET_REPORT, 0x0304, itfnum, (void*)msg, sizeof(msg), 5000);
    if ( res < 0 ) fatal("usb_control_msg(read master)");
    for ( int i=0; i<6; ++i ) current_ba.b[i] = msg[10+i];
    break;
  }
  }

  bdaddr_t ba;  // New pairing address.

  if ( usb_force_master && !mystr2ba(usb_force_master,&ba) )
    ;
  else {
    char ba_s[18];
    FILE *f = popen("hcitool dev", "r");
    if ( !f || fscanf(f, "%*s\n%*s %17s", ba_s)!=1 || mystr2ba(ba_s, &ba) )
      fatal("Unable to retrieve local bd_addr from `hcitool dev`.\n");
    pclose(f);
  }

  // Perform pairing.

  if ( ! bacmp(&current_ba, &ba) ) {
    fprintf(stderr, "  Already paired to %s\n", myba2str(&ba));
  } else {
    fprintf(stderr, "  Changing master from %s to %s\n",
	    myba2str(&current_ba), myba2str(&ba));
    switch ( dev->descriptor.idProduct ) {
    case PRODUCT_PSMOVE: {
      char msg[]= { 0x05, ba.b[0], ba.b[1], ba.b[2], ba.b[3],
		    ba.b[4], ba.b[5], 0x10, 0x01, 0x02, 0x12 };
      res = usb_control_msg
	(devh, USB_DIR_OUT | USB_TYPE_CLASS | USB_RECIP_INTERFACE,
	 USB_SET_REPORT, 0x0305, itfnum, msg, sizeof(msg), 5000);
      if ( res < 0 ) fatal("usb_control_msg(write master)");
      break;
    }
    }
  }

  fprintf(stderr, "  Now press the PS button.\n");
}

void usb_scan() {
  usb_init();
  if ( usb_find_busses() < 0 ) fatal("usb_find_busses");
  if ( usb_find_devices() < 0 ) fatal("usb_find_devices");
  struct usb_bus *busses = usb_get_busses();
  if ( ! busses ) fatal("usb_get_busses");
  
  struct usb_bus *bus;
  for ( bus=busses; bus; bus=bus->next ) {
    struct usb_device *dev;
    for ( dev=bus->devices; dev; dev=dev->next) {
      struct usb_config_descriptor *cfg;
      for ( cfg = dev->config;
	    cfg < dev->config + dev->descriptor.bNumConfigurations;
	    ++cfg ) {
	int itfnum;
	for ( itfnum=0; itfnum<cfg->bNumInterfaces; ++itfnum ) {
	  struct usb_interface *itf = &cfg->interface[itfnum];
	  struct usb_interface_descriptor *alt;
	  for ( alt = itf->altsetting;
		alt < itf->altsetting + itf->num_altsetting;
		++alt ) {
	    if ( dev->descriptor.idVendor == VENDOR_SONY &&
		  dev->descriptor.idProduct == PRODUCT_PSMOVE &&
		 alt->bInterfaceClass == 3 )
	      usb_pair_device(dev, itfnum);
	  }
	}
      }
    }
  }
}

#else // WITHOUT_USB

void usb_scan() { }

#endif

/**********************************************************************/
// Bluetooth HID devices

struct motion_dev {
  int index;
  bdaddr_t addr;
  enum { PSMOVE } type;
  int csk; 
  int isk;
  struct motion_dev *next;
};

#define L2CAP_PSM_HIDP_CTRL 0x11
#define L2CAP_PSM_HIDP_INTR 0x13

#define HIDP_TRANS_GET_REPORT    0x40
#define HIDP_TRANS_SET_REPORT    0x50
#define HIDP_DATA_RTYPE_INPUT    0x01
#define HIDP_DATA_RTYPE_OUTPUT   0x02
#define HIDP_DATA_RTYPE_FEATURE  0x03

// Incoming connections.

int l2cap_listen(const bdaddr_t *bdaddr, unsigned short psm) {
  int sk = socket(PF_BLUETOOTH, SOCK_SEQPACKET, BTPROTO_L2CAP);
  if ( sk < 0 ) fatal("socket");

  struct sockaddr_l2 addr = {
    .l2_family = AF_BLUETOOTH,
    .l2_bdaddr = *BDADDR_ANY,
    .l2_psm = htobs(psm),
    .l2_cid = 0,
  };
  if ( bind(sk, (struct sockaddr *)&addr, sizeof(addr)) < 0 ) {
    perror("bind");
    close(sk);
    return -1;
  }

  if ( listen(sk, 5) < 0 ) fatal("listen");
  return sk;
}

struct motion_dev *accept_device(int csk, int isk) {
  fprintf(stderr, "Incoming connection...\n");
  struct motion_dev *dev = malloc(sizeof(struct motion_dev));
  if ( ! dev ) fatal("malloc");

  dev->csk = accept(csk, NULL, NULL);
  if ( dev->csk < 0 ) fatal("accept(CTRL)");
  dev->isk = accept(isk, NULL, NULL);
  if ( dev->isk < 0 ) fatal("accept(INTR)");

  struct sockaddr_l2 addr;
  socklen_t addrlen = sizeof(addr);
  if ( getpeername(dev->isk, (struct sockaddr *)&addr, &addrlen) < 0 )
    fatal("getpeername");
  dev->addr = addr.l2_bdaddr;
  
  {
    // Distinguish SIXAXIS / DS3 / PSMOVE.
    unsigned char resp[64];
    char get03f2[] = { HIDP_TRANS_GET_REPORT | HIDP_DATA_RTYPE_FEATURE | 8,
		       0xf2, sizeof(resp), sizeof(resp)>>8 };
    send(dev->csk, get03f2, sizeof(get03f2), 0);  // 0301 is interesting too.
    int nr = recv(dev->csk, resp, sizeof(resp), 0);
    assert ( nr < 19 );
  }

  return dev;
}

// Outgoing connections.

int l2cap_connect(bdaddr_t *ba, unsigned short psm) {
  int sk = socket(PF_BLUETOOTH, SOCK_SEQPACKET, BTPROTO_L2CAP);
  if ( sk < 0 ) fatal("socket");

  struct sockaddr_l2 daddr = {
    .l2_family = AF_BLUETOOTH,
    .l2_bdaddr = *ba,
    .l2_psm = htobs(psm),
    .l2_cid = 0,
  };
  if ( connect(sk, (struct sockaddr *)&daddr, sizeof(daddr)) < 0 )
    fatal("connect");

  return sk;
}

struct motion_dev *connect_device(bdaddr_t *ba) {
  fprintf(stderr, "Connecting to %s\n", myba2str(ba));
  struct motion_dev *dev = malloc(sizeof(struct motion_dev));
  if ( ! dev ) fatal("malloc");
  dev->addr = *ba;
  dev->csk = l2cap_connect(ba, L2CAP_PSM_HIDP_CTRL);
  dev->isk = l2cap_connect(ba, L2CAP_PSM_HIDP_INTR);
  return dev;
}

/**********************************************************************/
// Device setup

void hidp_trans(int csk, char *buf, int len) {
  if ( send(csk, buf, len, 0) != len ) fatal("send(CTRL)");
  char ack;
  int nr = recv(csk, &ack, sizeof(ack), 0);
  if ( nr!=1 || ack!=0 ) fatal("ack");
}

void psmove_set_leds(struct motion_dev *dev, char r, char g, char b, char rumble)
{
    char report[] = {
      HIDP_TRANS_SET_REPORT | HIDP_DATA_RTYPE_OUTPUT,
      2, 0, r, g, b, 0, rumble,
    };

    hidp_trans(dev->csk, report, sizeof(report));
}

void setup_device(struct motion_dev *dev) {
  // Rumble when connection has been established
  psmove_set_leds(dev, 0x00, 0x00, 0x00, 0xff);
  psmove_set_leds(dev, 0x00, 0x00, 0x00, 0x00);

  fprintf(stderr, "New device %d %s\n",
	  dev->index, myba2str(&dev->addr));
}

/**********************************************************************/
// Reports

void parse_report_psmove(unsigned char *r, int len, int latest) {
  if ( r[0]==0x01 && len>=49 ) {
    int buttons = r[2] | (r[1] << 8) |
                  ((r[3] & 0x01) << 16) |
                  ((r[4] & 0xF0) << 13);

    if (verbose) {
        for (int i=0; i<BUTTON_COUNT; i++) {
            if (buttons & BUTTON_MASK(i)) {
                printf("%s ", BUTTON_NAMES[i]);
            }
        }
    }

    printf(" btns=0x%06x ", buttons);
    printf(" trig=%-6d ", r[6]);
    printf(" seq=%-2d", (r[4]&15)*2+(latest?1:0));
    int ai = latest ? 19 : 13;
    short aX = r[ai+0] + r[ai+1]*256 - 32768;
    short aY = r[ai+2] + r[ai+3]*256 - 32768;
    short aZ = r[ai+4] + r[ai+5]*256 - 32768;
    printf(" aX=%-6d aY=%-6d aZ=%-6d", aX,aY,aZ);
    int ri = latest ? 31 : 25;
    short gX = r[ri+0] + r[ri+1]*256 - 32768;
    short gY = r[ri+2] + r[ri+3]*256 - 32768;
    short gZ = r[ri+4] + r[ri+5]*256 - 32768;
    printf(" gX=%-6d gY=%-6d gZ=%-6d", gX,gY,gZ);
    short mX = (r[38]<<12) | (r[39]<<4);
    short mY = (r[40]<<8)  | (r[41]&0xf0);
    short mZ = (r[41]<<12) | (r[42]<<4);
    printf(" mX=%-5d mY=%-5d mZ=%-5d", mX>>4,mY>>4,mZ>>4);
  }
  printf("\n");
}

void parse_report(struct motion_dev *dev, unsigned char *r, int len) {
  printf("%d %s ", dev->index, myba2str(&dev->addr));
  parse_report_psmove(r, len, 0);
  printf("%d %s ", dev->index, myba2str(&dev->addr));
  parse_report_psmove(r, len, 1);

  fflush(stdout);
}

/**********************************************************************/
// Main

void usage() {
  fprintf(stderr, "Usage: linmctool [options] [BDADDR...]\n"
	  "  [--verbose]       Print debugging information\n"
	  "  [--binary]        Write raw reports to stdout\n"
	  "  [--dump-readable] Try to read all reports\n"
	  "  [--master BDADDR] Pair USB devices with this adapter\n");
  exit(1);
}

int main(int argc, char *argv[]) {
  struct motion_dev *devs = NULL;
  int next_devindex = 0;

  struct pollfd pfds = { fileno(stdin), POLLIN, 0 };
  char buf[100];
  int r, g, b;

  int csk = l2cap_listen(BDADDR_ANY, L2CAP_PSM_HIDP_CTRL);
  int isk = l2cap_listen(BDADDR_ANY, L2CAP_PSM_HIDP_INTR);

  if ( !(csk>=0 && isk>=0) ) {
      fprintf(stderr, "Unable to listen on HID PSMs (running as root?).");
      exit(1);
  }

  fprintf(stderr, "Waiting for Bluetooth connections.\n");

  for ( int i=1; i<argc; ++i )
    if      ( ! strcmp(argv[i], "--verbose") )       verbose = 1;
    else if ( ! strcmp(argv[i], "--text") )          output_type = TEXT;
    else if ( ! strcmp(argv[i], "--binary") )        output_type = BINARY;
    else if ( ! strcmp(argv[i], "--master") && i+1<argc )
      usb_force_master = argv[++i];
    else {
      bdaddr_t ba;
      if ( mystr2ba(argv[i], &ba) ) usage();
      struct motion_dev *dev = connect_device(&ba);
      dev->index = next_devindex++;
      dev->next = devs;
      devs = dev;
      setup_device(dev);
    }

  usb_scan();

  while ( !feof(stdin) ) {
    fd_set fds; FD_ZERO(&fds);
    int fdmax = 0;
    if ( csk >= 0 ) FD_SET(csk, &fds);
    if ( isk >= 0 ) FD_SET(isk, &fds);
    if ( csk > fdmax ) fdmax = csk;
    if ( isk > fdmax ) fdmax = isk;
    for ( struct motion_dev *dev=devs; dev; dev=dev->next ) {
      FD_SET(dev->csk, &fds); if ( dev->csk > fdmax ) fdmax = dev->csk;
      FD_SET(dev->isk, &fds); if ( dev->isk > fdmax ) fdmax = dev->isk;
    }
    if ( select(fdmax+1,&fds,NULL,NULL,NULL) < 0 ) fatal("select");
    // Incoming connection ?
    if ( csk>=0 && FD_ISSET(csk,&fds) ) {
      struct motion_dev *dev = accept_device(csk, isk);
      dev->index = next_devindex++;
      dev->next = devs;
      devs = dev;
      setup_device(dev);
    }

    if (poll(&pfds, 1, 0)) {
        while (poll(&pfds, 1, 0)) {
            if (fgets(buf, sizeof(buf), stdin) == NULL) { break; }
            sscanf(buf, "%d %d %d\n", &r, &g, &b);
        }
        // Randomly light up the PS Move Bulb when receiving a report
        psmove_set_leds(devs, r, g, b, 0x00);
    }

    // Incoming input report ?
    for ( struct motion_dev *dev=devs; dev; dev=dev->next )
      if ( FD_ISSET(dev->isk, &fds) ) {
	unsigned char report[256];
	int nr = recv(dev->isk, report, sizeof(report), 0);
	if ( nr <= 0 ) {
	  fprintf(stderr, "%d disconnected\n", dev->index);
	  close(dev->csk); close(dev->isk);
	  struct motion_dev **pdev;
	  for ( pdev=&devs; *pdev!=dev; pdev=&(*pdev)->next ) ;
	  *pdev = dev->next;
	  free(dev);
	} else {
	  if ( verbose ) {
	    fprintf(stderr, "[%d]", nr);
	    for ( int i=0; i<nr; ++i ) fprintf(stderr, " %02x", report[i]);
	    fprintf(stderr, "\n");
	  }
	  if ( report[0] == 0xa1 ) {
              switch (output_type) {
                  case BINARY:
                      write(1, report, nr);
                      break;
                  default:
                  case TEXT:
                      parse_report(dev, report+1, nr-1);
                      break;
              }
	  }
	}
      }
  }

  return 0;
}