#include "augment_renderer.h"
#include "config.h"
#include "frame_buffer_helper.hpp"
#include "stereo_camera.hpp"
#include "texture_renderer.h"
#include "video_encoder.h"

#include "third_party/scope_guard.hpp"

#include <imgui.h>
#include <imgui_impl_glfw.h>
#include <imgui_impl_opengl3.h>

#include <glad/gl.h>
#include <GLFW/glfw3.h>

#include <fmt/chrono.h>
#include <fmt/format.h>
#include <spdlog/spdlog.h>

#include <cassert>
#include <cstdlib>

CUcontext cuda_ctx;

int main() {

    spdlog::set_level(spdlog::level::trace);

    // setup glfw and main window
    glfwSetErrorCallback([](int error, const char *desc) {
        SPDLOG_ERROR("GLFW error: code = {}, description = {}", error, desc);
        assert(false);
    });
    auto ret = glfwInit();
    assert(ret == GLFW_TRUE);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 4);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 6);
    glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);
    auto main_window = glfwCreateWindow(main_window_width, main_window_height, "RemoteAR", nullptr, nullptr);
    assert(main_window != nullptr);
    glfwMakeContextCurrent(main_window);
//    glfwSwapInterval(1);

    // load opengl functions
    auto version = gladLoadGL(glfwGetProcAddress);
    assert(version > 0);
    SPDLOG_INFO("Loaded OpenGL {}.{}", GLAD_VERSION_MAJOR(version), GLAD_VERSION_MINOR(version));

#ifndef NDEBUG
    // log opengl error
    glEnable(GL_DEBUG_OUTPUT);
    glDebugMessageCallback([](GLenum source, GLenum type, GLuint id, GLenum severity,
                              GLsizei length, const GLchar *message, const void *user_data) {
        if (type == GL_DEBUG_TYPE_ERROR) {
            SPDLOG_ERROR("OpenGL error: type = {}, severity = {}, message = {}", type, severity, message);
            assert(false);
        }
    }, nullptr);
#endif

    // setup imgui context
    IMGUI_CHECKVERSION();
    ImGui::CreateContext();
    auto io = ImGui::GetIO();
    io.ConfigFlags |= ImGuiConfigFlags_NavEnableKeyboard;
    ImGui::StyleColorsDark();
    ImGui_ImplGlfw_InitForOpenGL(main_window, true);
    ImGui_ImplOpenGL3_Init();

    // setup cuda context
    cuInit(0);
    create_cuda_context(&cuda_ctx);

    // working staffs
    stereo_camera camera;
    texture_renderer tex_renderer;

    frame_buffer_helper output_fbo;
    output_fbo.initialize(output_frame_width, output_frame_height);

    video_encoder encoder;
    encoder.initialize();

    augment_renderer left_ar, right_ar;
    left_ar.initialize(&tex_renderer);
    right_ar.initialize(&tex_renderer);
    left_ar.set_background(&camera.left_rgb_image);
    right_ar.set_background(&camera.right_rgb_image);

    int camera_fps = default_camera_fps;
    float exposure_time_ms = default_camera_exposure_time_ms;
    float analog_gain = default_camera_analog_gain;
    float output_bitrate_mbps = default_video_stream_bitrate / 1e6f;

    FILE *video_save_file = nullptr;
    auto video_save_file_closer = sg::make_scope_guard([&]() {
        if (video_save_file == nullptr) return;
        fclose(video_save_file);
    });

    // main loop
    while (!glfwWindowShouldClose(main_window)) {

        glfwPollEvents();

        ImGui_ImplOpenGL3_NewFrame();
        ImGui_ImplGlfw_NewFrame();
        ImGui::NewFrame();

        ImGui::ShowDemoWindow();

        if (ImGui::Begin("Remote AR Control")) {

            // extra actions to make consistency
            if (!camera.is_capturing() && encoder.is_encoding()) {
                encoder.stop_encode();
            }

            // camera control
            if (ImGui::CollapsingHeader("Camera")) {
                ImGui::PushID("Camera");

                // camera actions
                ImGui::SeparatorText("Actions");
                if (!camera.is_opened()) {
                    if (ImGui::Button("Open")) {
                        camera.open();
                    }
                } else {
                    if (ImGui::Button("Close")) {
                        camera.close();
                    }
                    ImGui::SameLine();
                    if (!camera.is_capturing()) {
                        if (ImGui::Button("Start")) {
                            camera.start_capture(1000 * exposure_time_ms, analog_gain, camera_fps);
                        }
                    } else {
                        if (ImGui::Button("Stop")) {
                            camera.stop_capture();
                        }
                    }
                }

                // camera configs
                if (camera.is_opened()) {
                    ImGui::SeparatorText("Configs");

                    // don't allow config change wile camera is capturing
                    if (camera.is_capturing()) {
                        ImGui::BeginDisabled();
                    }

                    ImGui::PushItemWidth(200);
                    ImGui::SliderInt("Frame Rate (fps)", &camera_fps, 1, 60);
                    ImGui::DragFloat("Exposure Time (ms)", &exposure_time_ms,
                                     0.1, 1, 1e3f / (float) camera_fps, "%.01f");
                    ImGui::DragFloat("Analog Gain (dB)", &analog_gain, 0.1, 0, 24, "%.01f");
                    ImGui::PopItemWidth();

                    if (camera.is_capturing()) {
                        ImGui::EndDisabled();
                    }
                }

                ImGui::PopID();
            }

            // video streamer control
            if (camera.is_capturing() && ImGui::CollapsingHeader("Video Streamer")) {
                ImGui::PushID("Streamer");

                ImGui::SeparatorText("Actions");
                if (!encoder.is_encoding()) {
                    if (ImGui::Button("Start")) {
                        // create save file
                        assert(video_save_file == nullptr);
                        auto file_name = fmt::format("record_{:%Y_%m_%d_%H_%M_%S}.hevc",
                                                     std::chrono::system_clock::now());
                        video_save_file = fopen(file_name.c_str(), "wb");
                        encoder.start_encode(output_fbo.tex_width, output_fbo.tex_height,
                                             camera_fps, (int) (output_bitrate_mbps * 1e6));
                    }
                } else {
                    if (ImGui::Button("Close")) {
                        encoder.stop_encode();
                        fclose(video_save_file);
                        video_save_file = nullptr;
                    }
                }

                if (encoder.is_encoding()) {
                    ImGui::BeginDisabled();
                }

                ImGui::PushItemWidth(200);
                ImGui::DragFloat("Bitrate (Mbps)", &output_bitrate_mbps, 0.1, 1, 20, "%.01f");
                ImGui::PopItemWidth();

                if (encoder.is_encoding()) {
                    ImGui::EndDisabled();
                }

                ImGui::PopID();
            }

        }
        ImGui::End();
        ImGui::Render();

        std::chrono::high_resolution_clock::time_point start_time;
        if (camera.is_capturing()) {
            camera.retrieve_raw_images();
            start_time = std::chrono::high_resolution_clock::now();
            camera.debayer_images();
        }

        if (encoder.is_encoding()) {
            // draw frame for streaming
            glBindFramebuffer(GL_DRAW_FRAMEBUFFER, output_fbo.fbo);
            glViewport(0, 0, output_fbo.tex_width, output_fbo.tex_height);
            left_ar.render({-1, -1, 1, 2});
            right_ar.render({0, -1, 1, 2});

            // encode frame
            output_fbo.download_pixels();
            void *frame_data;
            size_t frame_length;
            encoder.encode_frame(output_fbo.pbo_res, &frame_data, &frame_length);

            SPDLOG_TRACE("Time used: {}ms", std::chrono::duration_cast<std::chrono::milliseconds>(
                    std::chrono::high_resolution_clock::now() - start_time).count());

            // save encoded frame
            fwrite(frame_data, frame_length, 1, video_save_file);
        }

        int frame_width, frame_height;
        glBindFramebuffer(GL_DRAW_FRAMEBUFFER, 0);
        glfwGetFramebufferSize(main_window, &frame_width, &frame_height);
        glViewport(0, 0, frame_width, frame_height);
        glClear(GL_COLOR_BUFFER_BIT);

        if (camera.is_capturing()) {
            // draw frame in the screen
            left_ar.render({-1, 1, 2, -2});
        }

        ImGui_ImplOpenGL3_RenderDrawData(ImGui::GetDrawData());
        glfwSwapBuffers(main_window);

        if (camera.is_capturing()) {
            glFlush();
        }
    }

    // cleanup
    std::atexit([]() {
        cuCtxDestroy(cuda_ctx);
    });

    ImGui_ImplOpenGL3_Shutdown();
    ImGui_ImplGlfw_Shutdown();
    ImGui::DestroyContext();

    glfwDestroyWindow(main_window);
    glfwTerminate();

    return 0;
}