#include "io_pipeline.h"

#include "camera.h"
#include "flash.h"
#include "main.h"
#include "pipeline.h"
#include "process_pipeline.h"
#include "state.h"
#include <assert.h>
#include <errno.h>
#include <fcntl.h>
#include <glib.h>
#include <math.h>
#include <stdio.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/prctl.h>

mp_state_io state_io;

MPCamera *mpcamera = NULL;

static MPPipeline *pipeline;
static GSource *capture_source;
static bool pipeline_changed = true;

typedef struct invoke_set_control {
        MPControl *control;
        int32_t int_value;
        bool bool_value;
} invoke_set_control;

static void
setup(MPPipeline *pipeline, const void *data)
{
        prctl(PR_SET_NAME, "megapixels-io", NULL, NULL, NULL);
}

void
mp_io_pipeline_start()
{
        mp_process_pipeline_start();
        pipeline = mp_pipeline_new();
        mp_pipeline_invoke(pipeline, setup, NULL, 0);
}

void
mp_io_pipeline_stop()
{
        if (capture_source) {
                g_source_destroy(capture_source);
        }
        mp_pipeline_free(pipeline);
        mp_process_pipeline_stop();
}

/*
 * Update state from IO -> Process
 */
static void
update_process_pipeline()
{
        if (!pipeline_changed) {
                return;
        }
        pipeline_changed = false;
        // Grab the latest control values
        if (!state_io.gain.manual && state_io.gain.control.id) {
                state_io.gain.value = mp_camera_control_get_int32(&state_io.gain.control);
        }

        if (!state_io.exposure.manual && state_io.exposure.control.id) {
                state_io.exposure.value = mp_camera_control_get_int32(&state_io.exposure.control);
        }

        float balance_red = 1.0f;
        float balance_blue = 1.0f;
        if (state_io.red.control.id && state_io.blue.control.id) {
                int red = mp_camera_control_get_int32(&state_io.red.control);
                int blue = mp_camera_control_get_int32(&state_io.blue.control);
                balance_red = (float)red / (float)state_io.red.max;
                balance_blue = (float)blue / (float)state_io.blue.max;
        }

        mp_state_proc new_state = {
                .camera = state_io.camera,
                .configuration = state_io.configuration,
                .burst_length = state_io.burst_length,
                .preview_width = state_io.preview_width,
                .preview_height = state_io.preview_height,
                .device_rotation = state_io.device_rotation,

                .gain.control = state_io.gain.control,
                .gain.auto_control = state_io.gain.auto_control,
                .gain.value = state_io.gain.value,
                .gain.max = state_io.gain.max,
                .gain.manual = state_io.gain.manual,

                .exposure.control = state_io.exposure.control,
                .exposure.auto_control = state_io.exposure.auto_control,
                .exposure.value = state_io.exposure.value,
                .exposure.max = state_io.exposure.max,
                .exposure.manual = state_io.exposure.manual,

                .focus.control = state_io.focus.control,
                .focus.auto_control = state_io.focus.auto_control,
                .focus.value = state_io.focus.value,
                .focus.max = state_io.focus.max,
                .focus.manual = state_io.focus.manual,

                .balance = { balance_red, 1.0f, balance_blue },

                .flash_enabled = state_io.flash_enabled,
        };

        mp_process_pipeline_update_state(&new_state);
}

static void
focus(MPPipeline *pipeline, const void *data)
{
        state_io.trigger_af = true;
}

void
mp_io_pipeline_focus()
{
        mp_pipeline_invoke(pipeline, focus, NULL, 0);
}

static void
set_control_int32(MPPipeline *pipeline, const void *data)
{
        const invoke_set_control *control_data = (const invoke_set_control *)data;
        mp_camera_control_set_int32(control_data->control, control_data->int_value);
}

void
mp_io_pipeline_set_control_int32(MPControl *control, uint32_t value)
{
        invoke_set_control data = { 0 };
        data.control = control;
        data.int_value = value;

        mp_pipeline_invoke(
                pipeline, set_control_int32, &data, sizeof(invoke_set_control));
}

static void
capture(MPPipeline *pipeline, const void *data)
{
        float gain_norm;

        // Disable the autogain/exposure while taking the burst
        mp_camera_control_set_int32(&state_io.gain.auto_control, 0);
        mp_camera_control_set_int32(&state_io.exposure.auto_control, V4L2_EXPOSURE_MANUAL);

        // Get current gain to calculate a burst length;
        // with low gain there's 3, with the max automatic gain of the ov5640
        // the value seems to be 248 which creates a 5 frame burst
        // for manual gain you can go up to 11 frames
        state_io.gain.value =
                mp_camera_control_get_int32(&state_io.gain.control);
        gain_norm = (float)state_io.gain.value / (float)state_io.gain.max;
        state_io.burst_length = (int)fmax(sqrtf(gain_norm) * 10, 2) + 1;
        state_io.burst_length = MAX(1, state_io.burst_length);
        state_io.captures_remaining = state_io.burst_length;

        // Change camera mode for capturing
        mp_process_pipeline_sync();
        mp_camera_stop_capture(mpcamera);
        struct v4l2_format format = { 0 };
        libmegapixels_select_mode(state_io.camera, state_io.mode_capture, &format);
        state_io.flush_pipeline = true;

        mp_camera_start_capture(mpcamera);

        // Enable flash
        if (state_io.flash_enabled) {
                mp_flash_enable(state_io.camera);
        }

        update_process_pipeline();

        mp_process_pipeline_capture();
}

void
mp_io_pipeline_capture()
{
        mp_pipeline_invoke(pipeline, capture, NULL, 0);
}

static void
release_buffer(MPPipeline *pipeline, const uint32_t *buffer_index)
{
        mp_camera_release_buffer(mpcamera, *buffer_index);
}

void
mp_io_pipeline_release_buffer(uint32_t buffer_index)
{
        mp_pipeline_invoke(pipeline,
                           (MPPipelineCallback)release_buffer,
                           &buffer_index,
                           sizeof(uint32_t));
}

static pid_t focus_continuous_task = 0;
static pid_t start_focus_task = 0;
static void
start_focus()
{
        // only run 1 manual focus at once
        if (!mp_camera_check_task_complete(mpcamera, start_focus_task) ||
            !mp_camera_check_task_complete(mpcamera, focus_continuous_task))
                return;

        if (state_io.focus.control.id) {
                focus_continuous_task = mp_camera_control_set_bool_bg(
                        mpcamera, &state_io.focus.control, 1);
        } else if (state_io.can_af_trigger) {
                // TODO improve
                MPControl auto_focus_start_control;
                auto_focus_start_control.id = V4L2_CID_AUTO_FOCUS_START;
                auto_focus_start_control.fd = state_io.camera->sensor_fd;
                start_focus_task = mp_camera_control_set_bool_bg(
                        mpcamera, &auto_focus_start_control, 1);
        }
}

static void
update_controls()
{
        bool state_changed = false;
        // Don't update controls while capturing
        if (state_io.captures_remaining > 0) {
                return;
        }

        if (state_io.trigger_af) {
                state_io.trigger_af = false;
                start_focus();
        }

        if (state_io.gain.manual != state_io.gain.manual_req) {
                mp_camera_control_set_bool_bg(mpcamera,
                                              &state_io.gain.auto_control,
                                              !state_io.gain.manual_req);
                state_io.gain.manual = state_io.gain.manual_req;
                state_changed = true;
        }

        if ((state_io.gain.manual ||
             (!state_io.gain.manual && state_io.gain.auto_control.id == 0)) &&
            state_io.gain.value != state_io.gain.value_req) {
                mp_camera_control_set_int32_bg(mpcamera,
                                               &state_io.gain.control,
                                               state_io.gain.value_req);
                state_io.gain.value = state_io.gain.value_req;
                state_changed = true;
        }

        if (state_io.exposure.manual != state_io.exposure.manual_req) {
                mp_camera_control_set_bool_bg(mpcamera,
                                              &state_io.exposure.auto_control,
                                              state_io.exposure.manual_req ?
                                                      V4L2_EXPOSURE_MANUAL :
                                                      V4L2_EXPOSURE_AUTO);
                state_io.exposure.manual = state_io.exposure.manual_req;
                state_changed = true;
        }

        if (state_io.exposure.manual &&
            state_io.exposure.value != state_io.exposure.value_req) {
                mp_camera_control_set_int32_bg(mpcamera,
                                               &state_io.exposure.control,
                                               state_io.exposure.value_req);
                state_io.exposure.value = state_io.exposure.value_req;
                state_changed = true;
        }

        if (state_changed) {
                pipeline_changed = true;
                update_process_pipeline();
        }
}

static void
do_aaa()
{
        bool auto_exposure =
                !state_io.exposure.manual && state_io.exposure.auto_control.id == 0;

        if (auto_exposure) {
                int direction = state_io.stats.exposure;
                int step = 0;
                if (direction > 0) {
                        // Preview is too dark

                        // Try raising the exposure time first
                        if (state_io.exposure.value < state_io.exposure.max) {
                                step = state_io.exposure.value / 16;
                                state_io.exposure.value_req =
                                        state_io.exposure.value + (step * direction);
                                printf("Expose + %d\n", state_io.exposure.value_req);
                        } else {
                                // Raise sensor gain if exposure limit is hit
                                step = state_io.gain.value / 16;
                                state_io.gain.value_req =
                                        state_io.gain.value + (step * direction);
                                printf("Gain + %d\n", state_io.gain.value_req);
                        }
                } else if (direction < 0) {
                        // Preview is too bright

                        // Lower the sensor gain first to have less noise
                        if (state_io.gain.value > 0) {
                                step = state_io.gain.value / 16;
                                state_io.gain.value_req =
                                        state_io.gain.value + (step * direction);
                                printf("Gain - %d\n", state_io.gain.value_req);
                        } else {
                                // Shorten the exposure time to go even darker
                                step = state_io.exposure.value / 16;
                                state_io.exposure.value_req =
                                        state_io.exposure.value + (step * direction);
                                printf("Expose - %d\n", state_io.exposure.value_req);
                        }
                }
        }
}

static void
on_frame(MPBuffer buffer, void *_data)
{
        // Don't process frame when the window is not active, unless we're capturing an image,
        // in which case the flash window may be active instead of this window
        if (!check_window_active() && state_io.captures_remaining == 0) {
                return;
        }

        pipeline_changed = true;

        // Only update controls right after a frame was captured
        do_aaa();
        update_controls();

        // When the mode is switched while capturing we get a couple blank frames,
        // presumably from buffers made ready during the switch. Ignore these.
        if (state_io.flush_pipeline) {
                if (state_io.blank_frame_count < 20) {
                        // Only check a 10x10 area
                        size_t test_size =
                                MIN(10, state_io.camera->current_mode->width) *
                                MIN(10, state_io.camera->current_mode->height);

                        bool image_is_blank = true;
                        for (size_t i = 0; i < test_size; ++i) {
                                if (buffer.data[i] != 0) {
                                        image_is_blank = false;
                                }
                        }

                        if (image_is_blank) {
                                ++state_io.blank_frame_count;
                                return;
                        }
                } else {
                        printf("Blank image limit reached, resulting capture may be blank\n");
                }

                state_io.flush_pipeline = false;
                state_io.blank_frame_count = 0;
        }

        // Send the image off for processing
        mp_process_pipeline_process_image(buffer);

        if (state_io.captures_remaining > 0) {
                --state_io.captures_remaining;

                if (state_io.captures_remaining == 0) {
                        // Restore the auto exposure and gain if needed
                        if (!state_io.exposure.manual) {
                                mp_camera_control_set_int32_bg(
                                        mpcamera,
                                        &state_io.exposure.auto_control,
                                        V4L2_EXPOSURE_AUTO);
                        }

                        if (!state_io.gain.manual) {
                                mp_camera_control_set_bool_bg(
                                        mpcamera, &state_io.gain.auto_control, true);
                        }

                        // Go back to preview mode
                        mp_process_pipeline_sync();
                        mp_camera_stop_capture(mpcamera);
                        struct v4l2_format format = { 0 };
                        libmegapixels_select_mode(
                                state_io.camera, state_io.mode_preview, &format);
                        state_io.flush_pipeline = true;

                        mp_camera_start_capture(mpcamera);

                        // Disable flash
                        if (state_io.flash_enabled) {
                                mp_flash_disable(state_io.camera);
                        }

                        update_process_pipeline();
                }
        }
}

static void
init_controls()
{
        MPControl focus_control;
        if (mp_camera_query_control(
                state_io.camera->sensor_fd, V4L2_CID_FOCUS_ABSOLUTE, &focus_control)) {
                state_io.focus.control = focus_control;
        } else {
                state_io.focus.control.id = 0;
        }

        MPControl auto_focus_control;
        if (mp_camera_query_control(state_io.camera->sensor_fd, V4L2_CID_FOCUS_AUTO, &auto_focus_control)) {
                mp_camera_control_set_bool_bg(
                        mpcamera, &auto_focus_control, true);
                state_io.focus.auto_control = auto_focus_control;
        } else {
                state_io.focus.auto_control.id = 0;
        }

        state_io.can_af_trigger = mp_camera_query_control(
                state_io.camera->sensor_fd, V4L2_CID_AUTO_FOCUS_START, NULL);

        MPControl gain_control;
        if (mp_camera_query_control(state_io.camera->sensor_fd, V4L2_CID_GAIN, &gain_control)) {
                state_io.gain.control = gain_control;
                state_io.gain.max = gain_control.max;
        } else if (mp_camera_query_control(
                           state_io.camera->sensor_fd, V4L2_CID_ANALOGUE_GAIN, &gain_control)) {
                state_io.gain.control = gain_control;
                state_io.gain.max = gain_control.max;
        } else {
                state_io.gain.max = 0;
                state_io.gain.control.id = 0;
        }
        if (state_io.gain.control.id) {
                state_io.gain.value = mp_camera_control_get_int32(&state_io.gain.control);
        } else {
                state_io.gain.value = 0;
        }

        MPControl auto_gain_control;
        if (mp_camera_query_control(state_io.camera->sensor_fd, V4L2_CID_AUTOGAIN, &auto_gain_control)) {
                state_io.gain.auto_control = auto_gain_control;
                state_io.gain.manual =
                        mp_camera_control_get_bool(&auto_gain_control) == 0;
        } else {
                state_io.gain.auto_control.id = 0;
        }

        MPControl exposure_control;
        if (mp_camera_query_control(state_io.camera->sensor_fd, V4L2_CID_EXPOSURE, &exposure_control)) {
                state_io.exposure.control = exposure_control;
                state_io.exposure.max = exposure_control.max;
                state_io.exposure.value = mp_camera_control_get_int32(&exposure_control);
        } else {
                state_io.exposure.control.id = 0;
        }

        MPControl auto_exposure_control;
        if (mp_camera_query_control(
                    state_io.camera->sensor_fd, V4L2_CID_EXPOSURE_AUTO, &auto_exposure_control)) {
                state_io.exposure.auto_control = auto_exposure_control;
                state_io.exposure.manual =
                        mp_camera_control_get_int32(&auto_exposure_control) == V4L2_EXPOSURE_MANUAL;
        } else {
                state_io.exposure.auto_control.id = 0;
        }

        MPControl red_control;
        if (mp_camera_query_control(
                    state_io.camera->sensor_fd, V4L2_CID_RED_BALANCE, &red_control)) {
                state_io.red.control = red_control;
                state_io.red.max = red_control.max;
        } else {
                state_io.red.control.id = 0;
        }

        MPControl blue_control;
        if (mp_camera_query_control(
                    state_io.camera->sensor_fd, V4L2_CID_BLUE_BALANCE, &blue_control)) {
                state_io.blue.control = blue_control;
                state_io.blue.max = blue_control.max;
        } else {
                state_io.blue.control.id = 0;
        }

        pipeline_changed = true;
        update_process_pipeline();
}

/*
 * State transfer from Main -> IO
 */
static void
update_state(MPPipeline *pipeline, const mp_state_io *new_state)
{
        if (state_io.camera != new_state->camera) {
                if (state_io.camera != NULL) {
                        mp_process_pipeline_sync();
                        mp_camera_stop_capture(mpcamera);
                        libmegapixels_close(state_io.camera);
                }

                if (capture_source) {
                        g_source_destroy(capture_source);
                        capture_source = NULL;
                }

                state_io.camera = new_state->camera;
                if (state_io.camera) {
                        libmegapixels_open(state_io.camera);
                        mpcamera = mp_camera_new(state_io.camera);
                        state_io.mode_preview = NULL;
                        state_io.mode_capture = NULL;
                        float score = 0;
                        int area_preview =
                                state_io.preview_width * state_io.preview_height;
                        if (area_preview == 0) {
                                area_preview = 1280 * 720;
                        }
                        for (int m = 0; m < state_io.camera->num_modes; m++) {
                                float mscore = 0;
                                if (state_io.camera->modes[m]->rate > 29) {
                                        mscore += 1;
                                }
                                int mode_area = state_io.camera->modes[m]->width *
                                                state_io.camera->modes[m]->height;
                                mscore += 1.0f -
                                          (float)(ABS(mode_area - area_preview) /
                                                  area_preview);
                                if (mscore > score) {
                                        state_io.mode_preview =
                                                state_io.camera->modes[m];
                                        score = mscore;
                                }
                        }
                        long area = 0;
                        for (int m = 0; m < state_io.camera->num_modes; m++) {
                                long this_pixels = state_io.camera->modes[m]->width *
                                                   state_io.camera->modes[m]->height;

                                if (this_pixels > area) {
                                        area = this_pixels;
                                        state_io.mode_capture =
                                                state_io.camera->modes[m];
                                }
                        }
                        if (state_io.mode_preview == NULL &&
                            state_io.mode_capture != NULL) {
                                // If no fast preview mode is available, make due
                                // with slow modes.
                                state_io.mode_preview = state_io.mode_capture;
                        }
                        if (state_io.mode_preview != NULL) {
                                if (state_io.camera->video_fd == 0) {
                                        libmegapixels_open(state_io.camera);
                                }
                                struct v4l2_format format = { 0 };
                                libmegapixels_select_mode(state_io.camera,
                                                          state_io.mode_preview,
                                                          &format);
                        }

                        mp_camera_start_capture(mpcamera);
                        capture_source = mp_pipeline_add_capture_source(
                                pipeline, mpcamera, on_frame, NULL);

                        init_controls();
                }
        }
        state_io.configuration = new_state->configuration;
        state_io.burst_length = new_state->burst_length;
        state_io.preview_width = new_state->preview_width;
        state_io.preview_height = new_state->preview_height;
        state_io.device_rotation = new_state->device_rotation;

        if (state_io.camera) {
                state_io.gain.value = new_state->gain.value;
                state_io.gain.value_req = new_state->gain.value_req;
                state_io.gain.manual = new_state->gain.manual;
                state_io.gain.manual_req = new_state->gain.manual_req;

                state_io.exposure.value = new_state->exposure.value;
                state_io.exposure.value_req = new_state->exposure.value_req;
                state_io.exposure.manual = new_state->exposure.manual;
                state_io.exposure.manual_req = new_state->exposure.manual_req;

                state_io.focus.value = new_state->focus.value;
                state_io.focus.value_req = new_state->focus.value_req;
                state_io.focus.manual = new_state->focus.manual;
                state_io.focus.manual_req = new_state->focus.manual_req;

                state_io.flash_enabled = new_state->flash_enabled;

                state_io.stats.exposure = new_state->stats.exposure;
                state_io.stats.temp = new_state->stats.temp;
                state_io.stats.tint = new_state->stats.tint;
                state_io.stats.focus = new_state->stats.focus;
        }

        update_process_pipeline();
}

void
mp_io_pipeline_update_state(const mp_state_io *state)
{
	if (!pipeline) {
		printf("no pipeline\n");
		exit(1);
	}
        mp_pipeline_invoke(pipeline,
                           (MPPipelineCallback)update_state,
                           state,
                           sizeof(mp_state_io));
}