#include <arpa/inet.h>
#include <GLFW/glfw3.h>
#include <imgui.h>
#include <imgui_impl_glfw.h>
#include <imgui_impl_opengl3.h>
#include <implot.h>
#include <initializer_list>
#include <json-c/json.h>
#include <limits.h>
#include <math.h>
#include <signal.h>
#include <threads.h>
#include "logger.c"

#define FETCH_INTERVAL 10
#define FETCH_INTERVAL_SUBDIV 1000

#define COLOR_BG ImVec4(0.f, 0.f, 0.f, 1.f)
#define FONT_SIZE 12.f
#define FONT_SOURCE "/usr/share/fonts/droid/DroidSansMono.ttf"
#define GLSL_VERSION "#version 130"
#define WINDOW_STYLE (ImGuiWindowFlags_NoBackground | ImGuiWindowFlags_NoDecoration | ImGuiWindowFlags_NoMove | ImGuiWindowFlags_NoSavedSettings)

#define DEFAULT_ENTRIES 2000
#define DEFAULT_NTH 1
#define DEFAULT_X_AXIS 0
#define DEFAULT_X_LOW 0
#define DEFAULT_X_HIGH INT_MAX
#define DEFAULT_HM_LEFT 0
#define DEFAULT_HM_PIXEL_COUNT 0x400

#define PLOT_MIN_COLS 1
#define PLOT_MAX_COLS 8

enum Status {
    STATUS_STOPPED,
    STATUS_RUNNING,
    STATUS_STOPPING,
};

struct Plot {
    Plot(const char *name, const char *section, std::initializer_list<const char *> cols)
        : name(name), section(section), cols(cols.begin()) {}

    const char *name;
    const char *section;
    const char *const *cols;
};

// Window globals
GLFWwindow *g_window;
ImGuiIO *g_io;
ImGuiStyle *g_imgui_style;
ImPlotStyle *g_implot_style;

// Data globals
const char *g_x_axes[] = {
    "rowid",
    "steps",
#define FOR_CORE(i) "cycl_" #i,
    FOR_CORES
#undef FOR_CORE
};

int g_status;
int g_entries = DEFAULT_ENTRIES;
int g_nth = DEFAULT_NTH;
int g_x_axis = DEFAULT_X_AXIS;
int g_x_low = DEFAULT_X_LOW;
int g_x_high = DEFAULT_X_HIGH;
int g_hm_left = DEFAULT_HM_LEFT;
int g_hm_pixel_count = DEFAULT_HM_PIXEL_COUNT;
int g_hm_pixel_pow; // calculate on startup
int g_x_current;
thrd_t g_data_fetching_thread;

// Plot globals
Plot g_plots[] = {
    Plot("cycl", "general", {
#define FOR_CORE(i) "cycl_" #i,
        FOR_CORES
#undef FOR_CORE
        nullptr,
    }),
    Plot("mall", "general", {
#define FOR_CORE(i) "mall_" #i,
        FOR_CORES
#undef FOR_CORE
        nullptr,
    }),
    Plot("pnum", "general", {
#define FOR_CORE(i) "pnum_" #i,
        FOR_CORES
#undef FOR_CORE
        nullptr,
    }),
    Plot("ppop", "general", {
#define FOR_CORE(i) "pfst_" #i, "plst_" #i,
        FOR_CORES
#undef FOR_CORE
        nullptr,
    }),
    Plot("ambs", "general", {
#define FOR_CORE(i) "amb0_" #i, "amb1_" #i,
        FOR_CORES
#undef FOR_CORE
        nullptr,
    }),
    Plot("eevs", "general", {
#define FOR_CORE(i) "emb0_" #i, "emb1_" #i, "eliv_" #i, "edea_" #i,
        FOR_CORES
#undef FOR_CORE
        nullptr,
    }),
};

#define PLOT_COUNT (int)(sizeof(g_plots) / sizeof(g_plots[0]))

// Layout globals
float g_data_col_width;
bool g_data_col_visible = true;
int g_plot_cols = 2;
int g_plot_col_selected;
int g_plot_row_selected;
Plot *g_plot_cells[PLOT_MAX_COLS][PLOT_COUNT];
Plot *g_plot_selected = &g_plots[0];
bool g_plot_maximized;

// ----------------------------------------------------------------------------
// Data functions
// ----------------------------------------------------------------------------
int data_calc_max_hm_pixel_pow(void) {
    return (int)floor(log2((float)(MVEC_SIZE - g_hm_left) / (float)g_hm_pixel_count));
}

void data_clamp(int *field, int low, int high) {
    assert(field);
    if (*field < low) *field = low;
    if (*field > high) *field = high;
}

void data_touched(void) {
    data_clamp(&g_entries, 1, DEFAULT_ENTRIES);
    data_clamp(&g_nth, DEFAULT_NTH, INT_MAX);
    data_clamp(&g_x_low, DEFAULT_X_LOW, INT_MAX);
    data_clamp(&g_x_high, g_x_low + 1, DEFAULT_X_HIGH);
 #if !defined(MVEC_LOOP)
    data_clamp(&g_hm_left, DEFAULT_HM_LEFT, MVEC_SIZE);
 #endif
    data_clamp(&g_hm_pixel_count, 1, DEFAULT_HM_PIXEL_COUNT);
    data_clamp(&g_hm_pixel_pow, 0, data_calc_max_hm_pixel_pow());
    g_x_current = 0;
}

void data_reset_values(void) {
    g_entries = DEFAULT_ENTRIES;
    g_nth = DEFAULT_NTH;
    g_x_axis = DEFAULT_X_AXIS;
    g_x_low = DEFAULT_X_LOW;
    g_x_high = DEFAULT_X_HIGH;
    g_hm_left = DEFAULT_HM_LEFT;
    g_hm_pixel_count = DEFAULT_HM_PIXEL_COUNT;
    g_hm_pixel_pow = data_calc_max_hm_pixel_pow();
}

void data_reset_plot_cells(void) {
    for (int i = 0; i < PLOT_MAX_COLS; i++) {
        for (int j = 0; j < PLOT_COUNT; j++) {
            g_plot_cells[i][j] = nullptr;
        }
    }
}

void data_fetch(void) {
    struct json_object *request = json_object_new_object();
    json_object_object_add(request, "request", json_object_new_string("data"));
    json_object_object_add(request, "entries", json_object_new_int(g_entries));
    json_object_object_add(request, "nth", json_object_new_int(g_nth));
    json_object_object_add(request, "x-axis", json_object_new_string(g_x_axes[g_x_axis]));
    json_object_object_add(request, "x-low", json_object_new_int(g_x_low));
    json_object_object_add(request, "x-high", json_object_new_int(g_x_high));
    json_object_object_add(request, "hm-left", json_object_new_int(g_hm_left));
    json_object_object_add(request, "hm-pixel-count", json_object_new_int(g_hm_pixel_count));
    json_object_object_add(request, "hm-pixel-pow", json_object_new_int(g_hm_pixel_pow));
    json_object_object_add(request, "x-current", json_object_new_int(g_x_current));
    const char *request_str = json_object_to_json_string(request);

    g_info("Sending request to server: %s", request_str);
    int socket_fd = socket(AF_INET, SOCK_STREAM, 0);
    struct sockaddr_in socket_addr;
    memset(&socket_addr, 0, sizeof(struct sockaddr_in));
    socket_addr.sin_family = AF_INET;
    socket_addr.sin_port = htons(PORT);
    inet_pton(AF_INET, IP, &socket_addr.sin_addr);
    if (connect(socket_fd, (struct sockaddr *)&socket_addr, sizeof(struct sockaddr_in))) assert(false);
    json_object_to_fd(socket_fd, request, 0);
    shutdown(socket_fd, SHUT_WR);

    struct json_object *response = json_object_from_fd(socket_fd);
    const char *response_str = json_object_to_json_string(response);
    g_info("Server responded with: %s", response_str);

    json_object_put(request);
    json_object_put(response);
}

int data_fetching_thread(void *data) {
    (void)data;

    assert(g_status == STATUS_RUNNING || g_status == STATUS_STOPPING);

    while (g_status == STATUS_RUNNING) {
        data_fetch();

        for (int i = 0; i < FETCH_INTERVAL_SUBDIV && g_status == STATUS_RUNNING; i++) {
            usleep((FETCH_INTERVAL * 1000000) / FETCH_INTERVAL_SUBDIV);
        }
    }

    assert(g_status == STATUS_STOPPING);
    g_status = STATUS_STOPPED;
    return 0;
}

void data_start_fetching(void) {
    g_info("Starting data fetching thread");
    g_status = STATUS_RUNNING;
    thrd_create(&g_data_fetching_thread, (thrd_start_t)data_fetching_thread, nullptr);
}

void data_stop_fetching(void) {
    assert(g_status == STATUS_RUNNING);
    g_info("Stopping data fetching thread");
    g_status = STATUS_STOPPING;
    thrd_join(g_data_fetching_thread, nullptr);
}

// ----------------------------------------------------------------------------
// GUI functions
// ----------------------------------------------------------------------------
void gui_print_data_col(void) {
    const ImGuiViewport *viewport = ImGui::GetMainViewport();
    const ImVec2 next_win_pos = viewport->Pos;
    const ImVec2 next_win_size = ImVec2(-1.f, viewport->Size.y);

    ImGui::SetNextWindowPos(next_win_pos);
    ImGui::SetNextWindowSize(next_win_size);
    ImGui::Begin("data-col", nullptr, WINDOW_STYLE);
    g_data_col_width = ImGui::GetWindowWidth();

    ImGui::SeparatorText("SALIS data client");
    ImGui::LabelText("name", NAME);
    ImGui::LabelText("seed", "%#lx", SEED);
    ImGui::LabelText("server", IP ":" PORT_STR);
    ImGui::LabelText("arch", ARCH);
    ImGui::LabelText("cores", "%d", CORES);
    ImGui::LabelText("mvec-size", "%#lx", MVEC_SIZE);
 #if defined(MVEC_LOOP)
    ImGui::LabelText("mvec-loop", "true");
 #else
    ImGui::LabelText("mvec-loop", "false");
 #endif
    ImGui::LabelText("data-push", "%#lx", DATA_PUSH_INTERVAL);

    ImGui::SeparatorText("Data fields");

    switch (g_status) {
    case STATUS_STOPPED:
        if (ImGui::InputInt("entries", &g_entries, 0, 0, ImGuiInputTextFlags_CharsDecimal)) data_touched();
        if (ImGui::InputInt("nth", &g_nth, 0, 0, ImGuiInputTextFlags_CharsDecimal)) data_touched();

        if (ImGui::BeginCombo("x-axis", g_x_axes[g_x_axis])) {
            for (int i = 0; i < CORES + 2; i++) {
                if (ImGui::Selectable(g_x_axes[i], g_x_axis == i)) {
                    g_x_axis = i;
                    data_touched();
                }
            }

            ImGui::EndCombo();
        }

        if (ImGui::InputInt("x-low", &g_x_low, 0, 0, ImGuiInputTextFlags_CharsDecimal)) data_touched();
        if (ImGui::InputInt("x-high", &g_x_high, 0, 0, ImGuiInputTextFlags_CharsDecimal)) data_touched();
        if (ImGui::InputInt("hm-left", &g_hm_left, 0, 0, ImGuiInputTextFlags_CharsDecimal)) data_touched();
        if (ImGui::InputInt("hm-pxl-count", &g_hm_pixel_count, 0, 0, ImGuiInputTextFlags_CharsDecimal)) data_touched();
        if (ImGui::InputInt("hm-pxl-pow", &g_hm_pixel_pow, 0, 0, ImGuiInputTextFlags_CharsDecimal)) data_touched();
        if (ImGui::Button("Run", ImVec2(-1.f, 0.f))) data_start_fetching();
        if (ImGui::Button("Reset", ImVec2(-1.f, 0.f))) data_reset_values();

        break;
    case STATUS_RUNNING:
    case STATUS_STOPPING:
        ImGui::LabelText("entries", "%d", g_entries);
        ImGui::LabelText("nth", "%d", g_nth);
        ImGui::LabelText("x-axis", "%s", g_x_axes[g_x_axis]);
        ImGui::LabelText("x-low", "%d", g_x_low);
        ImGui::LabelText("x-high", "%d", g_x_high);
        ImGui::LabelText("hm-left", "%d", g_hm_left);
        ImGui::LabelText("hm-pxl-count", "%d", g_hm_pixel_count);
        ImGui::LabelText("hm-pxl-pow", "%d", g_hm_pixel_pow);

        if (g_status == STATUS_RUNNING) {
            if (ImGui::Button("Stop", ImVec2(-1.f, 0.f))) data_stop_fetching();
            ImGui::LabelText("##", "Running");
        } else {
            ImGui::LabelText("##", "Stopping");
        }
    }

    ImGui::SeparatorText("Layout");
    ImGui::SliderInt("cols", &g_plot_cols, PLOT_MIN_COLS, PLOT_MAX_COLS);

    ImGui::End();
}

void gui_print_plots(void) {
    const char *section_current = g_plots[0].section;
    const char *section_next = nullptr;
    bool plots_covered[PLOT_COUNT] = { 0 };

    const ImGuiViewport *viewport = ImGui::GetMainViewport();
    const ImVec2 next_win_pos = g_data_col_visible ? ImVec2(g_data_col_width, viewport->Pos.y) : viewport->Pos;
    const ImVec2 next_win_size = g_data_col_visible ? ImVec2(viewport->Size.x - g_data_col_width, -1.f) : ImVec2(viewport->Size.x, -1.f);

    ImGui::SetNextWindowPos(next_win_pos);
    ImGui::SetNextWindowSize(next_win_size);
    ImGui::Begin("plots", nullptr, WINDOW_STYLE);

    int col = 0;
    int row = 0;

    while (section_current) {
        ImGui::SeparatorText(section_current);
        ImGui::BeginTable("plots-table", g_plot_cols);

        for (int i = 0; i < PLOT_COUNT; i++) {
            if (g_plots[i].section != section_current) {
                section_next = (!section_next && !plots_covered[i]) ? g_plots[i].section : section_next;
                continue;
            }

            ImGui::TableNextColumn();

            if (&g_plots[i] == g_plot_selected) {
                g_plot_col_selected = col;
                g_plot_row_selected = row;
                g_implot_style->Colors[ImPlotCol_FrameBg] = g_imgui_style->Colors[ImGuiCol_FrameBg];
            }

            if (ImPlot::BeginPlot(g_plots[i].name)) {
                //int test_x[] = {0,1,2,3};
                //int test_y1[] = {1,2,3,4};
                //int test_y2[] = {2,4,8,16};
                //ImPlot::PlotLine("test1", test_x, test_y1, 4);
                //ImPlot::PlotLine("test2", test_x, test_y2, 4);
                ImPlot::EndPlot();
            }

            g_implot_style->Colors[ImPlotCol_FrameBg] = COLOR_BG;

            g_plot_cells[col][row] = &g_plots[i];
            col = (col + 1) % g_plot_cols;
            row += col ? 0 : 1;

            plots_covered[i] = true;
        }

        section_current = section_next;
        section_next = nullptr;
        ImGui::EndTable();

        row += col ? 1 : 0;
        col = 0;
    }

    ImGui::End();
}

void gui_print_plot_maximized(void) {
    const ImGuiViewport *viewport = ImGui::GetMainViewport();
    ImGui::SetNextWindowPos(viewport->Pos);
    ImGui::SetNextWindowSize(viewport->Size);
    ImGui::Begin("plot-fullscreen", nullptr, WINDOW_STYLE);

    if (ImPlot::BeginPlot(g_plot_selected->name, viewport->Size)) {
        //int test_x[] = {0,1,2,3};
        //int test_y1[] = {1,2,3,4};
        //int test_y2[] = {2,4,8,16};
        //ImPlot::PlotLine("test1", test_x, test_y1, 4);
        //ImPlot::PlotLine("test2", test_x, test_y2, 4);
        ImPlot::EndPlot();
    }

    ImGui::End();
}

void gui_print(void) {
    if (g_plot_maximized) {
        gui_print_plot_maximized();
    } else {
        if (g_data_col_visible) gui_print_data_col();
        gui_print_plots();
    }
}

// ----------------------------------------------------------------------------
// Main functions
// ----------------------------------------------------------------------------
void sig_handler(int signo) {
    (void)signo;

    g_warn("Signal received, will stop SALIS data client...");
    if (g_status == STATUS_RUNNING) data_stop_fetching();
    glfwSetWindowShouldClose(g_window, GLFW_TRUE);
}

void glfw_error_callback(int error, const char* description) {
    g_warn("GLFW error %d: %s", error, description);
}

void glfw_key_callback(GLFWwindow* window, int key, int scancode, int action, int mods) {
    (void)window;
    (void)scancode;

    if (action != GLFW_PRESS) return;

    if (g_plot_maximized) {
        switch (mods) {
        case GLFW_MOD_CONTROL:
            switch (key) {
            case GLFW_KEY_C:
                glfwSetWindowShouldClose(g_window, GLFW_TRUE);
                break;
            }

            break;

        case 0:
            switch (key) {
            case GLFW_KEY_F:
                g_plot_maximized = !g_plot_maximized;
                break;
            }

            break;
        }

        return;
    }

    switch (mods) {
    case GLFW_MOD_CONTROL:
        switch (key) {
        case GLFW_KEY_C:
            glfwSetWindowShouldClose(g_window, GLFW_TRUE);
            break;
        case GLFW_KEY_N:
            g_data_col_visible = !g_data_col_visible;
            break;
        case GLFW_KEY_LEFT:
            g_plot_cols -= g_plot_cols > 1 ? 1 : 0;
            data_reset_plot_cells();
            break;
        case GLFW_KEY_RIGHT:
            g_plot_cols += g_plot_cols < PLOT_MAX_COLS ? 1 : 0;
            data_reset_plot_cells();
            break;
        }

        break;

    case 0:
        switch (key) {
        case GLFW_KEY_LEFT:
            g_plot_col_selected -= g_plot_col_selected ? 1 : 0;
            g_plot_selected = g_plot_cells[g_plot_col_selected][g_plot_row_selected];
            break;
        case GLFW_KEY_RIGHT:
            g_plot_col_selected += (g_plot_col_selected < PLOT_MAX_COLS - 1 && g_plot_cells[g_plot_col_selected + 1][g_plot_row_selected]) ? 1 : 0;
            g_plot_selected = g_plot_cells[g_plot_col_selected][g_plot_row_selected];
            break;
        case GLFW_KEY_UP:
            g_plot_row_selected -= g_plot_row_selected ? 1 : 0;
            g_plot_selected = g_plot_cells[g_plot_col_selected][g_plot_row_selected];
            break;
        case GLFW_KEY_DOWN:
            g_plot_row_selected += (g_plot_row_selected < PLOT_COUNT - 1 && g_plot_cells[g_plot_col_selected][g_plot_row_selected + 1]) ? 1 : 0;
            g_plot_selected = g_plot_cells[g_plot_col_selected][g_plot_row_selected];
            break;
        case GLFW_KEY_F:
            g_plot_maximized = !g_plot_maximized;
            break;
        }

        break;
    }
}

int main(int argc, char **argv) {
    (void)argc;
    (void)argv;

    signal(SIGINT, sig_handler);
    signal(SIGTERM, sig_handler);

    g_info("Starting SALIS data client");

    g_info("Initializing GLFW");
    glfwSetErrorCallback(glfw_error_callback);
    glfwInitHint(GLFW_WAYLAND_LIBDECOR, GLFW_WAYLAND_DISABLE_LIBDECOR);
    if (!glfwInit()) assert(false);

    float scale = ImGui_ImplGlfw_GetContentScaleForMonitor(glfwGetPrimaryMonitor());
    g_window = glfwCreateWindow((int)(800 * scale), (int)(600 * scale), "SALIS data client", nullptr, nullptr);
    assert(g_window);
    glfwSetKeyCallback(g_window, glfw_key_callback);
    glfwMakeContextCurrent(g_window);
    glfwSwapInterval(1); // enable vsync

    g_info("Initializing ImGui");
    IMGUI_CHECKVERSION();
    ImGui::CreateContext();
    ImPlot::CreateContext();

    g_io = &ImGui::GetIO();
    g_io->Fonts->AddFontFromFileTTF(FONT_SOURCE);
    g_io->IniFilename = nullptr;

    g_imgui_style = &ImGui::GetStyle();
    g_imgui_style->Colors[ImGuiCol_WindowBg] = COLOR_BG;
    g_imgui_style->FontScaleDpi = scale;
    g_imgui_style->FontSizeBase = FONT_SIZE;
    g_imgui_style->ItemSpacing = ImVec2(g_imgui_style->ItemSpacing.x, 2.f);
    g_imgui_style->ScaleAllSizes(scale);

    g_implot_style = &ImPlot::GetStyle();
    g_implot_style->Colors[ImPlotCol_FrameBg] = COLOR_BG;

    ImGui_ImplGlfw_InitForOpenGL(g_window, true);
    ImGui_ImplOpenGL3_Init(GLSL_VERSION);

    g_hm_pixel_pow = data_calc_max_hm_pixel_pow();

    // Main loop
    while (!glfwWindowShouldClose(g_window)) {
        glfwPollEvents();

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

        gui_print();

        ImGui::Render();
        int display_w, display_h;
        glfwGetFramebufferSize(g_window, &display_w, &display_h);
        glViewport(0, 0, display_w, display_h);
        glClearColor(0.f, 0.f, 0.f, 1.f);
        glClear(GL_COLOR_BUFFER_BIT);
        ImGui_ImplOpenGL3_RenderDrawData(ImGui::GetDrawData());
        glfwSwapBuffers(g_window);
    }

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

    g_info("Stopping SALIS data client");
    glfwDestroyWindow(g_window);
    glfwTerminate();
    return 0;
}