#include "module_generic_analog_demod.h"
#include "common/audio/audio_sink.h"
#include "common/dsp/io/wav_writer.h"
#include "core/config.h"
#include "imgui/imgui.h"
#include "logger.h"
#include <volk/volk.h>

namespace generic_analog
{
    GenericAnalogDemodModule::GenericAnalogDemodModule(std::string input_file, std::string output_file_hint, nlohmann::json parameters) : BaseDemodModule(input_file, output_file_hint, parameters)
    {
        name = "Generic Analog Demodulator (WIP)";
        show_freq = false;
        play_audio = satdump::satdump_cfg.shouldPlayAudio();

        constellation.d_hscale = 1.0; // 80.0 / 100.0;
        constellation.d_vscale = 0.5; // 20.0 / 100.0;

        MIN_SPS = 1;
        MAX_SPS = 1e9;

        upcoming_symbolrate = d_symbolrate;
        if (d_parameters.contains("bandwidth"))
        {
            upcoming_symbolrate = d_parameters["bandwidth"];
            settings_changed = true;
        }

        if (d_parameters.contains("record_demod_audio"))
        {
            record_demod_audio = d_parameters["record_demod_audio"];
        }

        // modulation_type
        if (d_parameters.contains("modulation_type"))
        {
            const std::string &mod = d_parameters["modulation_type"];
            if (mod == "NFM")
            {
                modulation_type = ModulationType::NFM;
            }
            else if (mod == "AM")
            {
                modulation_type = ModulationType::AM;
            }
        }
    }

    void GenericAnalogDemodModule::init()
    {
        BaseDemodModule::initb();

        // Resampler to BW
        //    res = std::make_shared<dsp::RationalResamplerBlock<complex_t>>(agc->output_stream, d_symbolrate, final_samplerate);

        // Quadrature demod
        //  qua = std::make_shared<dsp::QuadratureDemodBlock>(res->output_stream, dsp::hz_to_rad(d_symbolrate / 2, d_symbolrate));
    }

    GenericAnalogDemodModule::~GenericAnalogDemodModule() {}

    void GenericAnalogDemodModule::process()
    {
        if (input_data_type == satdump::pipeline::DATA_FILE)
            filesize = file_source->getFilesize();
        else
            filesize = 0;

        const bool output_to_file = output_data_type == satdump::pipeline::DATA_FILE || record_demod_audio;

        std::string output_file_hint = d_output_file_hint;
        if (output_to_file)
        {
            // append satellite NORAD number to filename
            if (d_parameters.contains("satellite_norad"))
            {
                output_file_hint.push_back('_');
                output_file_hint.append(make_safe_string(std::to_string(d_parameters["satellite_norad"].get<nlohmann::json::number_unsigned_t>())));
            }

            // append satellite name to filename
            if (d_parameters.contains("satellite_name") && d_parameters["satellite_name"].get<nlohmann::json::string_t>().size() > 0)
            {
                output_file_hint.push_back('_');
                output_file_hint.append(make_safe_string(d_parameters["satellite_name"]));
            }

            // append satellite downlink frequency to filename
            if (d_parameters.contains("satellite_frequency"))
            {
                output_file_hint.push_back('_');
                output_file_hint.append(make_safe_string(std::to_string(d_parameters["satellite_frequency"].get<nlohmann::json::number_unsigned_t>()) + "Hz"));
            }

            data_out = std::ofstream(output_file_hint + ".wav", std::ios::binary);
            d_output_file = output_file_hint + ".wav";
        }

        logger->info("Using input baseband " + d_input_file);
        logger->info("Demodulating to " + output_file_hint + ".wav");
        logger->info("Buffer size : " + std::to_string(d_buffer_size));

        time_t lastTime = 0;

        // Start
        BaseDemodModule::start();
        //    res->start();
        //    qua->start();

        // Buffers to wav
        int16_t *output_wav_buffer = new int16_t[d_buffer_size * 100];
        int16_t *output_wav_buffer_resamp = new int16_t[d_buffer_size * 200];
        uint64_t final_data_size = 0;
        dsp::WavWriter wave_writer(data_out);
        if (output_to_file)
            wave_writer.write_header(audio_samplerate, 1);

        std::shared_ptr<audio::AudioSink> audio_sink;
        if (input_data_type != satdump::pipeline::DATA_FILE && audio::has_sink())
        {
            enable_audio = true;
            audio_sink = audio::get_default_sink();
            audio_sink->set_samplerate(audio_samplerate);
            audio_sink->start();
        }

        /////////////
        dsp::RationalResamplerBlock<complex_t> input_resamp(nullptr, d_symbolrate, final_samplerate);
        dsp::QuadratureDemodBlock quad_demod(nullptr, dsp::hz_to_rad(d_symbolrate / 2, d_symbolrate));
        complex_t *work_buffer_complex = dsp::create_volk_buffer<complex_t>(d_buffer_size);
        float *work_buffer_float = dsp::create_volk_buffer<float>(d_buffer_size);

        int dat_size = 0;
        while (demod_should_run())
        {
            dat_size = agc->output_stream->read();

            if (dat_size <= 0)
            {
                agc->output_stream->flush();
                continue;
            }

#if 1
            proc_mtx.lock();

            if (settings_changed)
            {
                if (upcoming_symbolrate > 0)
                {
                    d_symbolrate = upcoming_symbolrate;
                    input_resamp.set_ratio(d_symbolrate, final_samplerate);
                    quad_demod.set_gain(dsp::hz_to_rad(d_symbolrate / 2, d_symbolrate));
                }

                settings_changed = false;
            }

            int nout = input_resamp.process(agc->output_stream->readBuf, dat_size, work_buffer_complex);

            if (modulation_type == ModulationType::NFM)
                nout = quad_demod.process(work_buffer_complex, nout, work_buffer_float);
            else if (modulation_type == ModulationType::AM)
                volk_32fc_magnitude_32f((float *)work_buffer_float, (lv_32fc_t *)work_buffer_complex, nout);

            // Into const
            constellation.pushFloatAndGaussian(work_buffer_float, nout);

            for (int i = 0; i < nout; i++)
            {
                if (work_buffer_float[i] > 1.0f)
                    work_buffer_float[i] = 1.0f;
                if (work_buffer_float[i] < -1.0f)
                    work_buffer_float[i] = -1.0f;
            }

            volk_32f_s32f_convert_16i(output_wav_buffer, (float *)work_buffer_float, 32767, nout);

            int final_out = audio::AudioSink::resample_s16(output_wav_buffer, output_wav_buffer_resamp, d_symbolrate, audio_samplerate, nout, 1);
            if (enable_audio && play_audio)
                audio_sink->push_samples(output_wav_buffer_resamp, final_out);
            if (output_to_file)
            {
                data_out.write((char *)output_wav_buffer_resamp, final_out * sizeof(int16_t));
                final_data_size += final_out * sizeof(int16_t);
            }
            else
            {
                output_fifo->write((uint8_t *)output_wav_buffer_resamp, final_out * sizeof(int16_t));
            }

            proc_mtx.unlock();
#else
            // Into const
            constellation.pushFloatAndGaussian(qua->output_stream->readBuf, qua->output_stream->getDataSize());

            for (int i = 0; i < dat_size; i++)
            {
                if (qua->output_stream->readBuf[i] > 1.0f)
                    qua->output_stream->readBuf[i] = 1.0f;
                if (qua->output_stream->readBuf[i] < -1.0f)
                    qua->output_stream->readBuf[i] = -1.0f;
            }

            volk_32f_s32f_convert_16i(output_wav_buffer, (float *)qua->output_stream->readBuf, 32767, dat_size);

            int final_out = audio::AudioSink::resample_s16(output_wav_buffer, output_wav_buffer_resamp, d_symbolrate, audio_samplerate, dat_size, 1);
            if (enable_audio && play_audio)
                audio_sink->push_samples(output_wav_buffer_resamp, final_out);
            if (output_data_type == DATA_FILE || record_live_audio)
            {
                data_out.write((char *)output_wav_buffer_resamp, final_out * sizeof(int16_t));
                final_data_size += final_out * sizeof(int16_t);
            }
            else
            {
                output_fifo->write((uint8_t *)output_wav_buffer_resamp, final_out * sizeof(int16_t));
            }
#endif

            agc->output_stream->flush();

            if (input_data_type == satdump::pipeline::DATA_FILE)
                progress = file_source->getPosition();

            if (time(NULL) % 10 == 0 && lastTime != time(NULL))
            {
                lastTime = time(NULL);
                logger->info("Progress " + std::to_string(round(((double)progress / (double)filesize) * 1000.0) / 10.0) + "%%");
            }
        }

        if (enable_audio)
            audio_sink->stop();

        // Finish up WAV
        if (output_to_file)
        {
            wave_writer.finish_header(final_data_size);
            data_out.close();
        }

        delete[] output_wav_buffer;
        delete[] output_wav_buffer_resamp;

        logger->info("Demodulation finished");

        if (input_data_type == satdump::pipeline::DATA_FILE)
            stop();
    }

    void GenericAnalogDemodModule::stop()
    {
        // Stop
        BaseDemodModule::stop();
        // res->stop();
        //   qua->stop();
        agc->output_stream->stopReader();
    }

    void GenericAnalogDemodModule::drawUI(bool window)
    {
        ImGui::Begin(name.c_str(), NULL, window ? 0 : NOWINDOW_FLAGS);

        ImGui::BeginGroup();
        constellation.draw(); // Constellation
        ImGui::EndGroup();

        ImGui::SameLine();

        ImGui::BeginGroup();
        {
            ImGui::Button("Settings", {200 * ui_scale, 20 * ui_scale});

            proc_mtx.lock();
            ImGui::SetNextItemWidth(200 * ui_scale);
            ImGui::InputInt("Bandwidth##bandwidthsetting", &upcoming_symbolrate);
            ImGui::Checkbox("Record Audio##analogoption", &record_demod_audio);

            ImGui::RadioButton("NFM##analogoption", reinterpret_cast<int *>(&modulation_type), ModulationType::NFM);
            ImGui::SameLine();
            ImGui::RadioButton("AM##analogoption", reinterpret_cast<int *>(&modulation_type), ModulationType::AM);
            style::beginDisabled();
            ImGui::RadioButton("WFM##analogoption", false);
            ImGui::SameLine();
            ImGui::RadioButton("USB##analogoption", false);
            ImGui::RadioButton("LSB##analogoption", false);
            // ImGui::SameLine();
            ImGui::SameLine();
            ImGui::RadioButton("CW##analogoption", false);
            style::endDisabled();

            if (ImGui::Button("Set###analogset"))
                settings_changed = true;
            proc_mtx.unlock();

            ImGui::Button("Signal", {200 * ui_scale, 20 * ui_scale});
            /* if (show_freq)
            {
                ImGui::Text("Freq : ");
                ImGui::SameLine();
                ImGui::TextColored(style::theme.orange, "%.0f Hz", display_freq);
            }
            snr_plot.draw(snr, peak_snr); */
            if (!d_is_streaming_input)
                if (ImGui::Checkbox("Show FFT", &show_fft))
                    fft_splitter->set_enabled("fft", show_fft);
            if (enable_audio)
            {
                const char *btn_icon, *label;
                ImVec4 color;
                if (play_audio)
                {
                    color = style::theme.green.Value;
                    btn_icon = u8"\uF028##aptaudio";
                    label = "Audio Playing";
                }
                else
                {
                    color = style::theme.red.Value;
                    btn_icon = u8"\uF026##aptaudio";
                    label = "Audio Muted";
                }

                ImGui::PushStyleColor(ImGuiCol_Text, color);
                if (ImGui::Button(btn_icon))
                    play_audio = !play_audio;
                ImGui::PopStyleColor();
                ImGui::SameLine();
                ImGui::TextUnformatted(label);
            }
        }
        ImGui::EndGroup();

        if (!d_is_streaming_input)
            ImGui::ProgressBar((double)progress / (double)filesize, ImVec2(ImGui::GetContentRegionAvail().x, 20 * ui_scale));

        drawStopButton();
        ImGui::End();
        drawFFT();
    }

    std::string GenericAnalogDemodModule::make_safe_string(const std::string &str)
    {
        std::string safe_str;
        for (const char c : str)
        {
            if (std::isalnum(c)) // character is alphanumeric
            {
                safe_str.push_back(c);
            }
            else if (safe_str.empty() || safe_str.back() != '_') // replace other characters with underscore
            {
                safe_str.push_back('_'); // but only once in a row
            }
        }
        if (!safe_str.empty() && safe_str.back() == '_') // remove trailing underscore
        {
            safe_str.pop_back();
        }
        return safe_str;
    }

    std::string GenericAnalogDemodModule::getID() { return "generic_analog_demod"; }

    std::shared_ptr<satdump::pipeline::ProcessingModule> GenericAnalogDemodModule::getInstance(std::string input_file, std::string output_file_hint, nlohmann::json parameters)
    {
        return std::make_shared<GenericAnalogDemodModule>(input_file, output_file_hint, parameters);
    }
} // namespace generic_analog