#![allow(dead_code)]

mod bfsk;
mod complex;
pub mod fft;
mod filtering;
mod iq;
mod math;
mod nco;
mod squelch;
mod ted;
mod units;
mod windows;

use egui_plot::{Legend, Line, Plot};
use hound::WavWriter;
use rand::{Rng, rand_core::le, seq::index::sample};
use std::{
    cell::{Cell, RefCell},
    collections::VecDeque,
    env::{self, args},
    fs::File,
    io::{BufWriter, Sink, Write, stdout},
    ops::DerefMut,
    sync::Arc,
    time::Duration,
};
use tokio::{join, net::UdpSocket, select, sync::mpsc::error::TryRecvError, time::timeout};

use crate::{
    bfsk::BFSKMod,
    complex::Complex32,
    filtering::{dc_block::DCBlocker, fir::FIRFilter},
    iq::IQSampler,
    nco::Nco,
    ted::elg::ELGate,
    units::frequency::hz_to_rad_per_sample,
};
use eframe::egui::{self, CentralPanel, Color32, debug_text::print};
use tokio::sync::RwLock;
use tokio::sync::mpsc::{Receiver, Sender, channel};

const BAUD_RATE: u32 = 1000;
const SAMPLE_RATE: u32 = 48000;

// Modulation parameters
const CENTER_FREQ: f32 = 1700.;
const DEVIATION: f32 = 500.;

pub enum SampleSenderCommand {
    Open,
    Close,
    Sample(f32),
}

pub trait SampleSender {
    async fn open_link(&mut self);
    async fn send_sample(&mut self, sample: f32);
    async fn close_link(&mut self);
}

struct WavSampleSender {
    writer: Option<WavWriter<BufWriter<File>>>,
}

impl Default for WavSampleSender {
    fn default() -> Self {
        Self { writer: None }
    }
}

impl SampleSender for WavSampleSender {
    async fn open_link(&mut self) {
        let spec = hound::WavSpec {
            channels: 1,
            sample_rate: SAMPLE_RATE,
            bits_per_sample: 16,
            sample_format: hound::SampleFormat::Int,
        };
        self.writer = Some(hound::WavWriter::create("audio/modulated.wav", spec).unwrap());
    }

    async fn send_sample(&mut self, sample: f32) {
        let out_sample = (sample * i16::MAX as f32) as i16;
        self.writer
            .as_mut()
            .unwrap()
            .write_sample(out_sample)
            .unwrap();
    }

    async fn close_link(&mut self) {
        self.writer = None;
    }
}

struct Transceiver {
    tx_stream: Sender<Vec<u8>>,
    rx_stream: Receiver<Vec<u8>>,

    eye_receiver: Receiver<Vec<f32>>,
}

impl Transceiver {
    pub async fn send(&self, data: Vec<u8>) {
        self.tx_stream.send(data).await;
    }

    pub fn get_sender(&self) -> Sender<Vec<u8>> {
        self.tx_stream.clone()
    }

    pub fn try_recv(&mut self) -> Result<Vec<u8>, TryRecvError> {
        self.rx_stream.try_recv()
    }

    pub fn try_recv_eye(&mut self) -> Result<Vec<f32>, TryRecvError> {
        self.eye_receiver.try_recv()
    }

    pub fn start(
        mut sample_stream: Receiver<f32>,
        mut sample_sender: Sender<SampleSenderCommand>,
    ) -> Self {
        let (mut eyes_tx, eyes_rx) = channel::<Vec<f32>>(1024);

        let (rx_stream_sender, rx_stream_receiver) = channel::<Vec<u8>>(1024);
        let (tx_stream_sender, mut tx_stream_receiver) = channel::<Vec<u8>>(1024);

        tokio::spawn(async move {
            let squelch_sum = 0.;
            loop {}
        });

        Self {
            eye_receiver: eyes_rx,

            tx_stream: tx_stream_sender,
            rx_stream: rx_stream_receiver,
        }
    }

    async fn squelch_detector(sample_stream: &mut Receiver<f32>) {
        let length = 200;
        let level = 0.4;
        let mut iq_sampler = IQSampler::new(hz_to_rad_per_sample(CENTER_FREQ, SAMPLE_RATE as f32));
        let mut squelch_sum = 0.;
        let mut i = 0;
        while let Some(smpl) = sample_stream.recv().await {
            let iq = iq_sampler.sample(smpl);
            squelch_sum += iq.mag() / length as f32;
            i += 1;

            if i >= length {
                if squelch_sum >= level {
                    return;
                }
                i = 0;
                squelch_sum = 0.;
            }
        }
    }

    pub async fn transmit(frame: Frame, samples_sender: &mut Sender<SampleSenderCommand>) {
        let bytes = frame.bytes();
        let mut bit_stream = bytes.iter().flat_map(|x| byte_to_bits(*x));
        let modulator = BFSKMod::new(
            (SAMPLE_RATE as f32 / BAUD_RATE as f32).round() as u32,
            hz_to_rad_per_sample(DEVIATION, SAMPLE_RATE as f32),
            &mut bit_stream,
        );

        let up_lo = Nco::new(hz_to_rad_per_sample(CENTER_FREQ, SAMPLE_RATE as f32));
        samples_sender.send(SampleSenderCommand::Open).await;
        for (m, up) in modulator.zip(up_lo) {
            let sample = m * up;
            samples_sender
                .send(SampleSenderCommand::Sample(sample.re))
                .await;
        }
        samples_sender.send(SampleSenderCommand::Close).await;
    }

    async fn receive(
        sample_stream: &mut Receiver<f32>,
        eye_sender: &mut Sender<Vec<f32>>,
    ) -> Result<Frame, FrameConstructionError> {
        let mut iq_sampler = IQSampler::new(hz_to_rad_per_sample(CENTER_FREQ, SAMPLE_RATE as f32));

        let samples_per_symbol = (SAMPLE_RATE as f32) / (BAUD_RATE as f32);

        let correllator_length = samples_per_symbol as usize;
        let mut pos_nco = Nco::new(hz_to_rad_per_sample(DEVIATION, SAMPLE_RATE as f32));
        let mut neg_nco = Nco::new(hz_to_rad_per_sample(-DEVIATION, SAMPLE_RATE as f32));
        let pos_ir = (0..correllator_length).map(|i| {
            pos_nco.step();
            pos_nco.cexp() * windows::blackmann(i as f32 / correllator_length as f32)
        });
        let neg_ir = (0..correllator_length).map(|i| {
            neg_nco.step();
            neg_nco.cexp() * windows::blackmann(i as f32 / correllator_length as f32)
        });
        let mut pos_correllator = FIRFilter::new(&pos_ir.collect::<Vec<_>>());
        let mut neg_correllator = FIRFilter::new(&neg_ir.collect::<Vec<_>>());
        pos_correllator.normalize_freq(hz_to_rad_per_sample(DEVIATION, SAMPLE_RATE as f32));
        neg_correllator.normalize_freq(hz_to_rad_per_sample(-DEVIATION, SAMPLE_RATE as f32));

        let mut matched_lowpass = FIRFilter::new(&vec![
            Complex32::new(1., 0.);
            samples_per_symbol as usize / 2
        ]);
        matched_lowpass.normalize_freq(hz_to_rad_per_sample(DEVIATION, SAMPLE_RATE as f32));
        //let mut dc_block = DCBlocker::new(0.999);
        let mut dc_block = DCBlocker::new(1.);

        let loop_i = 0.0;
        let loop_p = 0.1;
        let mut loop_ir = vec![Complex32::new(loop_i, 0.); samples_per_symbol as usize];
        loop_ir.push(Complex32::new(loop_p, 0.));
        let mut elg = ELGate::new(samples_per_symbol, FIRFilter::new(&loop_ir));

        // Frame reconstruction
        let mut last_byte = 0x00u8;
        let mut frame_constructor = FrameConstructor::new();
        let mut bit_count: Option<u32> = None;
        while let Some(sample) = sample_stream.recv().await {
            let iq = iq_sampler.sample(sample);
            let matched =
                matched_lowpass
                    .next_real(dc_block.next_real(
                        pos_correllator.next(iq).mag() - neg_correllator.next(iq).mag(),
                    ));
            if let Some((bit_sample, eye)) = elg.next_eye(matched) {
                let _ = eye_sender.send(eye).await;
                last_byte >>= 1;
                last_byte |= ((bit_sample > 0.) as u8) << 7;
                //last_byte <<= 1;
                //last_byte |= ((bit_sample < 0.) as u8);
                bit_count = bit_count.map(|x| x + 1);

                if let None = bit_count
                    && last_byte == 0xD8
                {
                    // Potential frame starts
                    last_byte = 0;
                    frame_constructor = FrameConstructor::new();
                    bit_count = Some(0);
                }

                if let Some(8) = bit_count {
                    let frame_opt = frame_constructor.add_byte(last_byte);
                    bit_count = Some(0);
                    //print!("{}", last_byte as char);
                    print!(".{:x}.", last_byte);
                    let _ = std::io::stdout().flush();
                    if let Ok(Some(Frame(ref frame_data))) = frame_opt {
                        println!("Got data");
                        return Ok(Frame(frame_data.to_vec()));
                    }

                    if let Err(()) = frame_opt {
                        // Erroneous frame
                        println!("Error");
                        return Err(());
                    }
                }
            }
        }
        return Err(());
    }
}

struct Frame(Vec<u8>);

type FrameConstructionError = ();
pub struct FrameConstructor {
    frame: Vec<u8>,
    frame_countdown: Option<u16>,
    checksum: u8,
    started: bool,
}

impl FrameConstructor {
    pub fn new() -> Self {
        Self {
            frame: Vec::new(),
            frame_countdown: None,
            checksum: 0u8,
            started: false,
        }
    }

    pub fn add_byte(&mut self, byte: u8) -> Result<Option<Frame>, FrameConstructionError> {
        if self.frame.is_empty() && byte != 0x4C && !self.started {
            println!("Wrong type {:x}", byte);
            self.started = true;
            return Err(());
        }

        if self.frame.is_empty() && byte == 0x4C && !self.started {
            self.started = true;
            return Ok(None);
        }

        self.frame.push(byte);

        // Retrieve length
        if self.frame.len() == 1 {
            self.frame_countdown = Some(self.frame[0] as u16);
            return Ok(None);
        }
        if self.frame.len() == 2 {
            *self.frame_countdown.as_mut().unwrap() |= (self.frame[1] as u16) << 8;
            return Ok(None);
        }

        if self.frame_countdown.unwrap() == 0 {
            // All data has been received
            if self.checksum == byte {
                return Ok(Some(Frame(self.frame.iter().skip(2).copied().collect())));
            }

            println!("Checksum failed");
            return Err(());
        }

        //self.frame.push(byte);
        self.checksum ^= byte;
        *self.frame_countdown.as_mut().unwrap() -= 1;

        Ok(None)
    }
}

impl Frame {
    pub fn bytes(&self) -> Vec<u8> {
        let mut output_bytes = vec![];

        // Initial training sequence
        output_bytes.append(&mut vec![0b01010101; 64]);

        // Preamble byte
        output_bytes.push(0xD8);

        let x = &self.0;
        assert!(x.len() < 65536, "Data size over MTU");

        output_bytes.push(0x4C); // DATA FRAME

        let len_u16 = x.len() as u16;
        output_bytes.push((len_u16 & 0xFF).try_into().unwrap());
        output_bytes.push(((len_u16 >> 8) & 0xFF).try_into().unwrap());

        let mut checksum = 0u8;
        x.iter().for_each(|x| checksum ^= x);
        output_bytes.extend(x.iter());

        output_bytes.push(checksum);

        // SEND EOT
        output_bytes.extend(std::iter::repeat_n(4, 32));
        output_bytes
    }
}

#[tokio::main]
async fn main() {
    //Transceiver::transmit(Frame::Data("Skibditoilet".repeat(100).bytes().collect::<Vec<_>>()), &mut WavSampleSender{}).await;
    //Transceiver::transmit(Frame::Ack, &mut WavSampleSender::default()).await;
    //return;

    let native_options = eframe::NativeOptions::default();
    let _ = eframe::run_native(
        "Egui",
        native_options,
        Box::new(|cc| Ok(Box::new(EguiApp::new(cc)))),
    );
}

//#[derive(Default)]

struct DummySampleSender();

impl SampleSender for DummySampleSender {
    async fn open_link(&mut self) {}
    async fn send_sample(&mut self, _sample: f32) {}
    async fn close_link(&mut self) {}
}

struct EguiApp {
    a_transceiver: Transceiver,
    b_transceiver: Transceiver,

    eyes_a: VecDeque<Vec<f32>>,
    eyes_b: VecDeque<Vec<f32>>,
}
impl EguiApp {
    fn new(_cc: &eframe::CreationContext<'_>) -> Self {
        let (up_a_sender, mut up_a_receiver) = channel::<SampleSenderCommand>(1024);
        let (down_a_sender, down_a_receiver) = channel::<f32>(1024);

        let (up_b_sender, mut up_b_receiver) = channel::<SampleSenderCommand>(1024);
        let (down_b_sender, down_b_receiver) = channel::<f32>(1024);

        let (a2b_tx, mut a2b_rx) = channel::<f32>(1024);
        let (b2a_tx, mut b2a_rx) = channel::<f32>(1024);

        let a_txrx = Transceiver::start(down_a_receiver, up_a_sender);
        let b_txrx = Transceiver::start(down_b_receiver, up_b_sender);

        // A dummy channel
        tokio::spawn(async move {
            //let rng = rand::thread_rng();
            let mut sending = false;
            loop {
                let noise = rand::random::<f32>() * 0.1;
                let mut sample = 0.;

                match up_a_receiver.try_recv() {
                    Ok(SampleSenderCommand::Open) => {
                        sending = true;
                        println!("open");
                    }
                    Ok(SampleSenderCommand::Close) => {
                        sending = false;
                        println!("close");
                    }
                    Ok(SampleSenderCommand::Sample(x)) => {
                        sample = x;
                    }
                    _ => {}
                }

                if sending {
                    // Flush receiver buffer but ignore
                    while let Ok(_) = b2a_rx.try_recv() {}

                    // Send to other
                    a2b_tx.send(sample + noise).await;
                } else if let Ok(down_sample) = b2a_rx.try_recv() {
                    down_a_sender.send(down_sample).await;
                }
            }
        });

        // B dummy channel
        tokio::spawn(async move {
            let mut sending = false;
            loop {
                let noise = rand::random::<f32>() * 0.1;
                let mut sample = 0.;

                match up_b_receiver.try_recv() {
                    Ok(SampleSenderCommand::Open) => {
                        sending = true;
                    }
                    Ok(SampleSenderCommand::Close) => {
                        sending = false;
                    }
                    Ok(SampleSenderCommand::Sample(x)) => {
                        sample = x;
                    }
                    _ => {}
                }

                if sending {
                    // Flush receiver buffer but ignore
                    while let Ok(_) = a2b_rx.try_recv() {}

                    // Send to other
                    b2a_tx.send(sample + noise).await;
                } else if let Ok(down_sample) = a2b_rx.try_recv() {
                    down_b_sender.send(down_sample).await;
                }
            }
        });

        EguiApp {
            a_transceiver: a_txrx,
            b_transceiver: b_txrx,

            eyes_a: VecDeque::new(),
            eyes_b: VecDeque::new(),
        }
    }
}

impl eframe::App for EguiApp {
    fn update(&mut self, ctx: &egui::Context, _frame: &mut eframe::Frame) {
        egui::CentralPanel::default().show(ctx, |ui| {
            let max_eyes = 100;

            while let Ok(eye) = self.a_transceiver.try_recv_eye() {
                self.eyes_a.push_back(eye);
            }
            while self.eyes_a.len() > max_eyes {
                self.eyes_a.pop_front();
            }

            while let Ok(eye) = self.b_transceiver.try_recv_eye() {
                self.eyes_b.push_back(eye);
            }
            while self.eyes_b.len() > max_eyes {
                self.eyes_b.pop_front();
            }

            ui.columns(2, |uis| {
                Plot::new("EyeA")
                    .legend(Legend::default())
                    .show(&mut uis[0], |plot_ui| {
                        //plot_ui.set_auto_bounds(Vec2b { x: false, y: false });
                        for eye in self.eyes_a.iter() {
                            let line = Line::new(
                                "EyeA",
                                eye.iter()
                                    .enumerate()
                                    .map(|(i, x)| [i as f64, *x as f64])
                                    .collect::<Vec<_>>(),
                            )
                            .color(Color32::LIGHT_GREEN);
                            plot_ui.line(line);
                        }
                    });

                if uis[0].button("Start").clicked() {
                    let snd = self.a_transceiver.get_sender();
                    tokio::spawn(async move {
                        let _ = snd
                            .send("Skibditoilet".repeat(10).as_bytes().to_vec())
                            .await;
                    });
                }

                Plot::new("EyeB")
                    .legend(Legend::default())
                    .show(&mut uis[1], |plot_ui| {
                        //plot_ui.set_auto_bounds(Vec2b { x: false, y: false });
                        for eye in self.eyes_b.iter() {
                            let line = Line::new(
                                "EyeB",
                                eye.iter()
                                    .enumerate()
                                    .map(|(i, x)| [i as f64, *x as f64])
                                    .collect::<Vec<_>>(),
                            )
                            .color(Color32::LIGHT_GREEN);
                            plot_ui.line(line);
                        }
                    });

                if uis[1].button("Start").clicked() {
                    let snd = self.b_transceiver.get_sender();
                    tokio::spawn(async move {
                        let _ = snd
                            .send("Skibditoilet".repeat(10).as_bytes().to_vec())
                            .await;
                    });
                }
            });
        }); // Central panel

        std::thread::sleep(Duration::from_millis(16));
        ctx.request_repaint();
    }
}

fn byte_to_bits(byte: u8) -> Vec<bool> {
    vec![
        byte & 1 == 1,
        (byte >> 1) & 1 == 1,
        (byte >> 2) & 1 == 1,
        (byte >> 3) & 1 == 1,
        (byte >> 4) & 1 == 1,
        (byte >> 5) & 1 == 1,
        (byte >> 6) & 1 == 1,
        (byte >> 7) & 1 == 1,
    ]
}

fn bits_to_byte(bits: &[bool]) -> u8 {
    bits[0] as u8
        | (bits[1] as u8) << 1
        | (bits[2] as u8) << 2
        | (bits[3] as u8) << 3
        | (bits[4] as u8) << 4
        | (bits[5] as u8) << 5
        | (bits[6] as u8) << 6
        | (bits[7] as u8) << 7
}