#include <chrono>

#include "DataLinkManager.h"
#include "esp_log.h"
#include "esp_timer.h"
#include "freertos/projdefs.h"
#include "freertos/semphr.h"
#include "esp_random.h"
#include "portmacro.h"

#define FRAME_DEQUEUE_TIMEOUT_MS 2000
#define FRAME_ENQUEUE_TIMEOUT_MS 50

void DataLinkManager::init_scheduler(){
    for (int i = 0; i < num_channels; i++){
        async_rx_queue_mutex[i] = xSemaphoreCreateMutex();
        rx_fragment_mutex[i] = xSemaphoreCreateMutex();
        sliding_window_mutex[i] = xSemaphoreCreateMutex();
        send_ack_queue_mutex[i] = xSemaphoreCreateMutex();

        ESP_LOGI(DEBUG_LINK_TAG, "Starting Frame Scheduler task for channel %d", i);
        auto args = (frame_scheduler_args*)malloc(sizeof(frame_scheduler_args));
        args->channel_id = i;
        args->that = this;
        xTaskCreate(DataLinkManager::frame_scheduler, "Scheduler", 4096, static_cast<void*>(args), 4, &scheduler_task);
    }

    xTaskCreate(DataLinkManager::receive_thread_main, "Receiver", 8192, static_cast<void*>(this), 5, &receive_task);
    xTaskCreate(DataLinkManager::send_ack_thread_main, "Send ACKs", 8192, static_cast<void*>(this), 5, &send_ack_task);
}

/**
 * @brief Schedules which frame to send
 *
 * Scheduler:
 * - All frames will be pushed to the back onto a queue
 * - When a generic frame sends a chunk, it will be pushed back to the queue for the next chunk to be sent
 *
 * Scheduling may change (above scheduler will lead to starvation of control frames depending on the number of generic frames/fragments to send)
 */
[[noreturn]] void DataLinkManager::frame_scheduler(void* args){
    const auto parsed_args = static_cast<frame_scheduler_args*>(args);
    uint8_t channel = parsed_args->channel_id;
    DataLinkManager* link_layer_obj = parsed_args->that;

    if (link_layer_obj == nullptr){
        ESP_LOGE(DEBUG_LINK_TAG, "Frame Scheduler failed to start due to invalid pointer");
        vTaskDelete(nullptr);
    }

    ESP_LOGI(DEBUG_LINK_TAG, "Starting Frame Scheduler task");
    while(!link_layer_obj->stop_tasks){
        link_layer_obj->scheduler_send(channel);
    }
    free(args);
    vTaskDelete(nullptr);
}

/**
 * @brief Pushes a frame to the scheduler
 *
 * @param frame
 * @param channel
 * @return esp_err_t
 */
esp_err_t DataLinkManager::push_frame_to_scheduler(SchedulerMetadata frame, uint8_t channel){
    if (frame.data == nullptr){
        ESP_LOGE(DEBUG_LINK_TAG, "Frame data is null");
        return ESP_ERR_INVALID_ARG;
    }

    if (channel >= num_channels){
        ESP_LOGE(DEBUG_LINK_TAG, "Invalid channels");
        return ESP_ERR_INVALID_ARG;
    }

    if (frame.data->size() == 0){
        ESP_LOGE(DEBUG_LINK_TAG, "Invalid Frame Length");

        return ESP_ERR_INVALID_ARG;
    }

    int64_t now = esp_timer_get_time();
    frame.enqueue_time_ns = now;

    frame_queue[channel]->enqueue(std::move(frame), std::chrono::milliseconds(FRAME_ENQUEUE_TIMEOUT_MS));

    // ESP_LOGI(DEBUG_LINK_TAG, "Pushed frame to queue on channel %d", channel);

    return ESP_OK;
}

/**
 * @brief Scheduler sending the actual frame at the top of the heap on a channel
 *
 * @return esp_err_t
 */
esp_err_t DataLinkManager::scheduler_send(uint8_t channel){
    if (phys_comms == nullptr){
        ESP_LOGE(DEBUG_LINK_TAG, "Failed to send frame due to no RMT object");
        return ESP_ERR_INVALID_ARG;
    }

    vTaskDelay(pdMS_TO_TICKS(10));  // the messages cannot be too close together

    SchedulerMetadata frame;

    if (auto maybe_frame = frame_queue[channel]->dequeue(std::chrono::milliseconds(FRAME_DEQUEUE_TIMEOUT_MS))) {
        frame = *maybe_frame;
    } else {
        // ESP_LOGI(DEBUG_LINK_TAG, "Scheduler queue for channel %d is empty", channel);
        return ESP_OK;
    }

    if (frame.data == nullptr){
        ESP_LOGE(DEBUG_LINK_TAG, "Data array does not exist");
        return ESP_ERR_INVALID_ARG;
    }

    if (this_board_id == PC_ADDR){
        ESP_LOGE(DEBUG_LINK_TAG, "This board is not assigned a board id");
        return ESP_ERR_INVALID_ARG;
    }

    esp_err_t res;
    bool isControlFrame = IS_CONTROL_FRAME(static_cast<uint8_t>(frame.header.type_flag));

    size_t frame_size = isControlFrame ? sizeof(ControlFrame) : sizeof(GenericFrame);
    uint8_t send_data[frame_size];

    // ESP_LOGI(DEBUG_LINK_TAG, "Sending frame on channel %d, type: %s", channel, isControlFrame ? "Control" : "Generic");

    if (isControlFrame){
        //control frame

        res = create_control_frame(frame.data->data(), frame.data->size(),
            make_control_frame_from_header(frame.header), send_data, &frame_size);

        if (res != ESP_OK){
            ESP_LOGE(DEBUG_LINK_TAG, "Failed to create control frame");
            return res;
        }

        // ESP_LOGI(DEBUG_LINK_TAG, "Sending control frame with %d bytes", frame_size);
        return scheduler_send_rmt(channel, frame, send_data, frame_size, false);
    } else {
        //generic frame
        if (frame.data->size() > (MAX_GENERIC_DATA_LEN)){
            //fragment here

            if (frame.timeout == 0){
                frame.timeout = GENERIC_FRAME_MIN_TIMEOUT;
            } else {
                frame.timeout--;
                res = push_frame_to_scheduler(frame, channel);
                if (res != ESP_OK){
                    ESP_LOGE(DEBUG_LINK_TAG, "Failed to schedule next generic frame fragment");
                    return res;
                }
                return ESP_OK;
            }

            if (static_cast<FrameType>(GET_TYPE(frame.header.type_flag)) != FrameType::MISC_UDP_GENERIC_TYPE){
                FrameAckRecord record = {
                    .last_ack = 0,
                    .total_frags = 0
                };

                res = get_record_sliding_window(channel, frame.header.receiver_id, frame.header.seq_num, &record);

                if (res != ESP_OK){
                    ESP_LOGE(DEBUG_LINK_TAG, "Failed to get sliding window ack record for board id %d seq num %d", frame.header.receiver_id, frame.header.seq_num);
                    return res;
                }

                if (record.last_ack == 0 && record.total_frags == 0){
                    //no acks has arrived yet - can only send fragment 1 ... fragment GENERIC_FRAME_SLIDING_WINDOW_SIZE (inclusive)
                    // ESP_LOGI(DEBUG_LINK_TAG, "no ack received yet for board id %d seq num %d", frame.header.receiver_id, frame.header.seq_num);
                    if (frame.curr_fragment < GENERIC_FRAME_SLIDING_WINDOW_SIZE){
                        frame.curr_fragment++;
                    } else {
                        frame.curr_fragment = 1;
                    }
                } else {
                    //some acks has arrived
                    // ESP_LOGI(DEBUG_LINK_TAG, "some ack received for board id %d seq num %d last ack %d total frags %d", frame.header.receiver_id, frame.header.seq_num, record.last_ack, record.total_frags);

                    //check if all acks are received
                    frame.last_ack = record.last_ack;
                    if (record.last_ack == record.total_frags){
                        //all acks received, can simply exit
                        // ESP_LOGI(DEBUG_LINK_TAG, "All acks recevied for board id %d seq num %d", frame.header.receiver_id, frame.header.seq_num);
                        complete_record_sliding_window(channel, frame.header.receiver_id, frame.header.seq_num);
                        return ESP_OK;
                    }

                    if (frame.curr_fragment - frame.last_ack < GENERIC_FRAME_SLIDING_WINDOW_SIZE && frame.curr_fragment - frame.last_ack < record.total_frags){
                        frame.curr_fragment++;
                    } else {
                        frame.curr_fragment = frame.last_ack;
                    }
                }
            } else {
                frame.curr_fragment++;
            }

            // ESP_LOGI(DEBUG_LINK_TAG, "current fragment to be sent for seq num %d is %d", frame.header.seq_num, frame.curr_fragment);

            //calculate data offset from curr_fragment
            uint16_t fragment_size = 0;

            if (frame.curr_fragment != (frame.header.frag_info >> 16)) {
                fragment_size = MAX_GENERIC_DATA_LEN;
            } else {
                fragment_size = frame.data->size() - (MAX_GENERIC_DATA_LEN * (frame.curr_fragment-1));
            }

            uint16_t curr_offset = MAX_GENERIC_DATA_LEN * (frame.curr_fragment - 1);

            // ESP_LOGI(DEBUG_LINK_TAG, "frame %d curr offset %d\n", frame.header.seq_num, curr_offset);
            // ESP_LOGI(DEBUG_LINK_TAG, "frame %d fragment size %d\n", frame.header.seq_num, fragment_size);

            frame.header.frag_info = (frame.header.frag_info & 0xFFFF0000) | frame.curr_fragment; //increment frag_num
            //create fragment
            res = create_generic_frame(frame.data->data(), fragment_size,
                make_generic_frame_from_header(frame.header), curr_offset, send_data, &frame_size);

            if (res != ESP_OK){
                ESP_LOGE(DEBUG_LINK_TAG, "Failed to create generic frame fragment");
                return res;
            }

            res = scheduler_send_rmt(channel, frame, send_data, frame_size, true);

            if (res != ESP_OK){
                ESP_LOGE(DEBUG_LINK_TAG, "Failed to send generic frame fragment");
                res = push_frame_to_scheduler(frame, channel);
                if (res != ESP_OK){
                    ESP_LOGE(DEBUG_LINK_TAG, "Failed to schedule next generic frame fragment");
                }
                return res;
            }

            //need to schedule the next fragment (if total_frags != frag_num)
            if ((frame.header.frag_info >> 16) > (frame.header.frag_info & 0xFFFF) || (frame.last_ack != (frame.header.frag_info >> 16) &&
            static_cast<FrameType>(GET_TYPE(frame.header.type_flag)) != FrameType::MISC_UDP_GENERIC_TYPE)){
                // frame.generic_frame_data_offset += fragment_size;
                // ESP_LOGI(DEBUG_LINK_TAG, "scheduling frame %d with frag_info 0x%X", frame.header.seq_num, frame.header.frag_info);
                res = push_frame_to_scheduler(frame, channel);
                if (res != ESP_OK){
                    ESP_LOGE(DEBUG_LINK_TAG, "Failed to schedule next generic frame fragment");
                    return res;
                }
            } else {
                //Done fragmenting, can free data array
                // ESP_LOGI(DEBUG_LINK_TAG, "finished fragmenting seq num %d frag_info 0x%X", frame.header.seq_num, frame.header.frag_info);
            }

        } else {
            //no fragmenting
            res = create_generic_frame(frame.data->data(), frame.data->size(),
                make_generic_frame_from_header(frame.header), 0, send_data, &frame_size);

            if (res != ESP_OK){
                ESP_LOGE(DEBUG_LINK_TAG, "Failed to create generic frame");
                return res;
            }
            return scheduler_send_rmt(channel, frame, send_data, frame_size, false);
        }
    }

    return ESP_OK;
}

esp_err_t DataLinkManager::scheduler_send_rmt(uint8_t channel, SchedulerMetadata frame, uint8_t* send_data, size_t frame_size, bool wait_for_tx_done){
    esp_err_t res;
    uint8_t channel_to_route = MAX_CHANNELS;

    // config is physical-layer specific; each IPhysicalLayer implementation
    // uses its own internal configuration set at init time, so nullptr is fine.
    if (frame.header.receiver_id == BROADCAST_ADDR){
        res = phys_comms->send(send_data, frame_size, nullptr, channel);
    } else {
        res = route_frame(frame.header.receiver_id, &channel_to_route);

        if (res != ESP_OK){
            ESP_LOGE(DEBUG_LINK_TAG, "Failed to find entry for %d", frame.header.receiver_id);
            return ESP_FAIL;
        }
        res = phys_comms->send(send_data, frame_size, nullptr, channel_to_route);
    }

    if (res != ESP_OK){
        ESP_LOGE(DEBUG_LINK_TAG, "Failed to send message");
        return ESP_FAIL;
    }

    return ESP_OK;
}

/**
 * @brief Increases the head of the sliding window associated with the board id and sequence number
 *
 * @param channel
 * @param board_id Receiving Board ID (the board who ACK'd)
 * @param seq_num
 * @param ack_record
 * @return esp_err_t
 */
esp_err_t DataLinkManager::inc_head_sliding_window(uint8_t channel, uint8_t board_id, uint16_t seq_num, FrameAckRecord* ack_record){
    if (ack_record == NULL){
        return ESP_ERR_INVALID_ARG;
    }

    if (ack_record->total_frags == 0 || ack_record->total_frags > MAX_GENERIC_NUM_FRAG
        || ack_record->last_ack == 0 || ack_record->total_frags < ack_record->last_ack){
        return ESP_ERR_INVALID_ARG;
    }

    if (sliding_window_mutex[channel] == NULL){
        return ESP_ERR_INVALID_STATE;
    }

    if (xSemaphoreTake(sliding_window_mutex[channel], pdMS_TO_TICKS(SLIDING_WINDOW_MUTEX_TIMEOUT_MS)) != pdTRUE){
        return ESP_ERR_TIMEOUT;
    }

    FrameAckRecord& record = sliding_window[channel][board_id][seq_num];

    if (record.last_ack > ack_record->last_ack){
        xSemaphoreGive(sliding_window_mutex[channel]);
        return ESP_ERR_INVALID_ARG;
    }

    record.last_ack = ack_record->last_ack;
    if (record.total_frags == 0){
        record.total_frags = ack_record->total_frags;
    }

    xSemaphoreGive(sliding_window_mutex[channel]);

    return ESP_OK;
}

/**
 * @brief Gets the current record associated with the board id and sequence number from the sliding window
 *
 * @param channel
 * @param board_id Receiving Board ID (the board who ACK'd)
 * @param seq_num
 * @param ack_record
 * @return esp_err_t
 */
esp_err_t DataLinkManager::get_record_sliding_window(uint8_t channel, uint8_t board_id, uint16_t seq_num, FrameAckRecord* ack_record){
    if (ack_record == NULL){
        return ESP_ERR_INVALID_ARG;
    }

    if (sliding_window_mutex[channel] == NULL){
        return ESP_ERR_INVALID_STATE;
    }

    if (xSemaphoreTake(sliding_window_mutex[channel], pdMS_TO_TICKS(SLIDING_WINDOW_MUTEX_TIMEOUT_MS)) != pdTRUE){
        return ESP_ERR_TIMEOUT;
    }

    if (sliding_window[channel][board_id].find(seq_num) == sliding_window[channel][board_id].end()){
        //record for this board id + seq number doesn't exist -- we don't want to create one
        xSemaphoreGive(sliding_window_mutex[channel]);
        ack_record->last_ack = 0;
        ack_record->total_frags = 0;
        return ESP_OK;
    }

    FrameAckRecord& record = sliding_window[channel][board_id][seq_num];

    ack_record->last_ack = record.last_ack;
    ack_record->total_frags = record.total_frags;

    xSemaphoreGive(sliding_window_mutex[channel]);

    return ESP_OK;
}

/**
 * @brief Removes the board id + sequence number record fromt the sliding window (map)
 *
 * @param channel
 * @param board_id Receiving Board ID (the board who ACK'd)
 * @param seq_num
 * @return esp_err_t
 */
esp_err_t DataLinkManager::complete_record_sliding_window(uint8_t channel, uint8_t board_id, uint16_t seq_num){
    if (sliding_window_mutex[channel] == NULL){
        return ESP_ERR_INVALID_STATE;
    }

    if (xSemaphoreTake(sliding_window_mutex[channel], pdMS_TO_TICKS(SLIDING_WINDOW_MUTEX_TIMEOUT_MS)) != pdTRUE){
        return ESP_ERR_TIMEOUT;
    }

    if (sliding_window[channel][board_id].find(seq_num) == sliding_window[channel][board_id].end()){
        //record for this board id + seq number doesn't exist -- we don't want to create one
        xSemaphoreGive(sliding_window_mutex[channel]);
        return ESP_ERR_INVALID_STATE;
    }

    sliding_window[channel][board_id].erase(seq_num);

    xSemaphoreGive(sliding_window_mutex[channel]);

    return ESP_OK;
}