#include <algorithm>
#include <alsa/asoundlib.h>
#include <cmath>
#include <complex>
#include <cstring>
#include <iostream>
#include <rtl-sdr.h>
#include <signal.h>
#include <vector>

#define DR_WAV_IMPLEMENTATION
#include "dr_wav.h"

#define DEFAULT_SAMPLE_RATE  2400000 // 2.4 MS/s
#define AUDIO_SAMPLE_RATE    48000
#define RTL_FM_DEMOD_BUF_LEN 16384 // 16kB

#include <stdio.h>
#include <stdlib.h>

volatile bool stop_flag = 0;

void handle_sigint(int sig)
{
    (void)sig;
    stop_flag = 1;
}

int main()
{
    signal(SIGINT, handle_sigint);

    rtlsdr_dev_t* dev = NULL;
    rtlsdr_open(&dev, 0);
    rtlsdr_set_center_freq(dev, 96700000);
    rtlsdr_set_sample_rate(dev, DEFAULT_SAMPLE_RATE);
    rtlsdr_set_tuner_gain_mode(dev, 1);
    rtlsdr_set_tuner_gain(dev, 400); // 40dB
    rtlsdr_reset_buffer(dev);

    snd_pcm_t* pcm_handle;
    snd_pcm_open(&pcm_handle, "default", SND_PCM_STREAM_PLAYBACK, 0);
    snd_pcm_set_params(pcm_handle, SND_PCM_FORMAT_S16_LE, SND_PCM_ACCESS_RW_INTERLEAVED, 1, AUDIO_SAMPLE_RATE, 1, 500000);

    drwav wav;
    drwav_data_format format = {
        .container = drwav_container_riff,
        .format = DR_WAVE_FORMAT_PCM,
        .channels = 1,
        .sampleRate = 48000,
        .bitsPerSample = 16};

    // Open file for writing
    drwav_init_file_write(&wav, "output.wav", &format, NULL);

    uint8_t rtl_buf[RTL_FM_DEMOD_BUF_LEN];
    float iq_buf[RTL_FM_DEMOD_BUF_LEN / 2];
    float demod_buf[RTL_FM_DEMOD_BUF_LEN / 2];
    int16_t audio_buf[8192];

    while (!stop_flag)
    {
        int n_read;
        rtlsdr_read_sync(dev, rtl_buf, RTL_FM_DEMOD_BUF_LEN, &n_read);
        int n_samples = n_read / 2;

        // i/q around 0
        for (int i = 0; i < n_samples; i++)
        {
            int8_t i_sample = rtl_buf[2 * i] - 127;
            int8_t q_sample = rtl_buf[2 * i + 1] - 127;
            iq_buf[i] = atan2(q_sample, i_sample);
        }

        // differential phase
        static float prev_phase = 0;
        for (int i = 0; i < n_samples; i++)
        {
            float dphi = iq_buf[i] - prev_phase;
            // keep dphi within -pi and pi (as phase is circular)
            if (dphi > M_PI)
                dphi -= 2 * M_PI;
            if (dphi < -M_PI)
                dphi += 2 * M_PI;
            prev_phase = iq_buf[i];

            // low pass
            static float y = 0;
            float x = dphi;
            y = y + (x - y) * 0.294f; // alpha 50us 48kHz
            iq_buf[i] = y;
        }

        // then decimate
        int out_index = 0;
        for (int i = 0; i < n_samples; i += 50)
        {
            if (out_index < 8192)
            {
                audio_buf[out_index++] = (int16_t)(iq_buf[i] * 32767.0f);
            }
        }

        snd_pcm_writei(pcm_handle, audio_buf, out_index);
        drwav_write_pcm_frames(&wav, out_index, audio_buf);
    }
    drwav_uninit(&wav);
    rtlsdr_close(dev);
}