#![allow(dead_code)]

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

use rand::{rand_core::le, seq::index::sample};
use std::{collections::VecDeque, time::Duration};
use tokio::{join, net::UdpSocket, select, 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, glow::SAMPLE_MASK_VALUE};
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 trait SampleSender {
    fn open_link(&mut self);
    fn send_samples(&mut self, samples: &[f32]);
    fn close_link(&mut self);
}

struct Transceiver {}

impl Transceiver {
    pub async fn start<T: SampleSender>(
        mut sample_stream: Receiver<f32>,
        mut tx_stream: Receiver<Vec<u8>>,
        mut rx_stream: Sender<Vec<u8>>,
        sample_sender: &mut T,
    ) {
        let mut resend: Option<Vec<u8>> = None;
        loop {
            select! {
                _ = Self::squelch_detector(&mut sample_stream) =>
                {
                    println!("Squelch UP");
                    select!
                    {
                        x = Self::receive(&mut sample_stream) =>
                        {
                            match x
                            {
                                Err(()) => {continue;},
                                Ok(Frame::Ack) =>
                                {
                                    resend = None;
                                }
                                Ok(Frame::Data(data)) =>
                                {
                                    rx_stream.send(data).await.unwrap();
                                    Self::transmit(Frame::Ack, sample_sender).await;
                                }
                            }
                        },
                        _ = tokio::time::sleep(Duration::from_secs(2)) => {continue;}, //TODO: 65
                                                                                       //sec
                                                                                       //timeout
                    }
                }, // End squelch
                data_opt = async
                {
                    tokio::time::sleep(Duration::from_secs(2)).await;
                    if let Some(resend_data) = resend.clone()
                    {
                        Some(resend_data)
                    }
                    else
                    {
                        tx_stream.recv().await
                    }
                }
                =>
                {
                    if let Some(data) = data_opt
                    {
                        Self::transmit(Frame::Data(data.clone()), sample_sender).await;
                        resend = Some(data);
                    }
                }
            }
        }
    }

    async fn squelch_detector(sample_stream: &mut Receiver<f32>) {
        let length = 500;
        let level = 0.01;
        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 {
            println!("sdkf");
            let iq = iq_sampler.sample(smpl);
            squelch_sum += iq.mag() / length as f32;
            i += 1;

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

    async fn transmit<T: SampleSender>(frame: Frame, samples_sender: &mut T) {
        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));

        let mut sample_buffer = vec![];
        for (m, up) in modulator.zip(up_lo) {
            let sample = m * up;
            sample_buffer.push(sample.re); // Project IQ
        }

        samples_sender.open_link();
        samples_sender.send_samples(&sample_buffer);
        samples_sender.close_link();
    }

    async fn receive(sample_stream: &mut Receiver<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(|_| {
            pos_nco.step();
            pos_nco.cexp()
        });
        let neg_ir = (0..correllator_length).map(|_| {
            neg_nco.step();
            neg_nco.cexp()
        });
        let mut pos_correllator = FIRFilter::new(&pos_ir.collect::<Vec<_>>());
        let mut neg_correllator = FIRFilter::new(&neg_ir.collect::<Vec<_>>());

        let mut matched_lowpass =
            FIRFilter::new(&vec![Complex32::new(1., 0.); samples_per_symbol as usize]);
        let mut dc_block = DCBlocker::new(0.995);

        let loop_i = 0.1;
        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 =
                dc_block
                    .next_real(matched_lowpass.next_real(
                        pos_correllator.next(iq).mag() - neg_correllator.next(iq).mag(),
                    ));
            if let Some(bit_sample) = elg.next(matched) {
                last_byte <<= 1;
                last_byte |= (bit_sample > 0.) as u8;
                bit_count = bit_count.map(|x| x + 1);

                if 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);
                    if let Ok(Some(Frame::Ack)) = frame_opt {
                        return Ok(Frame::Ack);
                    }

                    if let Ok(Some(Frame::Data(ref frame_data))) = frame_opt {
                        return Ok(Frame::Data(frame_data.to_vec()));
                    }

                    if let Err(()) = frame_opt {
                        // Erroneous frame
                        return Err(());
                    }
                }
            }
        }
        return Err(());
    }
}

enum Frame {
    Data(Vec<u8>),
    Ack,
}

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

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

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

        if self.frame.is_empty() && byte == 0xC4 {
            return Ok(Some(Frame::Ack));
        }

        if self.frame.is_empty() && byte == 0x4C {
            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::Data(
                    self.frame.iter().skip(2).copied().collect(),
                )));
            }

            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);

        // Command
        match self {
            Frame::Data(x) => {
                let mut checksum = 0u8;
                x.iter().for_each(|x| checksum ^= x);

                assert!(x.len() < 65536, "Data size over MTU");
                let len_u16 = x.len() as u16;
                output_bytes.push(0x4C); // DATA FRAME
                output_bytes.push((len_u16 & 0xFF).try_into().unwrap());
                output_bytes.push(((len_u16 >> 8) & 0xFF).try_into().unwrap());

                output_bytes.extend(x.iter());

                output_bytes.push(checksum);
            }
            Frame::Ack => {
                output_bytes.push(0xC4); // ACK FRAME
            }
        }

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

#[tokio::main]
async fn main() {
    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 EguiApp {}

struct DummySampleSender();

impl SampleSender for DummySampleSender {
    fn open_link(&mut self) {}
    fn send_samples(&mut self, samples: &[f32]) {}
    fn close_link(&mut self) {}
}

impl EguiApp {
    fn new(_cc: &eframe::CreationContext<'_>) -> Self {
        let (sample_tx, sample_rx) = channel::<f32>(1024);

        let (transmit_tx, transmit_rx) = channel::<Vec<u8>>(1024);
        let (receive_tx, receive_rx) = channel::<Vec<u8>>(1024);

        tokio::spawn(async move {
            Transceiver::start(sample_rx, transmit_rx, receive_tx, &mut DummySampleSender()).await;
        });

        tokio::spawn(async move {
            let sock = UdpSocket::bind("0.0.0.0:8080").await.unwrap();
            let mut buf = [0u8; 1024];
            let mut sample = 0i16;
            let mut byte_index = 0;

            loop {
                let len = sock.recv(&mut buf).await.unwrap();
                for x in buf.iter().take(len) {
                    sample |= (*x as i16) << (byte_index * 8);
                    byte_index += 1;
                    if byte_index >= 2 {
                        sample_tx
                            .send(sample as f32 / i16::MAX as f32)
                            .await
                            .unwrap();
                        byte_index = 0;
                        sample = 0i16;
                    }
                }
            }
        });

        Self {}
    }
}

impl eframe::App for EguiApp {
    fn update(&mut self, ctx: &egui::Context, _frame: &mut eframe::Frame) {
        egui::CentralPanel::default().show(ctx, |_ui| {});
    }
}

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
}