/* 
 * This program source code file is part of the KTag project, a DIY laser tag
 * game with customizable features and wide interoperability.
 *
 * 🛡 <https://ktag.clubk.club> 🃞
 *
 * Copyright © 2024-2026 Joseph P. Kearney and the KTag developers.
 *
 * This program is free software: you can redistribute it and/or modify it under
 * the terms of the GNU Affero General Public License as published by the Free
 * Software Foundation, either version 3 of the License, or (at your option) any
 * later version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE.  See the GNU Affero General Public License for more
 * details.
 *
 * There should be a copy of the GNU Affero General Public License in the LICENSE
 * file in the root of this repository.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <SystemK.h>
#include <System_Events.h>
#include <IR.h>
#include <driver/gpio.h>
#include <driver/rmt_tx.h>
#include <driver/rmt_rx.h>

#define IR_RESOLUTION_HZ (1 * 1000 * 1000) // 1MHz resolution, 1 tick = 1us

#define IR_TX_PIN GPIO_NUM_6
#define IR_TX_RIGHT_ENABLE GPIO_NUM_7
#define IR_TX_LEFT_ENABLE GPIO_NUM_5

#define IR_RX_LEFT_PIN GPIO_NUM_2
#define IR_RX_FORWARD_PIN GPIO_NUM_15
#define IR_RX_RIGHT_PIN GPIO_NUM_42

static const char *TAG = "IR";
static rmt_channel_handle_t Tx_Channel = NULL;

static gpio_config_t tx_right_enable_gpio_config = {
    .pin_bit_mask = (1ULL << IR_TX_RIGHT_ENABLE),
    .mode = GPIO_MODE_OUTPUT, // Set mode to output
    .pull_up_en = GPIO_PULLUP_ENABLE,
    .pull_down_en = GPIO_PULLDOWN_DISABLE,
    .intr_type = GPIO_INTR_DISABLE};

static gpio_config_t tx_left_enable_gpio_config = {
    .pin_bit_mask = (1ULL << IR_TX_LEFT_ENABLE),
    .mode = GPIO_MODE_OUTPUT, // Set mode to output
    .pull_up_en = GPIO_PULLUP_ENABLE,
    .pull_down_en = GPIO_PULLDOWN_DISABLE,
    .intr_type = GPIO_INTR_DISABLE};

// Create simple encoder
rmt_copy_encoder_config_t Copy_Encoder_Config;
rmt_encoder_handle_t Tx_Encoder;

rmt_transmit_config_t Tx_Config = {
    .loop_count = 0, // no transfer loop
};

rmt_receive_config_t Rx_Config = {
    .signal_range_min_ns = 5 * 1000,         // A pulse whose width is smaller than this threshold will be treated as glitch and ignored.
    .signal_range_max_ns = 20 * 1000 * 1000, // RMT will stop receiving if one symbol level has kept longer than this value.
    .flags.en_partial_rx = false};

static rmt_channel_handle_t Left_Rx_Channel = NULL;
static rmt_channel_handle_t Forward_Rx_Channel = NULL;
static rmt_channel_handle_t Right_Rx_Channel = NULL;

static QueueHandle_t Receive_Queue;

static TimedPulseTrain_T Left_Rxd_RMT_Data;
static TimedPulseTrain_T Forward_Rxd_RMT_Data;
static TimedPulseTrain_T Right_Rxd_RMT_Data;

volatile TickType_t Last_Time_Checked_In_Ticks = 0;
TagSensorLocation_T Active_Tag_Sensor = TAG_SENSOR_NONE;
SemaphoreHandle_t Tag_Sensor_Mutex;
static const TickType_t LOCKOUT_TIME_IN_TICKS = pdMS_TO_TICKS(100);

#define IR_RX_STACK_SIZE 4096
static StaticTask_t xTaskBuffer;
static StackType_t xStack[IR_RX_STACK_SIZE];
static TaskHandle_t IR_Rx_Task_Handle;

static TagPacket_T Shot_Packet;

typedef struct
{
    uint8_t count;
    TagSensorLocation_T location;
} RxNotification_T;

// Prevent double-receiving (or triple-receiving!) IR packets by locking out the other sensors for a short time once one sensor becomes active.
// A sensor is allowed if it is the already-active sensor or the lockout period has expired.
static BaseType_t Is_Sensor_Allowed(TagSensorLocation_T sensor_location, BaseType_t *xHigherPriorityTaskWoken)
{
    BaseType_t is_allowed = pdFALSE;
    TickType_t current_time_in_ticks = xTaskGetTickCountFromISR();

    if (xSemaphoreTakeFromISR(Tag_Sensor_Mutex, xHigherPriorityTaskWoken) == pdTRUE)
    {
        if ((Active_Tag_Sensor != TAG_SENSOR_NONE) && (Active_Tag_Sensor != sensor_location) &&
            ((current_time_in_ticks - Last_Time_Checked_In_Ticks) < LOCKOUT_TIME_IN_TICKS))
        {
            // We're in lockout period and this is not the active sensor--ignore this sensor.
            is_allowed = pdFALSE;
        }
        else
        {
            // It's OK to allow this sensor.
            Active_Tag_Sensor = sensor_location;
            Last_Time_Checked_In_Ticks = current_time_in_ticks;
            is_allowed = pdTRUE;
        }

        xSemaphoreGiveFromISR(Tag_Sensor_Mutex, xHigherPriorityTaskWoken);
    }

    return is_allowed;
}

static bool RMT_Rx_Left_Done_Callback(rmt_channel_handle_t channel, const rmt_rx_done_event_data_t *edata, void *user_data)
{
    BaseType_t xHigherPriorityTaskWoken = pdFALSE;

    if (Is_Sensor_Allowed(TAG_SENSOR_LEFT, &xHigherPriorityTaskWoken) == pdTRUE)
    {
        RxNotification_T notice = {.count = edata->num_symbols * 2,
                                   .location = TAG_SENSOR_LEFT};
        xQueueSendFromISR(Receive_Queue, &notice, &xHigherPriorityTaskWoken);
    }

    return xHigherPriorityTaskWoken;
}

static bool RMT_Rx_Forward_Done_Callback(rmt_channel_handle_t channel, const rmt_rx_done_event_data_t *edata, void *user_data)
{
    BaseType_t xHigherPriorityTaskWoken = pdFALSE;

    if (Is_Sensor_Allowed(TAG_SENSOR_FORWARD, &xHigherPriorityTaskWoken) == pdTRUE)
    {
        RxNotification_T notice = {.count = edata->num_symbols * 2,
                                   .location = TAG_SENSOR_FORWARD};
        xQueueSendFromISR(Receive_Queue, &notice, &xHigherPriorityTaskWoken);
    }

    return xHigherPriorityTaskWoken;
}

static bool RMT_Rx_Right_Done_Callback(rmt_channel_handle_t channel, const rmt_rx_done_event_data_t *edata, void *user_data)
{
    BaseType_t xHigherPriorityTaskWoken = pdFALSE;

    if (Is_Sensor_Allowed(TAG_SENSOR_RIGHT, &xHigherPriorityTaskWoken) == pdTRUE)
    {
        RxNotification_T notice = {.count = edata->num_symbols * 2,
                                   .location = TAG_SENSOR_RIGHT};
        xQueueSendFromISR(Receive_Queue, &notice, &xHigherPriorityTaskWoken);
    }

    return xHigherPriorityTaskWoken;
}

static const rmt_rx_event_callbacks_t left_cbs = {
    .on_recv_done = RMT_Rx_Left_Done_Callback,
};

static const rmt_rx_event_callbacks_t forward_cbs = {
    .on_recv_done = RMT_Rx_Forward_Done_Callback,
};

static const rmt_rx_event_callbacks_t right_cbs = {
    .on_recv_done = RMT_Rx_Right_Done_Callback,
};

static void IR_Receive_Task(void *param)
{
    KLOG_INFO(TAG, "IR Receive Task Started");

    while (true)
    {
        RxNotification_T notice;
        DecodedPacket_T *result = NULL;

        if (xQueueReceive(Receive_Queue, &notice, portMAX_DELAY) == pdPASS)
        {
            if (notice.location == TAG_SENSOR_FORWARD)
            {
                KLOG_INFO(TAG, "TAG_SENSOR_FORWARD Rx'd");
                Forward_Rxd_RMT_Data.count = notice.count;
                result = PROTOCOLS_MaybeDecodePacket(&Forward_Rxd_RMT_Data);
            }
            else if (notice.location == TAG_SENSOR_LEFT)
            {
                KLOG_INFO(TAG, "TAG_SENSOR_LEFT Rx'd");
                Left_Rxd_RMT_Data.count = notice.count;
                result = PROTOCOLS_MaybeDecodePacket(&Left_Rxd_RMT_Data);
            }
            else if (notice.location == TAG_SENSOR_RIGHT)
            {
                KLOG_INFO(TAG, "TAG_SENSOR_RIGHT Rx'd");
                Right_Rxd_RMT_Data.count = notice.count;
                result = PROTOCOLS_MaybeDecodePacket(&Right_Rxd_RMT_Data);
            }

            if (result != NULL)
            {
                if (result->Generic.type == DECODED_PACKET_TYPE_TAG_RECEIVED)
                {
                    KEvent_T tag_received_event = {.ID = KEVENT_TAG_RECEIVED, .Data = result};
                    Post_KEvent(&tag_received_event);
                }
                else if (result->Generic.type == DECODED_PACKET_TYPE_COMMAND_RECEIVED)
                {
                    KEvent_T command_received_event = {.ID = KEVENT_COMMAND_RECEIVED, .Data = result};
                    Post_KEvent(&command_received_event);
                }
            }
            else
            {
                KEvent_T near_miss_event = {.ID = KEVENT_NEAR_MISS, .Data = NULL};
                Post_KEvent(&near_miss_event);
            }

            // Start receiving again.
            // We need to reset all three channels for some reason--why?
            {
                ESP_ERROR_CHECK(rmt_disable(Forward_Rx_Channel));
                ESP_ERROR_CHECK(rmt_enable(Forward_Rx_Channel));
                ESP_ERROR_CHECK(rmt_receive(Forward_Rx_Channel, &Forward_Rxd_RMT_Data, sizeof(Forward_Rxd_RMT_Data.pulsetrain), &Rx_Config));
            }
            {
                ESP_ERROR_CHECK(rmt_disable(Left_Rx_Channel));
                ESP_ERROR_CHECK(rmt_enable(Left_Rx_Channel));
                ESP_ERROR_CHECK(rmt_receive(Left_Rx_Channel, &Left_Rxd_RMT_Data, sizeof(Left_Rxd_RMT_Data.pulsetrain), &Rx_Config));
            }
            {
                ESP_ERROR_CHECK(rmt_disable(Right_Rx_Channel));
                ESP_ERROR_CHECK(rmt_enable(Right_Rx_Channel));
                ESP_ERROR_CHECK(rmt_receive(Right_Rx_Channel, &Right_Rxd_RMT_Data, sizeof(Right_Rxd_RMT_Data.pulsetrain), &Rx_Config));
            }
        }
    }
}

static void Initialize_Receive_Task(void)
{
    Tag_Sensor_Mutex = xSemaphoreCreateMutex();

    if (Tag_Sensor_Mutex == NULL)
    {
        KLOG_ERROR(TAG, "Failed to create tag sensor mutex!");
    }

    IR_Rx_Task_Handle = xTaskCreateStaticPinnedToCore(
        IR_Receive_Task,      // Function that implements the task.
        "IR Rx",              // Text name for the task.
        IR_RX_STACK_SIZE,     // Stack size in words, not bytes.
        0,                    // Parameter passed into the task.
        tskIDLE_PRIORITY + 1, // Priority at which the task is created.
        xStack,               // Array to use as the task's stack.
        &xTaskBuffer,         // Variable to hold the task's data structure.
        APP_CPU_NUM);         // Core where the task should run.

    if (IR_Rx_Task_Handle == NULL)
    {
        KLOG_ERROR(TAG, "Failed to create IR Receive task!");
    }
}

void Initialize_IR()
{
    KLOG_INFO(TAG, "Initializing IR...");

    KLOG_INFO(TAG, "Creating RMT TX channel...");
    rmt_tx_channel_config_t tx_channel_cfg = {
        .gpio_num = IR_TX_PIN,             // GPIO number used by RMT TX channel.
        .clk_src = RMT_CLK_SRC_XTAL,       // Clock source of RMT TX channel, channels in the same group must use the same clock source
        .resolution_hz = IR_RESOLUTION_HZ, // Channel clock resolution, in Hz.
        .mem_block_symbols = 48,           // If DMA is not used, this field controls the size of the dedicated memory block owned by the channel.
                                           // THIS IS WRONG IN THE DOCS!  See https://github.com/espressif/esp-idf/issues/12084#issuecomment-1679881770.
        .trans_queue_depth = 1,            // Depth of the internal transaction queue.
        .flags.invert_out = false,         // Should the signal be inverted before sending it to the GPIO?
        .flags.with_dma = true,            // Should this channel use the DMA backend?
        .flags.io_loop_back = false,       // Should the signal output from the GPIO be fed to the input path as well?
        .flags.io_od_mode = false,         // Should the GPIO be configured in open-drain mode?
        .intr_priority = 1                 // RMT interrupt priority.  If set to 0 , then the driver will use a interrupt with low or medium priority (priority level may be one of 1, 2, or 3).
    };
    ESP_ERROR_CHECK(rmt_new_tx_channel(&tx_channel_cfg, &Tx_Channel));

    KLOG_INFO(TAG, "Modulate carrier to TX channel");
    rmt_carrier_config_t carrier_cfg = {
        .duty_cycle = 0.30,    // 30% duty cycle
        .frequency_hz = 38000, // 38KHz
    };
    ESP_ERROR_CHECK(rmt_apply_carrier(Tx_Channel, &carrier_cfg));

    ESP_ERROR_CHECK(rmt_enable(Tx_Channel));

    ESP_ERROR_CHECK(gpio_config(&tx_right_enable_gpio_config));
    KLOG_INFO(TAG, "Initialized right IR Tx enable as GPIO[%d].", IR_TX_RIGHT_ENABLE);

    ESP_ERROR_CHECK(gpio_config(&tx_left_enable_gpio_config));
    KLOG_INFO(TAG, "Initialized left IR Tx enable as GPIO[%d].", IR_TX_LEFT_ENABLE);

    ESP_ERROR_CHECK(rmt_new_copy_encoder(&Copy_Encoder_Config, &Tx_Encoder));

    KLOG_INFO(TAG, "Creating RMT Rx channels...");
    rmt_rx_channel_config_t rx_left_channel_cfg = {
        .gpio_num = IR_RX_LEFT_PIN,        // GPIO number used by RMT RX channel. Set to -1 if unused.
        .clk_src = RMT_CLK_SRC_XTAL,       // Clock source of RMT RX channel; channels in the same group must use the same clock source.
        .resolution_hz = IR_RESOLUTION_HZ, // Channel clock resolution, in Hz.
        .mem_block_symbols = 48,           // Size of memory block, in number of `rmt_symbol_word_t`, must be an even.
                                           // In the DMA mode, this field controls the DMA buffer size, it can be set to a large value (e.g. 1024);
                                           // In the normal mode, this field controls the number of RMT memory block that will be used by the channel.
        .flags.invert_in = true,           // Should the incoming signal be inverted before processing?
        .flags.with_dma = false,           // Should this channel use the DMA backend?
        .flags.io_loop_back = false,       // Should the signal output from the GPIO be fed to the input path as well?
        .intr_priority = 1                 // RMT interrupt priority.  If set to 0 , then the driver will use a interrupt with low or medium priority (priority level may be one of 1, 2, or 3).
    };
    ESP_ERROR_CHECK(rmt_new_rx_channel(&rx_left_channel_cfg, &Left_Rx_Channel));

    rmt_rx_channel_config_t rx_forward_channel_cfg = {
        .gpio_num = IR_RX_FORWARD_PIN,     // GPIO number used by RMT RX channel. Set to -1 if unused.
        .clk_src = RMT_CLK_SRC_XTAL,       // Clock source of RMT RX channel; channels in the same group must use the same clock source.
        .resolution_hz = IR_RESOLUTION_HZ, // Channel clock resolution, in Hz.
        .mem_block_symbols = 48,           // Size of memory block, in number of `rmt_symbol_word_t`, must be an even.
                                           // In the DMA mode, this field controls the DMA buffer size, it can be set to a large value (e.g. 1024);
                                           // In the normal mode, this field controls the number of RMT memory block that will be used by the channel.
        .flags.invert_in = true,           // Should the incoming signal be inverted before processing?
        .flags.with_dma = false,           // Should this channel use the DMA backend?
        .flags.io_loop_back = false,       // Should the signal output from the GPIO be fed to the input path as well?
        .intr_priority = 1                 // RMT interrupt priority.  If set to 0 , then the driver will use a interrupt with low or medium priority (priority level may be one of 1, 2, or 3).
    };
    ESP_ERROR_CHECK(rmt_new_rx_channel(&rx_forward_channel_cfg, &Forward_Rx_Channel));

    rmt_rx_channel_config_t rx_right_channel_cfg = {
        .gpio_num = IR_RX_RIGHT_PIN,       // GPIO number used by RMT RX channel. Set to -1 if unused.
        .clk_src = RMT_CLK_SRC_XTAL,       // Clock source of RMT RX channel; channels in the same group must use the same clock source.
        .resolution_hz = IR_RESOLUTION_HZ, // Channel clock resolution, in Hz.
        .mem_block_symbols = 48,           // Size of memory block, in number of `rmt_symbol_word_t`, must be an even.
                                           // In the DMA mode, this field controls the DMA buffer size, it can be set to a large value (e.g. 1024);
                                           // In the normal mode, this field controls the number of RMT memory block that will be used by the channel.
        .flags.invert_in = true,           // Should the incoming signal be inverted before processing?
        .flags.with_dma = false,           // Should this channel use the DMA backend?
        .flags.io_loop_back = false,       // Should the signal output from the GPIO be fed to the input path as well?
        .intr_priority = 1                 // RMT interrupt priority.  If set to 0 , then the driver will use a interrupt with low or medium priority (priority level may be one of 1, 2, or 3).
    };
    ESP_ERROR_CHECK(rmt_new_rx_channel(&rx_right_channel_cfg, &Right_Rx_Channel));

    ESP_LOGI(TAG, "register RX done callbacks");
    Receive_Queue = xQueueCreate(1, sizeof(RxNotification_T));
    assert(Receive_Queue);

    ESP_ERROR_CHECK(rmt_rx_register_event_callbacks(Left_Rx_Channel, &left_cbs, 0));
    ESP_ERROR_CHECK(rmt_rx_register_event_callbacks(Forward_Rx_Channel, &forward_cbs, 0));
    ESP_ERROR_CHECK(rmt_rx_register_event_callbacks(Right_Rx_Channel, &right_cbs, 0));

    Initialize_Receive_Task();

    Left_Rxd_RMT_Data.receiver = TAG_SENSOR_LEFT;
    Forward_Rxd_RMT_Data.receiver = TAG_SENSOR_FORWARD;
    Right_Rxd_RMT_Data.receiver = TAG_SENSOR_RIGHT;

    ESP_ERROR_CHECK(rmt_enable(Left_Rx_Channel));
    ESP_ERROR_CHECK(rmt_receive(Left_Rx_Channel, &Left_Rxd_RMT_Data, sizeof(Left_Rxd_RMT_Data.pulsetrain), &Rx_Config));

    ESP_ERROR_CHECK(rmt_enable(Forward_Rx_Channel));
    ESP_ERROR_CHECK(rmt_receive(Forward_Rx_Channel, &Forward_Rxd_RMT_Data, sizeof(Forward_Rxd_RMT_Data.pulsetrain), &Rx_Config));

    ESP_ERROR_CHECK(rmt_enable(Right_Rx_Channel));
    ESP_ERROR_CHECK(rmt_receive(Right_Rx_Channel, &Right_Rxd_RMT_Data, sizeof(Right_Rxd_RMT_Data.pulsetrain), &Rx_Config));

    xEventGroupSetBits(Get_System_Events(), SYS_IR_READY);
    KLOG_INFO(TAG, "Initialization complete.");
}

static inline void PrintPulseTrainToConsole(TimedPulseTrain_T *train)
{
    for (uint_fast16_t i = 0; i < train->count; i += 2)
    {
        KLOG_DEBUG(TAG, "%2d: (%d, %4d) (%d, %4d)", i + 1, train->bitstream[i].symbol, train->bitstream[i].duration, train->bitstream[i + 1].symbol, train->bitstream[i + 1].duration);
        vTaskDelay(pdMS_TO_TICKS(10));
    }
}

SystemKResult_T Prepare_Tag(void)
{
    TimedPulseTrain_T *Shot_Buffer;

    uint8_t team_ID;
    uint8_t player_ID;
    uint8_t weapon_ID;

    (void)SETTINGS_get_uint8_t(SYSTEMK_SETTING_TEAMID, &team_ID);
    (void)SETTINGS_get_uint8_t(SYSTEMK_SETTING_PLAYERID, &player_ID);
    (void)SETTINGS_get_uint8_t(SYSTEMK_SETTING_WEAPONID, &weapon_ID);
    Weapon_t weapon = GetWeaponFromID(weapon_ID);

    Shot_Packet.player_ID = player_ID;
    Shot_Packet.team_ID = team_ID;
    Shot_Packet.color = (uint32_t)PROTOCOLS_GetColor(weapon.Protocol, team_ID, player_ID);
    Shot_Packet.protocol = weapon.Protocol;
    Shot_Packet.damage = weapon.Damage_Per_Shot;

    Shot_Buffer = PROTOCOLS_EncodePacket(&Shot_Packet);

    KLOG_INFO(TAG, "Tag prepared (%u pulses):", Shot_Buffer->count);
    PrintPulseTrainToConsole(Shot_Buffer);

    Enable_IR_Transmit();

    return SYSTEMK_RESULT_SUCCESS;
}

SystemKResult_T Send_Tag(void)
{
    TimedPulseTrain_T *Shot_Buffer;

    Shot_Buffer = PROTOCOLS_EncodePacket(&Shot_Packet);

    ESP_ERROR_CHECK(rmt_transmit(Tx_Channel, Tx_Encoder, Shot_Buffer->pulsetrain, sizeof(rmt_symbol_word_t) * Shot_Buffer->count, &Tx_Config));

    KEvent_T tag_sent_event = {.ID = KEVENT_TAG_SENT, .Data = (void *)0x00};
    Post_KEvent(&tag_sent_event);

    return SYSTEMK_RESULT_SUCCESS;
}

void Enable_IR_Transmit()
{
    gpio_set_level(IR_TX_RIGHT_ENABLE, 0);
    gpio_set_level(IR_TX_LEFT_ENABLE, 0);
}

void Disable_IR_Transmit()
{
    gpio_set_level(IR_TX_RIGHT_ENABLE, 1);
    gpio_set_level(IR_TX_LEFT_ENABLE, 1);
}
