#include <libmegapixels.h>
#include <stdio.h>
#include <limits.h>
#include <string.h>
#include <linux/videodev2.h>
#include <sys/ioctl.h>
#include <stdlib.h>
#include <sys/types.h>
#include <errno.h>
#include <sys/mman.h>
#include <sys/time.h>
#include <getopt.h>
#include <ctype.h>
#include <math.h>

struct buffer {
		void *start;
		size_t length;
};
struct buffer *buffers;

struct control {
		uint32_t id;

		int64_t min;
		int64_t max;
		uint64_t step;
		int64_t default_value;
		int32_t value;
};

int
xioctl(int fd, int request, void *arg)
{
	int r;
	do {
		r = ioctl(fd, request, arg);
	} while (r == -1 && errno == EINTR);
	return r;
}

void
usage(char *name)
{
	fprintf(stderr, "Usage: %s [-h] [-n offset] [-c camera] [-o file]\n", name);
	fprintf(stderr, "Measure the linearity of the sensor response\n\n");
	fprintf(stderr, "Arguments:\n");
	fprintf(stderr, "  -n count    Number of datapoint to gather, takes 1 second per point\n");
	fprintf(stderr, "  -c camera   Use a specific camera number\n");
	fprintf(stderr, "  -m modenum  Use another camera mode than the first\n");
	fprintf(stderr, "  -o file     File to store the calibration in\n");
	fprintf(stderr, "  -h          Display this help text\n");
}

void
brightness(const uint8_t *buffer, size_t length, libmegapixels_mode *mode, float *red, float *green, float *blue)
{
	// Get the offset to a single line of pixels at the middle of the frame
	size_t line = libmegapixels_mode_width_to_bytes(mode->format, mode->width);
	size_t stride = line + libmegapixels_mode_width_to_padding(mode->format, mode->width);
	size_t offset = stride * (mode->height / 2);

	unsigned long long sum_r = 0, sum_g = 0, sum_b = 0;
	unsigned int total = 0;
	for (size_t i = 0; i < line; i += 2) {
		uint8_t p1 = buffer[offset + i];
		uint8_t p2 = buffer[offset + i + 1];
		uint8_t p3 = buffer[offset + i + stride];
		//uint8_t p4 = buffer[offset + i + 1 + stride];
		total++;

		switch (mode->v4l_pixfmt) {
			case V4L2_PIX_FMT_SGRBG8:
				sum_r += p2;
				sum_g += p1;
				sum_b += p3;
				break;
			default:
				// TODO: Implement the other modes....
				fprintf(stderr, "Unsupported v4l pixfmt\n");
				exit(1);
		}
	}

	float max = (float) pow(2, libmegapixels_format_bits_per_pixel(mode->format)) - 1.0f;
	*red = (float) sum_r / (float) total / max;
	*green = (float) sum_g / (float) total / max;
	*blue = (float) sum_b / (float) total / max;
}

int
get_control(int sensor_fd, struct control *control)
{
	struct v4l2_ext_control ctrl = {};
	ctrl.id = control->id;

	struct v4l2_ext_controls ctrls = {
		.ctrl_class = 0,
		.which = V4L2_CTRL_WHICH_CUR_VAL,
		.count = 1,
		.controls = &ctrl,
	};

	if (xioctl(sensor_fd, VIDIOC_G_EXT_CTRLS, &ctrls) == -1) {
		if (errno != EINVAL) {
			fprintf(stderr, "VIDIOC_G_EXT_CTRLS\n");
		}
		return 0;
	}

	control->value = ctrl.value;

	return 1;
}

int
set_control(int sensor_fd, struct control *control)
{
	struct v4l2_ext_control ctrl = {};
	ctrl.id = control->id;
	ctrl.value = control->value;

	if (control->value > control->max || control->value < control->min) {
		fprintf(stderr, "Value %d is out of range for %ld..%ld\n", control->value, control->min, control->max);
		return 0;
	}

	struct v4l2_ext_controls ctrls = {
		.ctrl_class = 0,
		.which = V4L2_CTRL_WHICH_CUR_VAL,
		.count = 1,
		.controls = &ctrl,
	};

	if (xioctl(sensor_fd, VIDIOC_S_EXT_CTRLS, &ctrls) == -1) {
		if (errno != EINVAL) {
			fprintf(stderr, "VIDIOC_S_EXT_CTRLS\n");
		}
		return 0;
	}

	control->value = ctrl.value;

	return 1;
}

int
query_control(int sensor_fd, struct control *control)
{
	struct v4l2_query_ext_ctrl ctrl = {};
	ctrl.id = control->id;

	if (xioctl(sensor_fd, VIDIOC_QUERY_EXT_CTRL, &ctrl) == -1) {
		if (errno != EINVAL) {
			fprintf(stderr, "VIDIOC_QUERY_EXT_CTRL\n");
		}
		return 0;
	}

	control->min = ctrl.minimum;
	control->max = ctrl.maximum;
	control->step = ctrl.step;
	control->default_value = ctrl.default_value;

	return get_control(sensor_fd, control);
}

int
main(int argc, char *argv[])
{
	int c;
	int camera_id = 0;
	long res;
	char *end;
	char *outfile = NULL;
	int mode_idx = 0;
	int step = 0;
	int steps = 10;

	while ((c = getopt(argc, argv, "hc:n:o:m:")) != -1) {
		switch (c) {
			case 'c':
				res = strtol(optarg, &end, 10);
				if (end == optarg || end == NULL || *end != '\0') {
					fprintf(stderr, "Invalid number for -c\n");
					return 1;
				}
				camera_id = (int) res;
				break;
			case 'o':
				outfile = optarg;
				break;
			case 'm':
				res = strtol(optarg, &end, 10);
				if (end == optarg || end == NULL || *end != '\0') {
					fprintf(stderr, "Invalid number for -m\n");
					return 1;
				}
				mode_idx = (int) res;
				break;
			case 'n':
				res = strtol(optarg, &end, 10);
				if (end == optarg || end == NULL || *end != '\0') {
					fprintf(stderr, "Invalid number for -n\n");
					return 1;
				}
				steps = (int) res;
				break;
			case 'h':
				usage(argv[0]);
				return 0;
				break;
			case '?':
				if (optopt == 'd' || optopt == 'l') {
					fprintf(stderr, "Option -%c requires an argument.\n", optopt);
				} else if (isprint(optopt)) {
					fprintf(stderr, "Unknown option '-%c'\n", optopt);
				} else {
					fprintf(stderr, "Unknown option character x%x\n", optopt);
				}
				return 1;
			default:
				usage(argv[0]);
				return 1;
		}
	}


	char configpath[PATH_MAX];
	int ret = libmegapixels_find_config(configpath);
	libmegapixels_devconfig *config = {0};
	libmegapixels_init(&config);

	if (ret) {
		printf("Using config: %s\n", configpath);
		libmegapixels_load_file(config, configpath);
	} else {
		fprintf(stderr, "No config found for this device\n");
	}
	libmegapixels_load_uvc(config);

	if (config->count == 0) {
		fprintf(stderr, "No valid camera configuration\n");
		return 1;
	}

	if (camera_id > config->count - 1) {
		fprintf(stderr, "Camera id %d does not exist\n", camera_id);
		return 1;
	}

	libmegapixels_camera *camera = config->cameras[camera_id];
	if (libmegapixels_open(camera) != 0) {
		fprintf(stderr, "Could not open default camera\n");
		return 1;
	}
	if (mode_idx > camera->num_modes) {
		fprintf(stderr, "Invalid mode index: %d\n", mode_idx);
	}

	libmegapixels_mode *mode = camera->modes[mode_idx];
	struct v4l2_format format = {0};
	unsigned int frame_size = libmegapixels_select_mode(camera, mode, &format);
	if (frame_size == 0) {
		fprintf(stderr, "Could not select mode\n");
		return 1;
	}

	// Get the handles to the required V4L controls for the calibration
	struct control shutter = {.id = V4L2_CID_EXPOSURE};
	struct control gain = {.id = V4L2_CID_ANALOGUE_GAIN};
	if (!query_control(camera->sensor_fd, &shutter)) {
		fprintf(stderr, "Could not query V4L2_CID_EXPOSURE\n");
		return 1;
	}
	printf("Exposure: %ld..%ld step %lu val %d\n", shutter.min, shutter.max, shutter.step, shutter.value);
	if (!query_control(camera->sensor_fd, &gain)) {
		fprintf(stderr, "Could not query V4L2_CID_ANALOGUE_GAIN\n");
		return 1;
	}
	printf("Gain: %ld..%ld step %lu val %d\n", gain.min, gain.max, gain.step, gain.value);

	// Set the controls to the initial state
	shutter.value = shutter.max;
	if (!set_control(camera->sensor_fd, &shutter)) {
		fprintf(stderr, "Could not set the shutter to max\n");
		return 1;
	}

	// Do the reqular V4L2 stuff to get a frame
	struct v4l2_capability cap;
	int mplanes = 0;
	enum v4l2_buf_type buftype = V4L2_BUF_TYPE_VIDEO_CAPTURE;
	if (ioctl(camera->video_fd, VIDIOC_QUERYCAP, &cap) != 0) {
		fprintf(stderr, "VIDIOC_QUERYCAP failed: %s\n", strerror(errno));
		return 1;
	}
	if (!(cap.capabilities & V4L2_CAP_STREAMING)) {
		fprintf(stderr, "Device does not support streaming i/o\n");
		return 1;
	}
	if (!(cap.capabilities & V4L2_CAP_VIDEO_CAPTURE)) {
		if (cap.capabilities & V4L2_CAP_VIDEO_CAPTURE_MPLANE) {
			mplanes = 1;
			buftype = V4L2_BUF_TYPE_VIDEO_CAPTURE_MPLANE;
		} else {
			fprintf(stderr, "Device does not support V4L2_CAP_VIDEO_CAPTURE\n");
			return 1;
		}
	}

	struct v4l2_requestbuffers req = {0};
	req.count = 4;
	req.type = buftype;
	req.memory = V4L2_MEMORY_MMAP;
	if (xioctl(camera->video_fd, VIDIOC_REQBUFS, &req) == -1) {
		fprintf(stderr, "VIDIOC_REQBUFS failed: %s\n", strerror(errno));
		return 1;
	}
	buffers = calloc(req.count, sizeof(*buffers));
	for (int i = 0; i < req.count; i++) {
		struct v4l2_buffer buf = {0};
		buf.type = buftype;
		buf.memory = V4L2_MEMORY_MMAP;
		buf.index = i;

		struct v4l2_plane planes[1];
		if (mplanes) {
			buf.m.planes = planes;
			buf.length = 1;
		}

		if (xioctl(camera->video_fd, VIDIOC_QUERYBUF, &buf) == -1) {
			fprintf(stderr, "VIDIOC_QUERYBUF failed: %s\n", strerror(errno));
			return 1;
		}

		unsigned int offset;
		if (mplanes) {
			buffers[i].length = planes[0].length;
			offset = planes[0].m.mem_offset;
		} else {
			buffers[i].length = buf.length;
			offset = buf.m.offset;
		}
		buffers[i].start = mmap(NULL, buffers[i].length, PROT_READ | PROT_WRITE, MAP_SHARED, camera->video_fd, offset);
		if (buffers[i].start == MAP_FAILED) {
			fprintf(stderr, "mmap() failed\n");
			return 1;
		}
	}

	for (int i = 0; i < req.count; i++) {
		struct v4l2_buffer qbuf = {0};
		qbuf.type = buftype;
		qbuf.memory = V4L2_MEMORY_MMAP;
		qbuf.index = i;

		if (mplanes) {
			struct v4l2_plane qplanes[1];
			qbuf.m.planes = qplanes;
			qbuf.length = 1;
		}

		if (xioctl(camera->video_fd, VIDIOC_QBUF, &qbuf) == -1) {
			fprintf(stderr, "VIDIOC_QBUF failed: %s\n", strerror(errno));
			return 1;
		}
	}
	enum v4l2_buf_type type = buftype;
	if (xioctl(camera->video_fd, VIDIOC_STREAMON, &type) == -1) {
		fprintf(stderr, "VIDIOC_STREAMON failed: %s\n", strerror(errno));
		return 1;
	}

	// Open the target file
	FILE *outf = fopen(outfile, "w");
	if (outf == NULL) {
		fprintf(stderr, "Could not open output file\n");
		return 1;
	}

	printf("Performing initial setup...\n");

	struct timeval t_start, t_now;
	gettimeofday(&t_start, NULL);
	int stage = 1;
	double point = 1.0;
	while (stage > 0) {
		while (1) {
			fd_set(fds);
			FD_ZERO(&fds);
			FD_SET(camera->video_fd, &fds);

			int sr = select(FD_SETSIZE, &fds, NULL, NULL, NULL);
			if (sr == -1) {
				if (errno == EINTR) {
					continue;
				}
				fprintf(stderr, "select() failed: %s\n", strerror(errno));
				return 1;
			}


			struct v4l2_buffer buf = {0};
			buf.type = buftype;
			buf.memory = V4L2_MEMORY_MMAP;
			if (mplanes) {
				struct v4l2_plane dqplanes[1];
				buf.m.planes = dqplanes;
				buf.length = 1;
			}
			if (xioctl(camera->video_fd, VIDIOC_DQBUF, &buf) == -1) {
				fprintf(stderr, "VIDIOC_DQBUF failed\n");
				return 1;
			}

			if (stage == 1) {
				// Setup stage to figure out initial brightness

				gettimeofday(&t_now, NULL);
				if (t_now.tv_sec - t_start.tv_sec > 1) {
					gettimeofday(&t_start, NULL);
					float red, green, blue;
					brightness(buffers[buf.index].start, buf.bytesused, mode, &red, &green, &blue);
					printf("Brightness: %f, %f, %f\n", red, green, blue);

					if (red == 1.0f || green == 1.0f || blue == 1.0f) {
						// Clipping the sensor. Lower gain
						if (gain.value == gain.min) {
							printf("! Lower the light source brightness, out of gain range\n");
						} else {
							gain.value -= gain.step;
							set_control(camera->sensor_fd, &gain);
						}
					} else if (red > 0.9 && green > 0.9 && blue > 0.9) {
						printf("Set up target hit, continue to calibration...\n\n\n");
						stage = 2;
					}
				}
			} else if (stage == 2) {
				gettimeofday(&t_now, NULL);
				if (t_now.tv_sec - t_start.tv_sec > 1) {
					gettimeofday(&t_start, NULL);

					float red, green, blue;
					brightness(buffers[buf.index].start, buf.bytesused, mode, &red, &green, &blue);
					printf("[%4d / %4d] %f: %f, %f, %f\n", step, steps, point, red, green, blue);
					fprintf(outf, "%f,%f,%f,%f\n", point, red, green, blue);
					step++;
					// Get next shutter value
					point = (double) step / (double) steps;
					point = 1.0 - point;
					uint32_t exposure = (uint32_t) (shutter.max * point);
					if (exposure < shutter.min) {
						exposure = shutter.min;
					}

					// Set the new shutter value for the next iteration
					shutter.value = exposure;
					set_control(camera->sensor_fd, &shutter);

					if (step == steps + 1) {
						stage = 0;
						fclose(outf);
					}
				}
			}

			if (xioctl(camera->video_fd, VIDIOC_QBUF, &buf) == -1) {
				fprintf(stderr, "VIDIOC_DQBUF failed\n");
				return 1;
			}
			break;
		}
	}
	return 0;
}