#include <errno.h>
#include <fcntl.h>
#include <getopt.h>
#include <limits.h>
#include <libgen.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <pthread.h>
#include <signal.h>
#include <sys/epoll.h>
#include <sys/ioctl.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/stat.h>

#define AEP_RESET   0xFF
#define AEP_MAGIC   0xFF
#define AEP_VERSION 0x01
#define AEP_VENDOR  0x03
#define AEP_RATE    0x05
#define AEP_CONFIG  0x07
#define AEP_START   0x09
#define AEP_STOP    0x0B
#define AEP_ACK     0xAC

#define AEP_POWER        0
#define AEP_VOLTAGE      1
#define AEP_CURRENT      2
#define AEP_CHANNEL_SIZE 3
#define AEP_NR_CHANNELS  3

#define AEP_SAMPLING_RATE 10000
#define AEP_MAX_PROBES 10

struct aep_frame {
        unsigned short frame;
	unsigned long timestamp;
        float channel[AEP_NR_CHANNELS][AEP_CHANNEL_SIZE];
};

struct aep_channel_stat {
        float savg[AEP_CHANNEL_SIZE];
        float avg[AEP_CHANNEL_SIZE];
        float min[AEP_CHANNEL_SIZE];
        float max[AEP_CHANNEL_SIZE];
        float energy[AEP_CHANNEL_SIZE];
};

struct aep_channel {
        float rshunt;
        struct aep_channel_stat stat;
};

struct aep_device {
	int fd;
	int index;
	const char *name;
	unsigned long total_frame;
	unsigned long frame_missed;
	unsigned short next_frame;
	struct aep_channel channel[AEP_NR_CHANNELS];
};

static struct option long_options[] = {
        { "average",   0, 0, 'a' },
	{ "raw",       0, 0, 'r' },
	{ "energy",    0, 0, 'e' },
	{ "samples",   0, 0, 's' },
	{ "time",      0, 0, 't' },
	{ "limit",     0, 0, 'l' },
	{ "probe",     0, 0, 'p' },
	{ "timestamp", 0, 0, 'T' },
        { "version",   0, 0, 'V' },
        { "help",      0, 0, 'h' },
        { 0, 0, 0, 0 }
};

struct aep_options {
	bool          average;
	bool          raw;
	bool          energy;
	unsigned long timestamp;
	unsigned long samples;
	unsigned long time;
	unsigned long limit;
	unsigned int  rshunt0[AEP_MAX_PROBES];
	unsigned int  rshunt1[AEP_MAX_PROBES];
	unsigned int  rshunt2[AEP_MAX_PROBES];
	const char   *probe[AEP_MAX_PROBES];
	const char   *name[AEP_MAX_PROBES];
	int nrprobe;
};

static int aep_getopt(int argc, char *argv[], struct aep_options *options)
{
        int c, i;
	struct timeval t;
	char *aux, *token;

        memset(options, 0, sizeof(*options));
 
        while (1) {
                int optindex = 0;

                c = getopt_long(argc, argv, "aers:t:l:p:TVh",
                                long_options, &optindex);
                if (c == -1)
                        break;

                switch (c) {
                case 'a':
                        options->average = true;
                        break;
                case 'r':
                        options->raw = true;
                        break;
		case 'e':
			options->energy = true;
			break;
                case 's':
                        options->samples = atol(optarg);
                        break;
                case 't':
			options->time = atol(optarg);
                        break;
                case 'l':
                        options->limit = atol(optarg);
                        break;
		case 'p':
			if (options->nrprobe == AEP_MAX_PROBES) {
				fprintf(stderr, "Maximum number of probes reached\n");
				return -1;
			}

			options->probe[options->nrprobe] = strdup(optarg);

			/* -p /dev/ttyACM0[,<R1>,<R2>,<R3>,[<name>]] */
			token = strstr(options->probe[options->nrprobe], ",");
			if (token) {
				*token = '\0';

				aux = token + 1;
				token = strstr(aux, ",");
				if (!token)
					return -1;
				*token = '\0';
				options->rshunt0[options->nrprobe] = atoi(aux);

				aux = token + 1;
				token = strstr(aux, ",");
				if (!token)
					return -1;
				*token = '\0';
				options->rshunt1[options->nrprobe] = atoi(aux);

				aux = token + 1;
				token = strstr(aux, ",");
				if (token) {
					*token = '\0';
					options->name[options->nrprobe] = token + 1;
				} else {
					options->name[options->nrprobe] =
						basename(strdup(options->probe[options->nrprobe]));
				}
				options->rshunt2[options->nrprobe] = atoi(aux);
			}

			if (!options->rshunt0[options->nrprobe])
				options->rshunt0[options->nrprobe] = 100;
			if (!options->rshunt1[options->nrprobe])
				options->rshunt1[options->nrprobe] = 100;
			if (!options->rshunt2[options->nrprobe])
				options->rshunt2[options->nrprobe] = 100;

			options->nrprobe++;

			break;
		case 'T':
			gettimeofday(&t, NULL);
			options->timestamp = t.tv_usec + (t.tv_sec * 1000000);
			break;
                case 'V':
			printf("Version 0.0.1\n");
                        break;
		case 'h':
			printf("Ask you neighbourg\n");
			break;
                case '?':
                        fprintf(stderr, "%s: Unknown option %c'.\n",
                                argv[0], optopt);
                default:
                        return -1;
                }
        }

	if (!options->average && !options->raw)
		options->raw = true;

	if (!options->nrprobe) {
		fprintf(stderr, "probe name missing.\n");
		return -1;
	}

	for (i = 0; i < options->nrprobe; i++) {
		fprintf(stderr, "Using probe '%s'\n", options->probe[i]);
	}

	return 0;
}
	
	
static inline float avg(float avg, float val, float nr)
{
        return avg + ((val - avg) / nr);
}

/*
 * Answers are always 64 bytes long, except for the ACK after the
 * start command which is 1 byte long
 */
static char answer[64];

static int aep_read_answer(int fd)
{
	char c = '\0';

again:
	if (read(fd, &c, sizeof(c)) != sizeof(c)) {
		fprintf(stderr, "Failed to read probe\n");
		return -1;
	}

	/* reverse-eng: probes can return some \0 before an ack */
	if (c == '\0')
		goto again;

	if (c != (char)AEP_ACK) {
		fprintf(stderr, "Probe returned an error\n");
		return -1;
	}

	if (read(fd, answer, sizeof(answer)) < 0) {
		fprintf(stderr, "failed to read magic number: %m\n");
		return -1;
	}

	return 0;
}

static int aep_read_vendor(int fd)
{
	if (aep_read_answer(fd))
		return -1;

	fprintf(stderr, "Vendor name: %s\n", answer);
	return 0;
}

static int aep_read_version(int fd)
{
	int val;

	if (aep_read_answer(fd))
		return -1;

	memcpy(&val, answer, sizeof(val));
	fprintf(stderr, "Device version: 0x%x\n", val);

	return 0;
}

static int aep_read_magic(int fd)
{
	const ssize_t magic_len = 8;
	char magic[magic_len];
	ssize_t i;
	char c = '0';

	while (c != (char)AEP_ACK) {
		if (read(fd, &c, sizeof(c)) < 0) {
			fprintf(stderr, "Failed to read ack char: %m\n");
			return -1;
		}
	}

	if (read(fd, magic, magic_len) != magic_len) {
		fprintf(stderr, "Failed to read magic message");
		return -1;
	}

	for (i = 0; i < magic_len; i++) {
		if (magic[i] != (char)AEP_MAGIC) {
			fprintf(stderr, "bad magic number\n");
			return 1;
		}
	}

	return 0;
}

static int aep_read_rate(int fd)
{
	int val;

	if (aep_read_answer(fd))
		return -1;

	memcpy(&val, answer, sizeof(val));
	fprintf(stderr, "Sampling rate: %d Hz\n", val);

	return 0;
}

static int aep_sync(int fd)
{
	ssize_t len;
	unsigned char zero[8] = { 0x0 };

	len = write(fd, zero, sizeof(zero));
	if (len < 0) {
		fprintf(stderr, "failed to sync empty frame\n");
		return -1;
	}

	return 0;
}

static int aep_write_char(int fd, char c)
{
	ssize_t len;

	len = write(fd, &c, sizeof(c));
	if (len < 0) {
		fprintf(stderr, "failed to send char 0x%x\n", (char)c);
		return -1;
	}

	return 0;
}

static int aep_stop(int fd)
{
	return aep_write_char(fd, AEP_STOP);
}

static int aep_channel_setup(int fd, int channel)
{
	/*
	 * Configuration packet is 3 bytes:
	 *
	 * - AEP_CONFIG
	 * - channel number (2 bits)
	 * - channel to acquire (3 bits)
	 *   - bit0 : enable power
	 *   - bit1 : enable voltage
	 *   - bit2 : enable current
	 */
	char cmd[] = {
		AEP_CONFIG, channel, 7
	};

	if (write(fd, cmd, sizeof(cmd)) != sizeof(cmd))
		return -1;

	if (aep_read_answer(fd))
		return -1;

	return 0;
}

static int aep_reset(int fd)
{
	if (aep_write_char(fd, AEP_RESET))
		return -1;

	return aep_read_magic(fd);
}

int aep_config(int fd)
{
	/*
	 * Initiate protocol communication
	 */
	if (aep_sync(fd))
		return -1;

	/* 
	 * Clear all channels, stop all captures, ...
	 * Magic number telling the probe is ready.
	 */
	/* 
	 * Clear all channels, stop all captures, ...
	 * Magic number telling the probe is ready.
	 */
	if (aep_reset(fd))
		return -1;

	/*
	 * Get the probe version
	 */
	if (aep_write_char(fd, AEP_VERSION))
		return -1;

	if (aep_read_version(fd))
		return -1;

	/*
	 * Get the probe and vendor name
	 */
	if (aep_write_char(fd, AEP_VENDOR))
		return -1;

	if (aep_read_vendor(fd))
		return -1;

	/*
	 * Get the current 10KHz is the only value
	 * possible ATM.
	 */
	if (aep_write_char(fd, AEP_RATE))
		return -1;

	if (aep_read_rate(fd))
		return -1;

	/* Config for channel 0 */
	if (aep_channel_setup(fd, 0))
		return -1;

	/* Config for channel 1 */
	if (aep_channel_setup(fd, 1))
		return -1;

	/* Config for channel 2 */
	if (aep_channel_setup(fd, 2))
		return -1;

	return 0;
}

static int aep_start(int fd)
{
	char ack;

	/* Start the acquisition sequence */
	if (aep_write_char(fd, AEP_START))
		return -1;

	/*
	 * We have here a special case. When we start the acquisition,
	 * the answer is no longer 64 bytes long but 1 byte.
	 * Why ? Dunno ...
	 */
	if (read(fd, &ack, sizeof(ack)) != sizeof(ack))
		return -1;

	if (ack != (char)AEP_ACK)
		return -1;

	return 0;
}

static inline unsigned short aep_char2short(unsigned char *value)
{
	return value[0] + (value[1] <<  8);
}

static inline float aep_char2float(unsigned char *value)
{
	return ((float)aep_char2short(value) / 1000.0);
}

/*
 * ValueCorrected = ValueMeasured * 100 / Rsh (mOhms)
 */
static inline float aep_correction(float rshunt, float value)
{
	return value * 100.0 / rshunt;
}

/*
 * We can't read more the 10000 samples at once
 */
static int aep_read_frame(struct aep_device *dev, struct aep_frame *frame,
			  int nr_frame, unsigned char *buffer)
{
	ssize_t len, frame_size;
	ssize_t nr_bytes, bytes_read;
	ssize_t aux = 0;
	struct timeval now;
	int i = 0;

	/*
	 * Each frame is 10 unsigned short:
	 *
	 *  - unsigned short frame_number
	 *  - unsigned short power0
	 *  - unsigned short volt0
	 *  - unsigned short current0
	 *
	 *  - unsigned short power1
	 *  - unsigned short volt1
	 *  - unsigned short current1
	 *
	 *  - unsigned short power2
	 *  - unsigned short volt2
	 *  - unsigned short current2
	 *
	 * For the sake of code simplicity, we assume all channels are
	 * enabled.
	 */
	frame_size = sizeof(unsigned short) * 10;

	if (ioctl(dev->fd, FIONREAD, &nr_bytes) < 0) {
		fprintf(stderr, "Failed to read the number of bytes available: %m\n");
		return -1;
	}

	/*
	 * The AEP provides us a truncated number of frames we have to
	 * handle.  We let the buffer non empty and read only the
	 * amount of frames which are non truncated. The
	 * EDGE_TRIGGERED option of the epoll will prevent to exit
	 * immediately after entering epoll_wait and will wait for more
	 * data.
	 */
	nr_frame = nr_bytes / frame_size;

	/*
	 * Count the number of bytes len to read the number of frames
	 * passed as parameter and allocate the buffer accordingly.
	 */
	len = nr_frame * frame_size;

	while (len) {

		bytes_read = read(dev->fd, &buffer[aux], len);
		if (bytes_read < 0) {

			/*
			 * Assumption: the AEP returns an error when
			 * the buffer is not ready to provide any data
			 * and returns EAGAIN, even if the fd is not
			 * flagged in a non blocking mode. We wait a
			 * 10 ms to let some new data to arrive before
			 * retrying again. That shouldn't happen.
			 */
			if (errno == EAGAIN) {
				usleep(10000);
				continue;
			}

			fprintf(stderr, "failed to read data: %m\n");
			return -1;
		}

		len -= bytes_read;
		aux += bytes_read;
	};

	/*
	 * Let's decode each frames and fill the values in the frame
	 * structure array.
	 */
	for (i = 0; i < nr_frame; i++) {

		gettimeofday(&now, NULL);

		frame[i].timestamp = now.tv_sec * 1000000 + now.tv_usec;

		/* Some conversion to the right byte ordering */
		frame[i].frame  = aep_char2short(&buffer[i * 20]);

		if (dev->next_frame != frame[i].frame) {

			/* Frame number wrapped */
			if (frame[i].frame < dev->next_frame)
				dev->frame_missed += (frame[i].frame + dev->next_frame) - USHRT_MAX;
			else
				dev->frame_missed += frame[i].frame - dev->next_frame;

			dev->next_frame = frame[i].frame;
		}

		dev->next_frame++;

		frame[i].channel[0][AEP_POWER]   = aep_correction(dev->channel[0].rshunt, aep_char2float(&buffer[(i * 20) + 2]));
		frame[i].channel[0][AEP_VOLTAGE] = aep_correction(dev->channel[0].rshunt, aep_char2float(&buffer[(i * 20) + 4]));
		frame[i].channel[0][AEP_CURRENT] = aep_correction(dev->channel[0].rshunt, aep_char2float(&buffer[(i * 20) + 6]));

		frame[i].channel[1][AEP_POWER]   = aep_correction(dev->channel[1].rshunt, aep_char2float(&buffer[(i * 20) + 8]));
		frame[i].channel[1][AEP_VOLTAGE] = aep_correction(dev->channel[1].rshunt, aep_char2float(&buffer[(i * 20) + 10]));
		frame[i].channel[1][AEP_CURRENT] = aep_correction(dev->channel[1].rshunt, aep_char2float(&buffer[(i * 20) + 12]));

		frame[i].channel[2][AEP_POWER]   = aep_correction(dev->channel[2].rshunt, aep_char2float(&buffer[(i * 20) + 14]));
		frame[i].channel[2][AEP_VOLTAGE] = aep_correction(dev->channel[2].rshunt, aep_char2float(&buffer[(i * 20) + 16]));
		frame[i].channel[2][AEP_CURRENT] = aep_correction(dev->channel[2].rshunt, aep_char2float(&buffer[(i * 20) + 18]));
	}

	return nr_frame;
}

static inline void aep_print_average(struct aep_device *dev)
{
	int i, j;

	printf("avg ");
	for (i = 0; i < AEP_NR_CHANNELS; i++)
		for (j = 0; j < AEP_CHANNEL_SIZE; j++)
			printf("%f ", dev->channel[i].stat.avg[j]);
	printf("\n");
}

static inline void aep_channel_stat(struct aep_device *dev,
				    struct aep_channel *channel, float *value)
{
	int i;
	
	for (i = 0; i < AEP_CHANNEL_SIZE; i++) {
		channel->stat.avg[i] = avg(channel->stat.avg[i],
					   value[i], dev->total_frame + 1);
	}
}

static bool inline aep_throttle(struct aep_device *dev, struct aep_options *opt)
{
	if (!opt->limit)
		return false;

	return !!(dev->total_frame % opt->limit);
}

static inline unsigned long aep_timestamp(unsigned long tbegin, unsigned long tend)
{
	return tend - tbegin;
}

static volatile sig_atomic_t intr = 0;

static void sighandler(int sig)
{
	intr = 1;
}

static void aep_print(struct aep_options *opt, struct aep_device *dev,
		      struct aep_frame *frame, char *buffer)
{
	int i, j;
	size_t len = 0;

	if (opt->average && !aep_throttle(dev, opt))
		aep_print_average(dev);

	if (opt->raw && !aep_throttle(dev, opt)) {

		len = sprintf(buffer, "%s ", dev->name);

		if (opt->timestamp)
			len += sprintf(buffer + len, "%lu ",
				       aep_timestamp(opt->timestamp, frame->timestamp));

		for (i = 0; i < AEP_NR_CHANNELS; i++) {

			for (j = 0; j < AEP_CHANNEL_SIZE; j++)
				len += sprintf(buffer + len, "%f ", frame->channel[i][j]);

			aep_channel_stat(dev, &dev->channel[i], frame->channel[i]);
		}
		
		len = sprintf(buffer + len, "\n");

		printf("%s", buffer);
	}
}

struct aep_thread_arg {
	struct aep_options *aep_opt;
	int index;
};

static void *aep_capture_thread(void *data)
{
	int i, epfd;
	int frame_read, nr_frame = 10000;
	unsigned char *buffer;
	char *output;

	struct aep_frame frame[nr_frame];
	struct aep_thread_arg *arg = data;
	struct aep_options *aep_opt = arg->aep_opt;
	struct aep_device aep_dev;
	struct epoll_event eevt;

	buffer = malloc(10000 * sizeof(struct aep_frame));
	if (!buffer) {
		fprintf(stderr, "Failed to allocate memory buffer\n");
		return NULL;
	}

	output = malloc(1024);
	if (!output) {
		fprintf(stderr, "Failed to allocate memory for output\n");
		return NULL;
	}

	memset(&aep_dev, 0, sizeof(aep_dev));
	memset(&eevt, 0, sizeof(eevt));

	aep_dev.fd = open(aep_opt->probe[arg->index], O_RDWR);
	if (aep_dev.fd < 0) {
		fprintf(stderr, "Failed to open '%s': %m\n", aep_opt->probe[arg->index]);
		return NULL;
	}
	
	aep_dev.index = arg->index;
	aep_dev.name = aep_opt->name[arg->index];
	aep_dev.channel[0].rshunt = aep_opt->rshunt0[arg->index];
	aep_dev.channel[1].rshunt = aep_opt->rshunt1[arg->index];
	aep_dev.channel[2].rshunt = aep_opt->rshunt2[arg->index];
	
	eevt.events = EPOLLIN | EPOLLET;
	eevt.data.ptr = &aep_dev;

	epfd = epoll_create(1);
	if (epfd < 0) {
		fprintf(stderr, "Failed to create poll fd: %m\n");
		return NULL;
	}

	if (epoll_ctl(epfd, EPOLL_CTL_ADD, aep_dev.fd, &eevt) < 0) {
		fprintf(stderr, "Failed to add fd to epoll: %m\n");
		return NULL;
	}

	if (aep_config(aep_dev.fd)) {
		fprintf(stderr, "Failed to configure probe\n");
		return NULL;
	}

	if (aep_start(aep_dev.fd)) {
		fprintf(stderr, "Failed to start probe\n");
		return NULL;
	}

	signal(SIGINT, sighandler);
again:
	while (!intr) {

		int nr_fds;
		struct aep_device *dev;

		nr_fds = epoll_wait(epfd, &eevt, 1, -1);
		if (nr_fds < 0) {
			if (errno == EINTR)
				goto again;
			fprintf(stderr, "Failed to epoll_wait: %m\n");
			return NULL;
		}

		if (!(eevt.events & EPOLLIN))
			continue;

		dev = (struct aep_device *)eevt.data.ptr;

		frame_read = aep_read_frame(dev, frame, nr_frame, buffer);
		if (frame_read < 0) {
			fprintf(stderr, "Failed to read frames\n");
			return NULL;
		}

		for (i = 0; i < frame_read; i++, dev->total_frame++)
			aep_print(aep_opt, dev, &frame[i], output);
	}

	if (aep_stop(aep_dev.fd))
		fprintf(stderr, "Failed to stop probe '%s'", aep_dev.name);

	if (aep_reset(aep_dev.fd))
		fprintf(stderr, "Failed to reset probe '%s'", aep_dev.name);

	if (aep_dev.frame_missed)
		fprintf(stderr, "%s missed %ld frames\n", aep_dev.name,
			aep_dev.frame_missed);

	return NULL;
}

int main(int argc, char *argv[])
{
	int i;
	struct aep_options aep_opt;
	struct aep_thread_arg *arg;

	pthread_t threads[AEP_MAX_PROBES];

	setbuf(stdout, NULL);

	if (aep_getopt(argc, argv, &aep_opt)) {
		fprintf(stderr, "Failed to get options\n");
		return 1;
	}

	for (i = 0; i < aep_opt.nrprobe; i++) {

		arg = malloc(sizeof(*arg));
		arg->aep_opt = &aep_opt;
		arg->index = i;

		pthread_create(&threads[i], NULL,
			       aep_capture_thread, arg);
	}

	for (i = 0; i < aep_opt.nrprobe; i++)
		pthread_join(threads[i], NULL);

	return 0;
}