/*
 * Copyright (C) 2023 by Matthias P. Braendli <matthias.braendli@mpb.li>
 *
 * Based on previous work by
 * Copyright (C) 2022 by Felix Erckenbrecht <eligs@eligs.de>
 * Copyright (C) 2016-2018 by Steve Markgraf <steve@steve-m.de>
 * Copyright (C) 2009 by Bartek Kania <mbk@gnarf.org>
 *
 * SPDX-License-Identifier: GPL-2.0+
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

use std::mem::size_of;
use std::sync::atomic::{AtomicBool, Ordering};
use std::{env, thread};
use std::io::{prelude::*, BufReader, BufWriter};
use std::fs::File;
use std::sync::{mpsc, Arc};
use getopts::Options;

mod fl2k;

const TRIG_TABLE_ORDER: usize = 8;
const TRIG_TABLE_SHIFT: usize = 32 - TRIG_TABLE_ORDER;
const TRIG_TABLE_LEN: usize = 1 << TRIG_TABLE_ORDER;

const PI: f32 = std::f32::consts::PI;
const INT32_MAX_AS_FLOAT: f32 = 0x1_0000_0000u64 as f32;
const ANG_INCR: f32 = INT32_MAX_AS_FLOAT / (2.0 * PI);

enum Waveform { Sine, Rect }

#[derive(Clone, Copy)]
enum Output { Debug, FL2K }

struct BufferDumper<T> {
    phantom : std::marker::PhantomData<T>,
    writer : BufWriter<File>,
}

impl<T> BufferDumper<T> {
    fn new(filename: &str) -> Self {
        let writer = BufWriter::new(File::create(filename).expect("create file"));
        BufferDumper { phantom: std::marker::PhantomData, writer }
    }

    fn write_buf(&mut self, buf: &[T]) -> std::io::Result<()> {
        let buf_u8: &[u8] = unsafe {
            std::slice::from_raw_parts(
                buf.as_ptr() as *const u8,
                buf.len() * size_of::<T>()
            )
        };

        self.writer.write_all(buf_u8)
    }
}

struct DDS {
    trig_table_quadrature: Vec<i8>,
    trig_table_inphase: Vec<i8>,

    /* instantaneous phase */
    phase: u32,
    /* phase increment */
    phase_step: u32,

    /* for phase modulation */
    phase_delta: i32,
    phase_slope: i32,
}

impl DDS {
    fn init(samp_rate: f32, freq: f32, phase: f32, waveform: Waveform) -> Self {
        let mut trig_table_inphase = Vec::with_capacity(TRIG_TABLE_LEN);
        let mut trig_table_quadrature = Vec::with_capacity(TRIG_TABLE_LEN);

        let incr = 1.0f32 / TRIG_TABLE_LEN as f32;
        for i in 0..TRIG_TABLE_LEN {
            let inphase = f32::cos(incr * i as f32 * 2.0 * PI) * 127.0;
            let quadrature = f32::sin(incr * i as f32 * 2.0 * PI) * 127.0;

            match waveform {
                Waveform::Sine => {
                    trig_table_inphase.push(f32::round(inphase) as i8);
                    trig_table_quadrature.push(f32::round(quadrature) as i8);
                }
                Waveform::Rect => {
                    trig_table_inphase.push(if inphase >= 0.0 { 127 } else { -127 });
                    trig_table_quadrature.push(if quadrature >= 0.0 { 127 } else { -127 });
                }
            }
        }

        let phase_step = (freq / samp_rate) * 2.0 * PI * ANG_INCR;

        DDS {
            trig_table_quadrature,
            trig_table_inphase,
            phase: f32::round(phase * ANG_INCR) as u32,
            phase_step: f32::round(phase_step) as u32,
            phase_delta: 0,
            phase_slope: 0,
        }
    }
}

fn print_usage(program: &str, opts: Options) {
    let brief = format!("Usage: {} FILE [options]", program);
    eprint!("{}", opts.usage(&brief));
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();

    let mut opts = Options::new();
    opts.optopt("f", "file", "Input file, containing signed 16-bit samples.", "FILE");
    opts.optopt("d", "device-index", "Select device index", "DEVINDEX");
    opts.optopt("c", "center-freq", "Center frequency in Hz (default: 1440 kHz)", "FREQ");
    opts.optopt("s", "samplerate", "Samplerate in Hz (default: 96 MS/s)", "RATE");
    opts.optopt("m", "mod-index", "Modulation index (default: 0.25)]", "FACTOR");
    opts.optopt("i", "input-rate", "Input baseband sample rate (default: 48000 Hz)", "RATE");
    opts.optopt("w", "waveform", "(sine|rect) default: rect", "WAVEFORM");
    opts.optflag("C", "device-count", "Return FL2K device count and quit");
    opts.optflag("D", "debug", "Write to debug files instead of FL2K");
    opts.optflag("h", "help", "print this help menu");
    let cli_args = match opts.parse(&args[1..]) {
        Ok(m) => { m }
        Err(f) => { panic!("{}", f.to_string()) }
    };
    if cli_args.opt_present("h") {
        print_usage(&program, opts);
        std::process::exit(1);
    }

    if cli_args.opt_present("C") {
        eprintln!("FL2K device count {}", fl2k::get_device_count());
        return;
    }

    let output = if cli_args.opt_present("D") {
        Output::Debug
    }
    else {
        Output::FL2K
    };

    let device_index: u32 = match cli_args.opt_str("d") {
        Some(s) => s.parse().expect("integer value"),
        None => 0,
    };

    let samp_rate: u32 = match cli_args.opt_str("s") {
        Some(s) => s.parse().expect("integer value"),
        None => 96_000_000,
    };

    let base_freq = match cli_args.opt_str("c") {
        Some(s) => s.parse().expect("floating point value"),
        None => 1_440_000.0,
    };

    let input_rate: u32 = match cli_args.opt_str("i") {
        Some(s) => s.parse().expect("integer value"),
        None => 48_000,
    };

    let modulation_index = match cli_args.opt_str("m") {
        Some(s) => s.parse().expect("floating point value"),
        None => 0.25,
    };

    let waveform = match cli_args.opt_str("w") {
        None => Waveform::Rect,
        Some(w) if w == "sine" => Waveform::Sine,
        Some(w) if w == "rect" => Waveform::Rect,
        _ => {
            eprintln!("Waveform must be 'sine' or 'rect'");
            print_usage(&program, opts);
            std::process::exit(1);
        }
    };

    let source_file_name = match cli_args.opt_str("f") {
        Some(f) => f,
        None => {
            eprintln!("Specify input file!");
            print_usage(&program, opts);
            std::process::exit(1);
        }
    };

    if samp_rate % input_rate != 0 {
        eprintln!("WARNING: input_rate freq does not divide sample rate.");
    }
    let rf_to_baseband_sample_ratio = samp_rate / input_rate;

    eprintln!("Input rate:       {} kHz", (input_rate as f32)/1e3);
    eprintln!("Samplerate:       {} MHz", (samp_rate as f32)/1e6);
    eprintln!("Center frequency: {} kHz", base_freq/1e3);

    let running = Arc::new(AtomicBool::new(true));

    let r = running.clone();
    ctrlc::set_handler(move || {
        r.store(false, Ordering::SeqCst);
    }).expect("Error setting Ctrl-C handler");

    let (input_samples_tx, input_samples_rx) = mpsc::sync_channel::<Vec<f32>>(2);
    let (pd_tx, pd_rx) = mpsc::sync_channel(2);
    let (iq_tx, iq_rx) = mpsc::sync_channel(2);

    let source_file = File::open(source_file_name).expect("open file");
    let mut source_file = BufReader::new(source_file);

    const BASEBAND_BUF_SAMPS: usize = 1024;

    // Read file and convert samples
    thread::spawn(move || {
        let mut in_debug = match output {
            Output::Debug => Some(BufferDumper::<f32>::new("debug-in.f32")),
            Output::FL2K => None,
        };

        while running.load(Ordering::SeqCst) {
            let mut buf = Vec::with_capacity(BASEBAND_BUF_SAMPS);
            buf.resize(BASEBAND_BUF_SAMPS, 0i16);

            let mut buf_u8: &mut [u8] = unsafe {
                std::slice::from_raw_parts_mut(
                    buf.as_mut_ptr() as *mut u8,
                    buf.len() * size_of::<i16>()
                )
            };

            match source_file.read(&mut buf_u8) {
                Ok(len) => {
                    if len == 0 {
                        if let Err(e_seek) = source_file.rewind() {
                            eprintln!("Failed to rewind source file: {}", e_seek);
                            break;
                        }
                    }
                    buf.resize(len / size_of::<i16>(), 0);
                },
                Err(e) => {
                    eprintln!("Failed to read source file: {}", e);
                    break;
                }
            }

            // Downmix stereo to mono, apply (v/2 + i16::MAX/2), and normalise to [-1.0; 1.0]
            let buf: Vec<f32> = buf
                .chunks_exact(2)
                .map(|v| ((v[0]/2 + v[1]/2)/2 + i16::MAX/2) as f32 / 32768.0)
                .collect();

            if let Some(w) = &mut in_debug {
                if let Err(e) = w.write_buf(&buf) {
                    eprintln!("Error writing debug-in.f32: {}", e);
                }
            }

            if let Err(_) = input_samples_tx.send(buf) {
                eprintln!("Quit read thread");
                break;
            }
        }

        eprintln!("Leaving input thread");
    });

    // Read samples, calculate PD and PDSLOPE
    thread::spawn(move || {
        let mut last_sample = 0f32;

        let mut debug_writer = match output {
            Output::Debug => Some(BufWriter::new(File::create("debug-pd.csv").expect("create file"))),
            Output::FL2K => None,
        };

        loop {
            let Ok(buf) = input_samples_rx.recv() else { break };

            let mut pd_buf = Vec::with_capacity(buf.len());

            /* What we do here is calculate a linear "slope" from
               the previous sample to this one. This is then used by
               the modulator to gently increase/decrease the phase
               with each sample without the need to recalculate
               the dds parameters. In fact this gives us a very
               efficient and pretty good interpolation filter. */

            for sample in buf {
                let slope = (sample - last_sample) / rf_to_baseband_sample_ratio as f32;

                let pd = last_sample * modulation_index * INT32_MAX_AS_FLOAT;

                const MIN_VAL: f32 = std::i32::MIN as f32;
                const MAX_VAL: f32 = std::i32::MAX as f32;

                if pd < MIN_VAL || pd > MAX_VAL {
                    panic!("pd out of bounds {}", pd);
                }

                let pdslope = slope * modulation_index * INT32_MAX_AS_FLOAT;
                if pdslope < MIN_VAL || pdslope > MAX_VAL {
                    panic!("pdslope out of bounds {}", pdslope);
                }

                if let Some(w) = &mut debug_writer {
                    writeln!(w, "{},{},{},{}", sample, slope, pd, pdslope)
                        .expect("write debug-pd.csv");
                }

                pd_buf.push((pd as i32, pdslope as i32));

                last_sample = sample;
            }

            if let Err(_) = pd_tx.send(pd_buf) {
                eprintln!("Quit pd thread");
                break;
            }
        }
        eprintln!("Leaving pd thread");
    });

    // Read PD and PDSLOPE, interpolate to higher rate
    thread::spawn(move || {
        let mut dds = DDS::init(samp_rate as f32, base_freq, 0.0, waveform);

        let mut debug_writer = match output {
            Output::Debug => Some(BufWriter::new(File::create("debug-dds.csv").expect("create file"))),
            Output::FL2K => None,
        };

        loop {
            let Ok(pd_buf) = pd_rx.recv() else { break };

            for (pd, pdslope) in pd_buf {
                dds.phase_delta = pd;
                dds.phase_slope = pdslope;

                let len = rf_to_baseband_sample_ratio as usize;
                let mut out_i = Vec::with_capacity(len);
                let mut out_q = Vec::with_capacity(len);
                for ix in 0..len {
                    // get current carrier phase, add phase mod, calculate table index
                    let phase_idx_i = dds.phase.overflowing_sub(dds.phase_delta as u32).0 >> TRIG_TABLE_SHIFT;
                    let phase_idx_q = dds.phase.overflowing_add(dds.phase_delta as u32).0 >> TRIG_TABLE_SHIFT;

                    if phase_idx_q > 255 || phase_idx_i > 255 {
                        panic!("Phase IDX out of bounds");
                    }

                    out_i.push(dds.trig_table_inphase[phase_idx_i as usize]);
                    out_q.push(dds.trig_table_quadrature[phase_idx_q as usize]);

                    if let Some(w) = &mut debug_writer {
                        writeln!(w, "{},{},{},{},{}", ix, dds.phase, dds.phase_delta, phase_idx_i, phase_idx_q)
                            .expect("write debug-dds.csv");
                    }

                    dds.phase = dds.phase.overflowing_add(dds.phase_step).0;
                    dds.phase_delta += dds.phase_slope;
                }

                if let Err(_) = iq_tx.send((out_i, out_q)) {
                    eprintln!("Quit dds thread");
                    break;
                }
            }
        }
        eprintln!("Leaving dds thread");
    });

    // Main thread, output to file/device
    match output {
        Output::FL2K => {
            let mut fl2k = fl2k::FL2K::open(device_index).expect("fl2k open");

            fl2k.set_sample_rate(samp_rate).expect("set fl2k sample rate");

            fl2k.start_tx().expect("fl2k start_tx");

            eprintln!("FL2K sample rate set to {}", fl2k.get_sample_rate().unwrap());

            loop {
                let Ok((i, q)) = iq_rx.recv() else { break };
                if fl2k.send(i, q) == false { break };
            }

            fl2k.stop_tx().expect("stop tx");
        }
        Output::Debug => {
            let mut out_file = BufferDumper::new("debug-out.i8");
            loop {
                let Ok((i_buf, q_buf)) = iq_rx.recv() else { break };

                if i_buf.len() != q_buf.len() {
                    panic!("i_buf and q_buf must have same length");
                }

                let mut buf = Vec::with_capacity(i_buf.len() * 2);
                for (i, q) in i_buf.iter().zip(q_buf) {
                    buf.push(*i);
                    buf.push(q);
                }

                let buf_u8: &[u8] = unsafe {
                    std::slice::from_raw_parts(
                        buf.as_ptr() as *const u8,
                        buf.len()
                    )
                };

                if let Err(e) = out_file.write_buf(buf_u8) {
                    eprintln!("Write output error: {}", e);
                    break;
                }
            }
        }
    }

    eprintln!("Leaving main thread");
}