#include "blk_box_drivers/ssegs.hpp"

#include "tm1640.hpp"
#include "pins.h"
#include "freertos/FreeRTOS.h"
#include "freertos/queue.h"
#include "freertos/event_groups.h"
#include "esp_log.h"

#include <atomic>
#include <cmath>
#include <tuple>

const static uint32_t TICKER_PERIOD_MS = 100;

const static uint8_t MODULE_IDX = 0;
const static uint8_t GAME_IDX = 4;

TM1640 ssegs(PIN_SSEG_CLK, PIN_SSEG_DAT);
const static size_t CMD_QUEUE_SIZE = 10;
QueueHandle_t cmd_queue;

std::atomic<int32_t> game_time = 0;
std::atomic<int32_t> module_time = 0;

// for notifying users of events
const static uint32_t EVENT_CMDS_FLUSHED = (1 << 0);
const static uint32_t EVENT_MODULE_POSITIVE = (1 << 1);
const static uint32_t EVENT_MODULE_ZERO_NEG = (1 << 2);
const static uint32_t EVENT_GAME_POSITIVE = (1 << 3);
const static uint32_t EVENT_GAME_ZERO_NEG = (1 << 4);

EventGroupHandle_t ssegs_event_group;

const static char* TAG = "ssegs";

/// Uses a compare-exchange loop to do a saturating subtraction on an atomic int32_t. Returns the old and new values.
std::pair<int32_t, int32_t> saturating_sub(std::atomic<int32_t>& v, int32_t sub) {
    int32_t cur = v.load(std::memory_order_relaxed);
    while (true) {
        int32_t desired = (cur <= sub) ? 0 : cur - sub;
        if (v.compare_exchange_weak(cur, desired,
                                    std::memory_order_relaxed)) {
            return {cur, desired}; // {old value, new value}
        }
        // cur is updated automatically with latest value on failure
    }
}

/// Updates the segment buffer to reflect the current time.
///
/// Returns `true` if the buffer has changed and needs to be redrawn.
///
/// `seg_buf.len()` should be >= 4.
bool update_segments(int32_t last_time, int32_t current_time, uint8_t seg_buf[4]) {
    const uint32_t MILLIS_10S = 100;
    const uint32_t SECOND = 1000;
    const uint32_t SECOND_10S = SECOND * 10;
    const uint32_t MINUTE = 60 * SECOND;
    const uint32_t MINUTE_10S = 10 * MINUTE;
    const uint32_t HOUR = 60 * MINUTE;
    const uint32_t HOUR_10S = 10 * HOUR;

    uint32_t time = std::abs(current_time);

    if (time > HOUR) {
        // HH.MM
        if ((current_time / MINUTE) == (last_time / MINUTE)) {
            // no change neccesary
            return false;
        }

        uint8_t h1 = (time / HOUR_10S) % 10;
        uint8_t h0 = (time / HOUR) % 10;
        uint8_t minutes = (time / MINUTE) % 60;
        uint8_t m1 = minutes / 10;
        uint8_t m0 = minutes % 10;

        seg_buf[0] = SSegController::FONT_HEX[h1];
        seg_buf[1] = SSegController::FONT_HEX[h0] | SSegController::BIT_MASK_DP;
        seg_buf[2] = SSegController::FONT_HEX[m1];
        seg_buf[3] = SSegController::FONT_HEX[m0];
        return true;
    } else if (time > MINUTE) {
        // MM.SS
        if ((current_time / SECOND) == (last_time / SECOND)) {
            // no change neccesary
            return false;
        }

        uint8_t m1 = (time / MINUTE_10S) % 10;
        uint8_t m0 = (time / MINUTE) % 10;
        uint8_t seconds = (time / SECOND) % 60;
        uint8_t s1 = seconds / 10;
        uint8_t s0 = seconds % 10;

        seg_buf[0] = SSegController::FONT_HEX[m1];
        seg_buf[1] = SSegController::FONT_HEX[m0] | SSegController::BIT_MASK_DP;
        seg_buf[2] = SSegController::FONT_HEX[s1];
        seg_buf[3] = SSegController::FONT_HEX[s0];
        return true;
    } else {
        // SS.m
        if ((current_time / MILLIS_10S) == (last_time / MILLIS_10S)) {
            // no change neccesary
            return false;
        }

        uint8_t s1 = (time / SECOND_10S) % 10;
        uint8_t s0 = (time / SECOND) % 10;
        uint8_t m1 = (time / MILLIS_10S) % 10;

        seg_buf[0] = 0; // unused digit
        seg_buf[1] = SSegController::FONT_HEX[s1];
        seg_buf[2] = SSegController::FONT_HEX[s0] | SSegController::BIT_MASK_DP;
        seg_buf[3] = SSegController::FONT_HEX[m1];
        return true;
    }

    return true;
}

static void timer_task(void* arg) {
    (void) arg;

    const TickType_t ticker_period_ticks = pdMS_TO_TICKS(TICKER_PERIOD_MS);

    ESP_LOGI(TAG, "sseg timer task starting...");

    bool game_en = false;
    bool game_running = false;
    bool game_rollover = true;
    bool module_en = false;
    bool module_running = false;
    bool module_rollover = false;

    uint8_t seg_buf[4] = {0};

    TickType_t last_wake_time = xTaskGetTickCount();
    SSegCommand cmd;

    while (true) {
        TickType_t elapsed = xTaskGetTickCount() - last_wake_time;
        if ((ticker_period_ticks > elapsed) && (xQueueReceive(cmd_queue, &cmd, ticker_period_ticks - elapsed) == pdPASS)) {
            // command received
            ESP_LOGI(TAG, "sseg command received");

            switch (cmd.type) {
                case SSegCommand::Type::SetIntensity: {
                    uint8_t intensity = std::get<uint8_t>(cmd.data);
                    ssegs.set_intensity(intensity);
                    break;
                }
                case SSegCommand::Type::EnableGameTimer: {
                    game_en = true;
                    int32_t game_time_val = game_time.load(std::memory_order_acquire);
                    if (update_segments(std::numeric_limits<int32_t>::max(), game_time_val, seg_buf)) {
                        ssegs.set_digits(GAME_IDX, seg_buf, 4);
                    }
                    break;
                }
                case SSegCommand::Type::DisableGameTimer: {
                    game_en = false;
                    game_running = false;
                    game_time.store(0, std::memory_order_release);
                    xEventGroupClearBits(ssegs_event_group, EVENT_GAME_POSITIVE | EVENT_GAME_ZERO_NEG);
                    for (uint8_t& seg : seg_buf) {
                        seg = 0;
                    }
                    ssegs.set_digits(GAME_IDX, seg_buf, 4);
                    break;
                }
                case SSegCommand::Type::StartGameTimer:
                    game_running = true;
                    break;
                case SSegCommand::Type::StopGameTimer:
                    game_running = false;
                    break;
                case SSegCommand::Type::SetGameTime: {
                    int32_t new_time = std::get<int32_t>(cmd.data);
                    int32_t last_time = game_time.exchange(new_time, std::memory_order_acq_rel);
                    if (new_time > 0) {
                        xEventGroupSetBits(ssegs_event_group, EVENT_GAME_POSITIVE);
                        xEventGroupClearBits(ssegs_event_group, EVENT_GAME_ZERO_NEG);
                    } else {
                        xEventGroupSetBits(ssegs_event_group, EVENT_GAME_ZERO_NEG);
                        xEventGroupClearBits(ssegs_event_group, EVENT_GAME_POSITIVE);
                    }
                    if (game_en) {
                        if (update_segments(last_time, new_time, seg_buf)) {
                            ssegs.set_digits(GAME_IDX, seg_buf, 4);
                        }
                    }
                    break;
                }
                case SSegCommand::Type::EnableModuleTimer: {
                    module_en = true;
                    int32_t module_time_val = module_time.load(std::memory_order_acquire);
                    if (update_segments(std::numeric_limits<int32_t>::max(), module_time_val, seg_buf)) {
                        ssegs.set_digits(MODULE_IDX, seg_buf, 4);
                    }
                    break;
                }
                case SSegCommand::Type::DisableModuleTimer: {
                    module_en = false;
                    module_running = false;
                    module_time.store(0, std::memory_order_release);
                    xEventGroupClearBits(ssegs_event_group, EVENT_MODULE_POSITIVE | EVENT_MODULE_ZERO_NEG);
                    for (uint8_t& seg : seg_buf) {
                        seg = 0;
                    }
                    ssegs.set_digits(MODULE_IDX, seg_buf, 4);
                    break;
                }
                case SSegCommand::Type::StartModuleTimer:
                    module_running = true;
                    break;
                case SSegCommand::Type::StopModuleTimer:
                    module_running = false;
                    break;
                case SSegCommand::Type::SetModuleTime: {
                    int32_t new_time = std::get<int32_t>(cmd.data);
                    int32_t last_time = module_time.exchange(new_time, std::memory_order_acq_rel);
                    if (new_time > 0) {
                        xEventGroupSetBits(ssegs_event_group, EVENT_MODULE_POSITIVE);
                        xEventGroupClearBits(ssegs_event_group, EVENT_MODULE_ZERO_NEG);
                    } else {
                        xEventGroupSetBits(ssegs_event_group, EVENT_MODULE_ZERO_NEG);
                        xEventGroupClearBits(ssegs_event_group, EVENT_MODULE_POSITIVE);
                    }
                    if (module_en) {
                        if (update_segments(last_time, new_time, seg_buf)) {
                            ssegs.set_digits(MODULE_IDX, seg_buf, 4);
                        }
                    }
                    break;
                }
                case SSegCommand::Type::SetGameRaw: {
                    std::array<uint8_t, 4> raw = std::get<std::array<uint8_t, 4>>(cmd.data);
                    ssegs.set_digits(GAME_IDX, raw.data(), 4);
                    break;
                }
                case SSegCommand::Type::SetGameDigit: {
                    auto [digit, value] = std::get<std::pair<uint8_t, uint8_t>>(cmd.data);
                    ssegs.set_digit(GAME_IDX + digit, value);
                    break;
                }
                case SSegCommand::Type::SetModuleRaw: {
                    std::array<uint8_t, 4> raw = std::get<std::array<uint8_t, 4>>(cmd.data);
                    ssegs.set_digits(MODULE_IDX, raw.data(), 4);
                    break;
                }
                case SSegCommand::Type::SetModuleDigit: {
                    auto [digit, value] = std::get<std::pair<uint8_t, uint8_t>>(cmd.data);
                    ssegs.set_digit(MODULE_IDX + digit, value);
                    break;
                }

                case SSegCommand::Type::SetGameRollover: {
                    bool rollover = std::get<bool>(cmd.data);
                    game_rollover = rollover;
                    break;
                }
                case SSegCommand::Type::SetModuleRollover: {
                    bool rollover = std::get<bool>(cmd.data);
                    module_rollover = rollover;
                    break;
                }
            }

            if (uxQueueMessagesWaiting(cmd_queue) == 0) {
                xEventGroupSetBits(ssegs_event_group, EVENT_CMDS_FLUSHED);
            }
        } else {
            // ticker finished
            last_wake_time += pdMS_TO_TICKS(TICKER_PERIOD_MS);

            bool update_module = module_en && module_running;
            bool update_game = game_en && game_running;

            // ESP_LOGI(TAG, "ticker ticked: update_game=%d, update_module=%d", update_game, update_module);

            if (update_module) {
                int32_t old_time;
                int32_t new_time;
                if (module_rollover) {
                    old_time = module_time.fetch_sub(TICKER_PERIOD_MS);
                    new_time = old_time - TICKER_PERIOD_MS; // fetch_sub returns old value
                } else {
                    std::tie(old_time, new_time) = saturating_sub(module_time, TICKER_PERIOD_MS);
                }

                if (new_time > 0) {
                    xEventGroupSetBits(ssegs_event_group, EVENT_MODULE_POSITIVE);
                    xEventGroupClearBits(ssegs_event_group, EVENT_MODULE_ZERO_NEG);
                } else {
                    xEventGroupSetBits(ssegs_event_group, EVENT_MODULE_ZERO_NEG);
                    xEventGroupClearBits(ssegs_event_group, EVENT_MODULE_POSITIVE);
                }

                if (update_segments(old_time, new_time, seg_buf)) {
                    ssegs.set_digits(MODULE_IDX, seg_buf, 4);
                }
                if (new_time == 0 && !module_rollover) {
                    // we've hit 0 and are not rolling over
                    module_running = false;
                }
            }
            if (update_game) {
                int32_t old_time;
                int32_t new_time;
                if (game_rollover) {
                    old_time = game_time.fetch_sub(TICKER_PERIOD_MS);
                    new_time = old_time - TICKER_PERIOD_MS; // fetch_sub returns old value
                } else {
                    std::tie(old_time, new_time) = saturating_sub(game_time, TICKER_PERIOD_MS);
                }

                if (new_time > 0) {
                    xEventGroupSetBits(ssegs_event_group, EVENT_GAME_POSITIVE);
                    xEventGroupClearBits(ssegs_event_group, EVENT_GAME_ZERO_NEG);
                } else {
                    xEventGroupSetBits(ssegs_event_group, EVENT_GAME_ZERO_NEG);
                    xEventGroupClearBits(ssegs_event_group, EVENT_GAME_POSITIVE);
                }

                if (update_segments(old_time, new_time, seg_buf)) {
                    ssegs.set_digits(GAME_IDX, seg_buf, 4);
                }
                if (new_time == 0 && !game_rollover) {
                    // we've hit 0 and are not rolling over
                    game_running = false;
                }
            }
        }
    }
}

// SSegController static method implementations
void SSegController::enable_game_timer() {
    xEventGroupClearBits(ssegs_event_group, EVENT_CMDS_FLUSHED);
    SSegCommand cmd = SSegCommand::EnableGameTimer();
    xQueueSend(cmd_queue, &cmd, portMAX_DELAY);
}

void SSegController::disable_game_timer() {
    xEventGroupClearBits(ssegs_event_group, EVENT_CMDS_FLUSHED);
    SSegCommand cmd = SSegCommand::DisableGameTimer();
    xQueueSend(cmd_queue, &cmd, portMAX_DELAY);
}

void SSegController::start_game_timer() {
    xEventGroupClearBits(ssegs_event_group, EVENT_CMDS_FLUSHED);
    SSegCommand cmd = SSegCommand::StartGameTimer();
    xQueueSend(cmd_queue, &cmd, portMAX_DELAY);
}

void SSegController::stop_game_timer() {
    xEventGroupClearBits(ssegs_event_group, EVENT_CMDS_FLUSHED);
    SSegCommand cmd = SSegCommand::StopGameTimer();
    xQueueSend(cmd_queue, &cmd, portMAX_DELAY);
}

void SSegController::set_game_time(int32_t millis) {
    // Align to TICKER_PERIOD_MS
    millis = millis - (millis % TICKER_PERIOD_MS);
    xEventGroupClearBits(ssegs_event_group, EVENT_CMDS_FLUSHED);
    SSegCommand cmd = SSegCommand::SetGameTime(millis);
    xQueueSend(cmd_queue, &cmd, portMAX_DELAY);
}

void SSegController::enable_module_timer() {
    xEventGroupClearBits(ssegs_event_group, EVENT_CMDS_FLUSHED);
    SSegCommand cmd = SSegCommand::EnableModuleTimer();
    xQueueSend(cmd_queue, &cmd, portMAX_DELAY);
}

void SSegController::disable_module_timer() {
    xEventGroupClearBits(ssegs_event_group, EVENT_CMDS_FLUSHED);
    SSegCommand cmd = SSegCommand::DisableModuleTimer();
    xQueueSend(cmd_queue, &cmd, portMAX_DELAY);
}

void SSegController::start_module_timer() {
    xEventGroupClearBits(ssegs_event_group, EVENT_CMDS_FLUSHED);
    SSegCommand cmd = SSegCommand::StartModuleTimer();
    xQueueSend(cmd_queue, &cmd, portMAX_DELAY);
}

void SSegController::stop_module_timer() {
    xEventGroupClearBits(ssegs_event_group, EVENT_CMDS_FLUSHED);
    SSegCommand cmd = SSegCommand::StopModuleTimer();
    xQueueSend(cmd_queue, &cmd, portMAX_DELAY);
}

void SSegController::set_module_time(int32_t millis) {
    // Align to TICKER_PERIOD_MS
    millis = millis - (millis % TICKER_PERIOD_MS);
    xEventGroupClearBits(ssegs_event_group, EVENT_CMDS_FLUSHED);
    SSegCommand cmd = SSegCommand::SetModuleTime(millis);
    xQueueSend(cmd_queue, &cmd, portMAX_DELAY);
}

void SSegController::set_game_raw(const std::array<uint8_t, 4>& segments) {
    xEventGroupClearBits(ssegs_event_group, EVENT_CMDS_FLUSHED);
    SSegCommand cmd = SSegCommand::SetGameRaw(segments);
    xQueueSend(cmd_queue, &cmd, portMAX_DELAY);
}

void SSegController::set_game_digit_raw(uint8_t digit, uint8_t segments) {
    xEventGroupClearBits(ssegs_event_group, EVENT_CMDS_FLUSHED);
    SSegCommand cmd = SSegCommand::SetGameDigit(digit, segments);
    xQueueSend(cmd_queue, &cmd, portMAX_DELAY);
}

void SSegController::set_module_raw(const std::array<uint8_t, 4>& segments) {
    xEventGroupClearBits(ssegs_event_group, EVENT_CMDS_FLUSHED);
    SSegCommand cmd = SSegCommand::SetModuleRaw(segments);
    xQueueSend(cmd_queue, &cmd, portMAX_DELAY);
}

void SSegController::set_module_digit_raw(uint8_t digit, uint8_t segments) {
    xEventGroupClearBits(ssegs_event_group, EVENT_CMDS_FLUSHED);
    SSegCommand cmd = SSegCommand::SetModuleDigit(digit, segments);
    xQueueSend(cmd_queue, &cmd, portMAX_DELAY);
}

void SSegController::game_timer_rollover(bool rollover) {
    xEventGroupClearBits(ssegs_event_group, EVENT_CMDS_FLUSHED);
    SSegCommand cmd = SSegCommand::SetGameRollover(rollover);
    xQueueSend(cmd_queue, &cmd, portMAX_DELAY);
}

void SSegController::module_timer_rollover(bool rollover) {
    xEventGroupClearBits(ssegs_event_group, EVENT_CMDS_FLUSHED);
    SSegCommand cmd = SSegCommand::SetModuleRollover(rollover);
    xQueueSend(cmd_queue, &cmd, portMAX_DELAY);
}

int32_t SSegController::get_game_time() {
    return game_time.load(std::memory_order_acquire);
}

int32_t SSegController::get_module_time() {
    return module_time.load(std::memory_order_acquire);
}

void SSegController::flush() {
    xEventGroupWaitBits(ssegs_event_group, EVENT_CMDS_FLUSHED, pdFALSE, pdTRUE, portMAX_DELAY);
}

void SSegController::wait_game_timer_done() {
    xEventGroupWaitBits(ssegs_event_group, EVENT_GAME_ZERO_NEG, pdTRUE, pdFALSE, portMAX_DELAY);
}

void SSegController::wait_module_timer_done() {
    xEventGroupWaitBits(ssegs_event_group, EVENT_MODULE_ZERO_NEG, pdTRUE, pdFALSE, portMAX_DELAY);
}

void SSegController::set_intensity(uint8_t intensity) {
    xEventGroupClearBits(ssegs_event_group, EVENT_CMDS_FLUSHED);
    SSegCommand cmd = SSegCommand::SetIntensity(intensity);
    xQueueSend(cmd_queue, &cmd, portMAX_DELAY);
}

void init_ssegs() {
    ssegs.init();

    cmd_queue = xQueueCreate(CMD_QUEUE_SIZE, sizeof(SSegCommand));
    if (cmd_queue == NULL) {
        ESP_LOGE(TAG, "Failed to create command queue!");
        return;
    }

    ssegs_event_group = xEventGroupCreate();
    if (ssegs_event_group == NULL) {
        ESP_LOGE(TAG, "Failed to create event group!");
        return;
    }

    xTaskCreate(timer_task, "ssegs_timer_task", 4096, NULL, 4, NULL);
}