
#include "src/pico_server_dispatcher.hpp"

extern SimpleLogger logger;

int pico_server_dispatcher::pico_stop()
{
    auto retval = ps4000aStop(handle);
    if (retval != 0)
    {
        sendErrorMsg(0x00000003, "Failed to stop ADC!");
        return -1;
    }

    pk_ptr->pack(magic_number);
    pk_ptr->pack(0xCC);

    return sendValidatedMsgpackData();
}

int pico_server_dispatcher::pico_status()
{
    sendStatusMsg(0x00000000, "OK");
    return 1;
}

int pico_server_dispatcher::pico_open()
{
    auto retval = ps4000aOpenUnit(&handle, NULL);
    if(retval != 0)
    {
        sendErrorMsg(retval, "Error while open picoscope!");
        return -1;
    }

    logger << LogPref::Flag(INFO) << "Pico device opened with handle: " << handle << endl;

    int16_t start{10};
    retval = ps4000aFlashLed(handle, start);
    if(retval != 0)
    {
        sendErrorMsg(retval, "Failed to flash ADC LED!");
        return -1;
    }

    pk_ptr->pack(magic_number);
    pk_ptr->pack(0xCC);

    return sendValidatedMsgpackData();
}

int pico_server_dispatcher::pico_close()
{
    auto retval = ps4000aCloseUnit(handle);
    if(retval != 0)
    {
        sendErrorMsg(retval, "Failed to flash ADC LED!");
        return -1;
    }

    pk_ptr->pack(magic_number);
    pk_ptr->pack(0xCC);

    return sendValidatedMsgpackData();
}

int pico_server_dispatcher::pico_get_params()
{
    pk_ptr->pack(magic_number);
    pk_ptr->pack(0xCD);
    pk_ptr->pack(this->sample_rate);
    pk_ptr->pack(this->measurement_times);
    pk_ptr->pack(this->number_channels);
    pk_ptr->pack((int) this->trig_channel);
    pk_ptr->pack((int) this->threshold);
    pk_ptr->pack((int) this->th_direction);
    pk_ptr->pack(this->trig_autoTrigger_ms);

    return sendValidatedMsgpackData();
}

int pico_server_dispatcher::pico_set_rate() 
{
    msgpack::object_handle oh;

    if(!upk_ptr->next(oh))
    {
        sendErrorMsg(0x10000001, "Invalid command data!");
        return -1;
    }

    this->sample_rate = oh.get().as<uint32_t>();

    pk_ptr->pack(magic_number);
    pk_ptr->pack(0xCC);

    return sendValidatedMsgpackData();
}

int pico_server_dispatcher::pico_set_points()
{
    msgpack::object_handle oh;

    if(!upk_ptr->next(oh))
    {
        sendErrorMsg(0x20000001, "Invalid command data!");
        return -1;
    }

    this->points_vector = oh.get().as<std::vector<uint32_t>>();
    this->points_value = *(std::max_element(points_vector.begin(), points_vector.end()));
    this->measurement_times = points_vector.size();

    pk_ptr->pack(magic_number);
    pk_ptr->pack(0xCC);

    return sendValidatedMsgpackData();
}

int pico_server_dispatcher::pico_set_times()
{
    msgpack::object_handle oh;

    if(!upk_ptr->next(oh))
    {
        sendErrorMsg(0x30000001, "Invalid command data!");
        return -1;
    }

    this->measurement_times = oh.get().as<uint32_t>();

    pk_ptr->pack(magic_number);
    pk_ptr->pack(0xCC);

    return sendValidatedMsgpackData();
}

int pico_server_dispatcher::pico_channel_configure() 
{
    msgpack::object_handle oh;

    if(!upk_ptr->next(oh))
    {
        sendErrorMsg(0x40000001, "Invalid command data!");
        return -1;
    }

    this->number_channels = oh.get().as<uint32_t>();
    if(this->number_channels < 1 && this->number_channels > 16)
    {
        sendErrorMsg(0x40000004, "Invalid number of channels!");
        return -1;
    }

    this->channels = create_channel(number_channels);
    conditions = new PS4000A_CONDITION[number_channels];
    for (uint32_t i = 0; i < number_channels; ++i) {
        conditions[i].source = channels[i];
        conditions[i].condition = PS4000A_CONDITION_FALSE;
    }

    logger << "Channels opened..." << endl;

    if(!upk_ptr->next(oh))
    {
        sendErrorMsg(0x40000002, "Invalid command data!");
        return -1;
    }

    std::vector<int> int_ranges = oh.get().as<std::vector<int>>();

    std::vector<enPicoConnectProbeRange> channel_ranges = {};
    for(int i = 0; i < this->number_channels; ++i)
    {
        channel_ranges.push_back(static_cast<enPicoConnectProbeRange>(int_ranges[i]));
    }

    if(!upk_ptr->next(oh))
    {
        sendErrorMsg(0x40000003, "Invalid command data!");
        return -1;
    }

    this->trig_channel = static_cast<PS4000A_CHANNEL>(oh.get().as<int>());

    auto retval = ps4000aSetChannel(handle, channels[this->trig_channel], true, PS4000A_DC, channel_ranges[this->trig_channel], 0);
    if(retval != 0)
    {
        sendErrorMsg(retval, "Failed to set trigger channel!");
        return -1;
    }

    logger << "Trigger channel set..." << endl;

    for(int i = 0; i < this->number_channels; i++)
    {
        if(i == this->trig_channel) continue;
        else {
            retval = ps4000aSetChannel(handle, channels[i], true, PS4000A_AC, channel_ranges[i], 0);
            if(retval != 0)
            {
                sendErrorMsg(retval, "Failed to set channel!");
                return -1;
            }
        }
    }

    logger << "Channels configured!" << endl;

    pk_ptr->pack(magic_number);
    pk_ptr->pack(0xCC);

    return sendValidatedMsgpackData();
}

int pico_server_dispatcher::pico_trigger_configure()
{
    msgpack::object_handle oh;

    if(!upk_ptr->next(oh))
    {
        sendErrorMsg(0x50000001, "Invalid command data!");
        return -1;
    }

    this->th_direction = static_cast<PS4000A_THRESHOLD_DIRECTION>(oh.get().as<int>());

    if(!upk_ptr->next(oh))
    {
        sendErrorMsg(0x50000002, "Invalid command data!");
        return -1;
    }

    this->trig_delay = 10;
    this->threshold = oh.get().as<int16_t>();

    if(!upk_ptr->next(oh))
    {
        sendErrorMsg(0x50000003, "Invalid command data!");
        return -1;
    }

    this->trig_autoTrigger_ms = oh.get().as<int16_t>();

    auto retval = ps4000aSetSimpleTrigger
    (
        handle,
        true,
        channels[this->trig_channel],
        this->threshold,
        this->th_direction,
        this->trig_delay,
        this->trig_autoTrigger_ms
    );

    if(retval != 0)
    {
        sendErrorMsg(retval, "Failed to configure trigger!");
        return -1;
    }

    pk_ptr->pack(magic_number);
    pk_ptr->pack(0xCC);

    return sendValidatedMsgpackData();
}

int pico_server_dispatcher::pico_set_premeasurement()
{
    msgpack::object_handle oh;
    if(!upk_ptr->next(oh))
    {
        sendErrorMsg(0x60000001, "Invalid command data!");
        return -1;
    }

    this->premeasurement_percentage = oh.get().as<uint32_t>();

    pk_ptr->pack(magic_number);
    pk_ptr->pack(0xCC);

    return sendValidatedMsgpackData();
}

int pico_server_dispatcher::pico_measurement()
{
    if(is_buffers_allocated)
    {
        for(size_t i = 0; i < measurement_times; ++i) {
            for(size_t j = 0; j < number_channels; ++j) {
                delete data_buffer_vec[i][j];
            }
        }
    }

    is_buffers_allocated = false;
    data_buffer_vec.assign(measurement_times, std::vector<int16_t*>(number_channels, nullptr));

    logger << LogPref::Flag(INFO) << "Get timebase:" << endl;
    uint32_t timebase = timebase_choice(sample_rate);
    logger << "Timebase = " << timebase << endl;

    logger << LogPref::Flag(INFO) << "Set data buffer:" << endl;
    for(size_t i = 0; i < measurement_times; ++i) {
        for(size_t j = 0; j < number_channels; ++j) {
            data_buffer_vec[i][j] = new int16_t[points_vector[i]];
            memset(data_buffer_vec[i][j], 0, points_vector[i]);
        }
    }
    is_buffers_allocated = true;

    int32_t* nSamples = new int32_t[measurement_times];

    auto retval = ps4000aMemorySegments(handle, measurement_times, nSamples);
    if(retval != 0)
    {
        sendErrorMsg(retval, "Failed to set memory segments!");
        return -1;
    }
    retval = ps4000aSetNoOfCaptures(handle, measurement_times);
    if(retval != 0)
    {
        sendErrorMsg(retval, "Failed to set memory segments!");
        return -1;
    }

    int32_t timeIntervalNanoseconds = 0;
    int32_t maxSamples = 0;
    int32_t segmentIndex = 0;

    retval = ps4000aGetTimebase(handle, timebase, points_value, &timeIntervalNanoseconds,
                                                    &maxSamples, segmentIndex);
    if(retval != 0)
    {
        sendErrorMsg(retval, "Failed to get timebase!");
        return -1;
    }

    int32_t noOfPreTriggerSamples = premeasurement_percentage;
    int32_t noOfPostTriggerSamples = points_value;

    logger << LogPref::Flag(INFO) << "Run block!" << endl;
    retval = ps4000aRunBlock(handle, noOfPreTriggerSamples, noOfPostTriggerSamples, timebase,
                                    nullptr, segmentIndex, nullptr, nullptr);
    if(retval != 0)
    {
        sendErrorMsg(retval, "Failed to run block!");
        return -1;
    }

    logger << LogPref::Flag(INFO) << "Run block done!" << endl;

    int16_t ready = 0;
    while (ready == 0) {
            retval = ps4000aIsReady(handle, &ready);
    }

    logger << LogPref::Flag(INFO) << "Data collection is starting!" << endl;

    for(size_t i = 0; i < measurement_times; ++i) {
        for(size_t j = 0; j < number_channels; ++j) {
            logger << LogPref::Flag(INFO) << "Set data buffer for channel " << j << endl;

            PS4000A_RATIO_MODE mode = PS4000A_RATIO_MODE_NONE;
            retval = ps4000aSetDataBuffer(handle, channels[j], data_buffer_vec[i][j], points_vector[i],
                                            i, mode);

            if(retval != 0) {
                logger << LogPref::Flag(ERROR) << "Failed to set data buffer for channel " << j << ". Code: " << retval << endl;
                sendErrorMsg(0x70000003, "Failed to set data buffer!");
                return -1;
            }
        }
    }

    logger << LogPref::Flag(INFO) << "Set data buffer done" << endl;

    logger << LogPref::Flag(INFO) << "Data collection start" << endl;
    uint32_t StartIndex = 0;
    uint32_t noOfSamples = static_cast<uint32_t>(points_value);
    int16_t overflow = 0;
    uint32_t DownSampleRatio = 1;

    retval = ps4000aGetValuesBulk(handle, &noOfSamples, 0, measurement_times-1, DownSampleRatio, PS4000A_RATIO_MODE_NONE, &overflow);
    if(retval != 0)
    {
        logger << LogPref::Flag(ERROR) << "Failed to get values. Code: " << retval << endl;
        sendErrorMsg(retval, "Failed to get values!");
        return -1;
    }

    logger << LogPref::Flag(INFO) << "Data collection done" << endl;
    logger << LogPref::Flag(INFO) << "Confirm data to socket" << endl;

    pk_ptr->pack(magic_number);
    pk_ptr->pack(0xCC);

    return sendValidatedMsgpackData();
}

std::vector<std::vector<int16_t>> pico_server_dispatcher::split_packets(int i, int j)
{
    std::vector<std::vector<int16_t>> packets = {};

    int n = ((float) points_vector[i]) / MAX_PACKET + 1;

    for(int k = 0; k < points_vector[i]; k = k + MAX_PACKET / sizeof(int16_t))
    {
        std::vector<int16_t> vec = {};

        if((k + MAX_PACKET / sizeof(int16_t)) > points_vector[i])
        {
            vec.assign(data_buffer_vec[i][j] + size_t(k), data_buffer_vec[i][j] + (size_t)(points_vector[i]));
            packets.push_back(vec);
        }
        else
        {
            vec.assign(data_buffer_vec[i][j] + (size_t)(k), data_buffer_vec[i][j] + (size_t)(k + MAX_PACKET / sizeof(int16_t)));
            packets.push_back(vec);
        }
    }

    return packets;
}

int pico_server_dispatcher::pico_send_points()
{
    logger << "Sending points" << endl;
    for(uint32_t i = 0; i < measurement_times; ++i)
    {
        for(uint32_t j = 0; j < number_channels; ++j)
        {
            logger << "Spliting packets..." << endl;
            std::vector<std::vector<int16_t>> packets = split_packets(i, j);
            logger << "Packets splited. No_Packets: " << packets.size() << endl;

            uint32_t num_packet = 1;
            uint32_t total_data = sizeof(int16_t) * points_vector[i];
            uint32_t total_packets = 1 + ((uint32_t) ((total_data * 1.0f) / MAX_PACKET));

            logger << "Ready to send data!" << endl;

            for(std::vector<int16_t> packet : packets)
            {
                #ifdef DEBUG
                    logger << "Packet: " << i << " " << j << ", size = " << packet.size() << endl;
                #endif
                pk_ptr->pack(magic_number);
                pk_ptr->pack(0xCD);
                pk_ptr->pack(i);
                pk_ptr->pack(j);
                pk_ptr->pack(num_packet);
                pk_ptr->pack(packets.size());
                pk_ptr->pack(packet);

                int ret = sendValidatedMsgpackData();
                if(ret != 1)
                {
                    sendErrorMsg(0x80000001, "Failed to send confirm packet!");
                    return -1;
                }

                #ifdef DEBUG
                    logger << "Data sent!" << endl;
                #endif

                ret = recvValidatedHeader();
                if(ret != 1)
                {
                    sendErrorMsg(0x80000002, "Failed to recv confirm packet!");
                    return -1;
                }

                if(command_code == 0x3D)
                {
                #ifdef DEBUG
                    logger << LogPref::Flag(INFO) << "Packet confirmed" << endl;
                #endif
                }
                else  {
                    sendErrorMsg(0x80000002, "Invalid confirm code!");
                    return -1;
                }

                num_packet++;
            }
        }
    }

    logger << "Data sended! Wait confirmation." << endl;
    pk_ptr->pack(magic_number);
    pk_ptr->pack(0xCE);

    int ret = sendValidatedMsgpackData();
    if(ret != 1)
    {
        sendErrorMsg(0x80000010, "Server failed to confirm packet!");
        return -1;
    }

    ret = recvValidatedHeader();
    if(ret != 1)
    {
        sendErrorMsg(0x80000011, "Client failed to confirm data!");
        return -1;
    }

    if(command_code == 0x3E)
    {
        logger << LogPref::Flag(INFO) << "Data confirmed" << endl;
    }
    else {
        return -1;
    }

    return 1;
}