/*
 *  puzzlecmd.cpp
 *
 *  Command-line program to test PuzzleFW firmware.
 *
 *  Joris van Rantwijk 2024
 */

#include <limits.h>
#include <stdio.h>
#include <chrono>
#include <optional>
#include <string>

#include <getopt.h>
#include <boost/asio.hpp>

#include "puzzlefw.hpp"
#include "logging.hpp"
#include "interrupt_manager.hpp"
#include "data_server.hpp"


/** Convert TriggerMode to string description. */
std::string trigger_mode_to_string(puzzlefw::TriggerMode mode)
{
    using puzzlefw::TriggerMode;
    switch (mode) {
        case TriggerMode::TRIG_AUTO: return "auto";
        case TriggerMode::TRIG_EXTERNAL: return "external";
        case TriggerMode::TRIG_EXTERNAL_ONCE: return "external-once";
        default: return "none";
    }
}


/** Show firmware status. */
void show_status(puzzlefw::PuzzleFwDevice& device)
{
    printf("Status:\n");

    printf("    timestamp =        %llu\n",
           (unsigned long long)device.get_timestamp());

    for (unsigned int i = 0; i < device.get_analog_channel_count(); i++) {
        unsigned int sample, min_sample, max_sample;
        sample = device.get_adc_sample(i);
        device.get_adc_range(i, min_sample, max_sample);
        printf("    channel %u =       %5u  (min = %u, max = %u)\n",
               i, sample, min_sample, max_sample);
    }

    uint32_t digital_state = device.get_digital_input_state();
    printf("    digital input =    %u %u %u %u\n",
           digital_state & 1,
           (digital_state >> 1) & 1,
           (digital_state >> 2) & 1,
           (digital_state >> 3) & 1);

    printf("    acquisition =      %s\n",
           device.is_acquisition_enabled() ? "on" : "off");

    printf("    channel mode =     %d channels\n",
           device.is_4channel_mode() ? 4 : 2);

    printf("    trigger mode =     %s, ch=%u, edge=%s\n",
           trigger_mode_to_string(device.get_trigger_mode()).c_str(),
           device.get_trigger_ext_channel(),
           device.get_trigger_ext_falling() ? "falling" : "rising");

    printf("    trigger delay =    %u * 8 ns\n",
           device.get_trigger_delay());

    printf("    record length =    %u samples\n",
           device.get_record_length());

    unsigned int divisor = device.get_decimation_factor();
    printf("    rate divisor =     %u (%u Sa/s)\n",
           divisor, 125000000 / divisor);

    printf("    averaging =        %s\n",
           device.is_averaging_enabled() ? "on" : "off");

    printf("    shift steps =      %u\n",
           device.get_shift_steps());

    printf("    timetagger mask =  0x%02x\n",
           device.get_timetagger_event_mask());

    printf("    ADC simulation =   %s\n",
           device.is_adc_simulation_enabled() ? "on" : "off");

    if (device.is_digital_simulation_enabled()) {
        uint32_t simulation_state = device.get_digital_simulation_state();
        printf("    digital sim =      %u %u %u %u\n",
               simulation_state & 1,
               (simulation_state >> 1) & 1,
               (simulation_state >> 2) & 1,
               (simulation_state >> 3) & 1);
    } else {
        printf("    digital sim =      off\n");
    }

    uint32_t dma_status = device.get_dma_status();
    printf("    DMA status =       %s, %s%s%s%s\n",
           device.is_dma_enabled() ? "enabled" : "disabled",
           (dma_status & 1) ? "busy" : "idle",
           (dma_status & 2) ? ", READ ERROR" : "",
           (dma_status & 4) ? ", WRITE ERROR" : "",
           (dma_status & 8) ? ", ADDRESS ERROR" : "");
}


/** Run TCP data server. */
void run_data_server(puzzlefw::PuzzleFwDevice& device)
{
    namespace asio = boost::asio;
    using namespace puzzlefw;

    asio::io_context io;

    // Catch Ctrl-C for controlled shut down.
    asio::signal_set signals(io, SIGINT);
    signals.async_wait(
        [&io](auto ec, int sig) {
            log(LOG_INFO, "Got SIGINT, stopping server");
            io.stop();
        });

    // Reserve 3/4 of the DMA buffer for analog acquisition data.
    // Reserve 1/4 of the DMA buffer for timetagger data.
    size_t acq_buf_size = 3 * 4096 * (device.dma_buffer_size() / 4096 / 4);
    size_t timetagger_buf_size = device.dma_buffer_size() - acq_buf_size;

    DmaWriteStream acq_stream(
        device,
        DmaWriteStream::DMA_ACQ,
        0,
        acq_buf_size);
    DmaWriteStream timetagger_stream(
        device,
        DmaWriteStream::DMA_TT,
        acq_buf_size,
        acq_buf_size + timetagger_buf_size);

    DataServer acq_server(io, acq_stream, 5001);
    DataServer timetagger_server(io, timetagger_stream, 5002);

    InterruptManager interrupt_manager(io, device);
    interrupt_manager.add_callback(
        [&acq_server](){ acq_server.handle_interrupt(); });
    interrupt_manager.add_callback(
        [&timetagger_server](){ timetagger_server.handle_interrupt(); });

    DmaErrorMonitor dma_error_monitor(
        io,
        device,
        std::chrono::milliseconds(100));

    // Disable DMA engine on exit from this function.
    struct ScopeGuard {
        PuzzleFwDevice& m_device;
        ScopeGuard(PuzzleFwDevice& device) : m_device(device) { }
        ~ScopeGuard() { m_device.set_dma_enabled(false); }
    } scope_guard(device);

    // Clear DMA errors, then enable DMA engine.
    device.clear_dma_errors();
    device.set_dma_enabled(true);

    // Enable data servers.
    acq_server.start_server();
    timetagger_server.start_server();

    log(LOG_INFO, "Running, press Ctrl-C to stop");
    io.run();
}


/** Parse integer. */
int parse_int(const char *arg, bool& ok)
{
    if (*arg == 0) {
        ok = false;
        return 0;
    }
    char *end;
    long v = strtol(arg, &end, 0);
    if (*end != 0) {
        ok = false;
        return 0;
    }
    if (v < INT_MIN || v > INT_MAX) {
        ok = false;
        return 0;
    }
    ok = true;
    return v;
}


int main(int argc, char **argv)
{
    using puzzlefw::PuzzleFwDevice;
    using puzzlefw::VersionInfo;

    enum Option {
        OPT_SHOW = 1,
        OPT_CLEAR_DMA, OPT_CLEAR_TIMESTAMP, OPT_CLEAR_RANGE,
        OPT_ACQUISITION_ON, OPT_ACQUISITION_OFF, OPT_CHANNELS,
        OPT_TRIGGER_NONE, OPT_TRIGGER_AUTO, OPT_TRIGGER_EXT,
        OPT_TRIGGER_EXT_ONCE,
        OPT_TRIGGER_CHANNEL, OPT_TRIGGER_RISING, OPT_TRIGGER_FALLING,
        OPT_TRIGGER_DELAY, OPT_TRIGGER,
        OPT_RECORD_LEN, OPT_DIVISOR, OPT_AVERAGE, OPT_DECIMATE, OPT_SHIFT,
        OPT_TIMETAGGER_CHANNELS, OPT_MARKER,
        OPT_ADC_SIM, OPT_DIG_SIM,
        OPT_SERVER
    };

    static const struct option options[] = {
        {"help",              0, 0, 'h'},
        {"show",              0, 0, OPT_SHOW},
        {"clear-dma",         0, 0, OPT_CLEAR_DMA},
        {"clear-timestamp",   0, 0, OPT_CLEAR_TIMESTAMP},
        {"clear-range",       0, 0, OPT_CLEAR_RANGE},
        {"acquisition-on",    0, 0, OPT_ACQUISITION_ON},
        {"acquisition-off",   0, 0, OPT_ACQUISITION_OFF},
        {"channels",          1, 0, OPT_CHANNELS},
        {"trigger-none",      0, 0, OPT_TRIGGER_NONE},
        {"trigger-auto",      0, 0, OPT_TRIGGER_AUTO},
        {"trigger-ext",       0, 0, OPT_TRIGGER_EXT},
        {"trigger-ext-once",  0, 0, OPT_TRIGGER_EXT_ONCE},
        {"trigger-channel",   1, 0, OPT_TRIGGER_CHANNEL},
        {"trigger-rising",    0, 0, OPT_TRIGGER_RISING},
        {"trigger-falling",   0, 0, OPT_TRIGGER_FALLING},
        {"trigger-delay",     1, 0, OPT_TRIGGER_DELAY},
        {"trigger",           0, 0, OPT_TRIGGER},
        {"record-len",        1, 0, OPT_RECORD_LEN},
        {"divisor",           1, 0, OPT_DIVISOR},
        {"average",           0, 0, OPT_AVERAGE},
        {"decimate",          0, 0, OPT_DECIMATE},
        {"shift",             1, 0, OPT_SHIFT},
        {"timetagger",        1, 0, OPT_TIMETAGGER_CHANNELS},
        {"marker",            0, 0, OPT_MARKER},
        {"adc-sim",           1, 0, OPT_ADC_SIM},
        {"dig-sim",           1, 0, OPT_DIG_SIM},
        {"server",            0, 0, OPT_SERVER},
        {nullptr,             0, 0, 0}
    };

    const char *usage_text =
        "Usage: %s {options...}\n"
        "Try '--help' for more information.\n";

    const char *help_text =
        "Usage: %s {options...}\n"
        "Command-line program to test PuzzleFW firmware.\n"
        "\n"
        "Options:\n"
        "  --help, -h            Show this help message.\n"
        "  --show                Show current firmware state.\n"
        "  --clear-dma           Stop DMA and clear DMA errors.\n"
        "  --clear-timestamp     Reset timestamp counter to zero.\n"
        "  --clear-range         Clear ADC min/max sample monitor.\n"
        "  --acquisition-on      Enable analog acquisition.\n"
        "  --acquisition-off     Disable analog acquisition.\n"
        "  --channels 2|4        Select 2-channel or 4-channel mode.\n"
        "  --trigger-none        Disable triggering.\n"
        "  --trigger-auto        Enable continuous triggering.\n"
        "  --trigger-ext         Enable external trigger.\n"
        "  --trigger-ext-once    Enable external trigger once.\n"
        "  --trigger-channel N   Select trigger input channel (range 0 to 3)."
        "\n"
        "  --trigger-rising      Trigger on rising edge.\n"
        "  --trigger-falling     Trigger on falling edge.\n"
        "  --trigger-delay N     Set trigger delay to N * 8 ns.\n"
        "  --trigger             Send a manual trigger event.\n"
        "  --record-len N        Set record length to N samples per trigger."
        "\n"
        "  --divisor N           Set sample rate to 125 MSa/s divided by N."
        "\n"
        "  --average             Reduce sample rate by summing samples.\n"
        "  --decimate            Reduce sample rate by decimating samples.\n"
        "  --shift N             Shift sample values right by N positions.\n"
        "  --timetagger MASK     Set bit mask of enabled timetagger events.\n"
        "                        0 = all events disabled\n"
        "                        1 = enable rising edge on channel 0\n"
        "                        2 = enable falling edge on channel 0\n"
        "                        4 = enable rising edge on channel 1\n"
        "                        ... 255 = all event types enabled\n"
        "  --marker              Emit a marker in the timetagger stream.\n"
        "  --adc-sim 0|1         Enable (1) or disable (0) ADC simulation.\n"
        "  --dig-sim MASK|off    Set simulated digital input state as bitmask"
        "\n"
        "                        of channels, or disable simulation.\n"
        "  --server              Run TCP server to send data.\n"
        "\n";

    struct {
        bool                show;
        bool                clear_dma;
        bool                clear_timestamp;
        bool                clear_range;
        std::optional<bool> acquisition_en;
        std::optional<int>  channels;
        std::optional<puzzlefw::TriggerMode> trigger_mode;
        std::optional<int>  trigger_channel;
        std::optional<bool> trigger_falling;
        std::optional<int>  trigger_delay;
        bool                trigger_force;
        std::optional<int>  record_len;
        std::optional<int>  divisor;
        std::optional<bool> averaging_en;
        std::optional<int>  shift_steps;
        std::optional<unsigned int> timetagger_channels;
        bool                marker;
        std::optional<bool> adc_sim;
        std::optional<int>  dig_sim;
        bool                server;
    } args = {};

    while (1) {
        int c = getopt_long(argc, argv, "h", options, nullptr);
        if (c == -1) {
            break;
        }

        bool ok;
        switch (c) {
            case 'h':
                printf(help_text, argv[0]);
                return 0;
            case OPT_SHOW:
                args.show = true;
                break;
            case OPT_CLEAR_DMA:
                args.clear_dma = true;
                break;
            case OPT_CLEAR_TIMESTAMP:
                args.clear_timestamp = true;
                break;
            case OPT_CLEAR_RANGE:
                args.clear_range = true;
                break;
            case OPT_ACQUISITION_ON:
                args.acquisition_en = true;
                break;
            case OPT_ACQUISITION_OFF:
                args.acquisition_en = false;
                break;
            case OPT_CHANNELS:
                args.channels = parse_int(optarg, ok);
                if (!ok || (args.channels != 2 && args.channels != 4)) {
                    fprintf(stderr, "ERROR: Invalid value for --channels\n");
                    return 1;
                }
                break;
            case OPT_TRIGGER_NONE:
                args.trigger_mode = puzzlefw::TRIG_NONE;
                break;
            case OPT_TRIGGER_AUTO:
                args.trigger_mode = puzzlefw::TRIG_AUTO;
                break;
            case OPT_TRIGGER_EXT:
                args.trigger_mode = puzzlefw::TRIG_EXTERNAL;
                break;
            case OPT_TRIGGER_EXT_ONCE:
                args.trigger_mode = puzzlefw::TRIG_EXTERNAL_ONCE;
                break;
            case OPT_TRIGGER_CHANNEL:
                args.trigger_channel = parse_int(optarg, ok);
                if (!ok
                        || args.trigger_channel < 0
                        || args.trigger_channel > 3) {
                    fprintf(stderr,
                            "ERROR: Invalid value for --trigger-channel\n");
                    return 1;
                }
                break;
            case OPT_TRIGGER_RISING:
                args.trigger_falling = false;
                break;
            case OPT_TRIGGER_FALLING:
                args.trigger_falling = true;
                break;
            case OPT_TRIGGER_DELAY:
                args.trigger_delay = parse_int(optarg, ok);
                if (!ok
                    || args.trigger_delay < 0
                    || args.trigger_delay > PuzzleFwDevice::MAX_TRIGGER_DELAY)
                {
                    fprintf(stderr,
                            "ERROR: Invalid value for --trigger-delay\n");
                    return 1;
                }
                break;
            case OPT_TRIGGER:
                args.trigger_force = true;
                break;
            case OPT_RECORD_LEN:
                args.record_len = parse_int(optarg, ok);
                if (!ok
                    || args.record_len < 1
                    || args.record_len > PuzzleFwDevice::MAX_RECORD_LENGTH) {
                    fprintf(stderr,
                            "ERROR: Invalid value for --record-len\n");
                    return 1;
                }
                break;
            case OPT_DIVISOR:
                args.divisor = parse_int(optarg, ok);
                if (!ok
                    || args.divisor < 1
                    || args.divisor > PuzzleFwDevice::MAX_DECIMATION_FACTOR) {
                    fprintf(stderr,
                            "ERROR: Invalid value for --divisor\n");
                    return 1;
                }
                break;
            case OPT_AVERAGE:
                args.averaging_en = true;
                break;
            case OPT_DECIMATE:
                args.averaging_en = false;
                break;
            case OPT_SHIFT:
                args.shift_steps = parse_int(optarg, ok);
                if (!ok || args.shift_steps < 0 || args.shift_steps > 8) {
                    fprintf(stderr,
                            "ERROR: Invalid value for --shift-steps\n");
                    return 1;
                }
                break;
            case OPT_TIMETAGGER_CHANNELS:
                args.timetagger_channels = parse_int(optarg, ok);
                if (!ok
                        || args.timetagger_channels < 0
                        || args.timetagger_channels > 255) {
                    fprintf(stderr,
                            "ERROR: Invalid value for --timetagger\n");
                    return 1;
                }
                break;
            case OPT_MARKER:
                args.marker = true;
                break;
            case OPT_ADC_SIM:
                if (std::string("0") == optarg
                        || std::string("off") == optarg) {
                    args.adc_sim = false;
                } else if (std::string("1") == optarg
                        || std::string("on") == optarg) {
                    args.adc_sim = true;
                } else {
                    fprintf(stderr,
                            "ERROR: Invalid value for --adc-sim\n");
                    return 1;
                }
                break;
            case OPT_DIG_SIM:
                if (std::string("off") == optarg) {
                    args.dig_sim = -1;
                } else {
                    args.dig_sim = parse_int(optarg, ok);
                    if (!ok || args.dig_sim < 0 || args.dig_sim > 15) {
                        fprintf(stderr,
                                "ERROR: Invalid value for --dig-sim\n");
                        return 1;
                    }
                }
                break;
            case OPT_SERVER:
                args.server = true;
                break;
            default:
                fprintf(stderr, usage_text, argv[0]);
                return 1;
        }
    }

    if (optind < argc) {
        fprintf(stderr, "ERROR: Unexpected positional argument\n");
        fprintf(stderr, usage_text, argv[0]);
        return 1;
    }

    try {
        PuzzleFwDevice device;

        VersionInfo version = device.get_version_info();
        printf("Detected PuzzleFW firmware version %d.%d\n",
               version.major_version, version.minor_version);
        printf("    %u analog input channels\n",
               device.get_analog_channel_count());
        printf("    DMA buffer size: %zu bytes\n",
               device.dma_buffer_size());

        if (args.clear_dma) {
            printf("Disabling and clearing DMA engine ...\n");
            device.set_dma_enabled(false);
            device.clear_dma_errors();
        }

        if (args.clear_timestamp) {
            printf("Resetting timestamp counter ...\n");
            device.clear_timestamp();
        }

        if (args.clear_range) {
            printf("Resetting ADC min/max sample monitor ...\n");
            device.clear_adc_range();
        }

        if (args.acquisition_en.has_value()) {
            printf("%s analog acquisition ...\n",
                   args.acquisition_en.value() ? "Enabling" : "Disabling");
            device.set_acquisition_enabled(args.acquisition_en.value());
        }

        if (args.record_len.has_value()) {
            printf("Setting record length to %d samples per trigger ...\n",
                   args.record_len.value());
            device.set_record_length(args.record_len.value());
        }

        if (args.divisor.has_value()) {
            printf("Selecting sample rate divisor %d ...\n",
                   args.divisor.value());
            device.set_decimation_factor(args.divisor.value());
        }

        if (args.averaging_en.has_value()) {
            printf("Selecting downsampling via %s ...\n",
                   args.averaging_en.value() ? "averaging" : "decimation");
            device.set_averaging_enabled(args.averaging_en.value());
        }

        if (args.channels.has_value()) {
            printf("Selecting %d-channel mode ...\n", args.channels.value());
            device.set_4channel_mode(args.channels.value() == 4);
        }

        if (args.shift_steps.has_value()) {
            printf("Selecting %d right-shift steps ...\n",
                   args.shift_steps.value());
            device.set_shift_steps(args.shift_steps.value());
        }

        if (args.trigger_channel.has_value()) {
            printf("Selecting trigger channel %d ...\n",
                   args.trigger_channel.value());
            device.set_trigger_ext_channel(args.trigger_channel.value());
        }

        if (args.trigger_falling.has_value()) {
            printf("Selecting trigger on %s edge ...\n",
                   args.trigger_falling.value() ? "falling" : "rising");
            device.set_trigger_ext_falling(args.trigger_falling.value());
        }

        if (args.trigger_delay.has_value()) {
            printf("Selecting trigger delay %d * 8 ns ...\n",
                   args.trigger_delay.value());
            device.set_trigger_delay(args.trigger_delay.value());
        }

        if (args.trigger_mode.has_value()) {
            printf("Selecting trigger mode %s ...\n",
                   trigger_mode_to_string(args.trigger_mode.value()).c_str());
            device.set_trigger_mode(args.trigger_mode.value());
        }

        if (args.trigger_force) {
            printf("Sending forced trigger ...\n");
            device.trigger_force();
        }

        if (args.timetagger_channels.has_value()) {
            printf("Setting timetagger channel mask 0x%02x ...\n",
                   args.timetagger_channels.value());
            device.set_timetagger_event_mask(
                args.timetagger_channels.value());
        }

        if (args.marker) {
            printf("Emitting timetagger marker record ...\n");
            device.timetagger_mark();
        }

        if (args.adc_sim.has_value()) {
            printf("%s ADC simulation ...\n",
                   args.adc_sim.value() ? "Enabling" : "Disabling");
            device.set_adc_simulation_enabled(args.adc_sim.value());
        }

        if (args.dig_sim.has_value()) {
            if (args.dig_sim.value() < 0) {
                printf("Disabling digital input simulation ...\n");
                device.set_digital_simulation_enabled(false);
            } else {
                printf("Setting digital input simulation state 0x%x ...\n",
                       (unsigned int)args.dig_sim.value());
                device.set_digital_simulation_state(args.dig_sim.value());
                device.set_digital_simulation_enabled(true);
            }
        }

        if (args.show) {
            show_status(device);
        }

        if (args.server) {
            run_data_server(device);
        }

    } catch (std::exception& e) {
        fprintf(stderr, "ERROR: %s\n", e.what());
        return 1;
    }

    return 0;
}

/* end */