#![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 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, 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};
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 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 {}

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>>,
        mut sample_sender: T,
        mut eye_sender: Sender<Vec<f32>>,
    ) {
        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, &mut eye_sender) =>
                        {
                            // Flush channel
                            while sample_stream.recv().await.is_some() {};
                            match x
                            {
                                Err(()) => {continue;},
                                Ok(Frame::Ack) =>
                                {
                                    resend = None;
                                }
                                Ok(Frame::Data(data)) =>
                                {
                                    let _ = rx_stream.send(data).await;
                                    tokio::time::sleep(Duration::from_secs(1)).await;
                                    Self::transmit(Frame::Ack, &mut sample_sender).await;
                                }
                            }
                        },
                        _ = tokio::time::sleep(Duration::from_secs(100)) => {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()), &mut sample_sender).await;
                        resend = Some(data);
                    }
                }
            }
        }
    }

    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<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));
        samples_sender.open_link().await;
        for (m, up) in modulator.zip(up_lo) {
            let sample = m * up;
            samples_sender.send_sample(sample.re).await;
        }
        samples_sender.close_link().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::Ack)) = frame_opt {
                        println!("Got ack");
                        return Ok(Frame::Ack);
                    }

                    if let Ok(Some(Frame::Data(ref frame_data))) = frame_opt {
                        println!("Got data");
                        return Ok(Frame::Data(frame_data.to_vec()));
                    }

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

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

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 != 0xC4 && byte != 0x4C && !self.started {
            println!("Wrong type {:x}", byte);
            self.started = true;
            return Err(());
        }

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

        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::Data(
                    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);

        // Command
        match self {
            Frame::Data(x) => {
                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);
            }
            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() {
    //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 {
    eye_receiver_a: Receiver<Vec<f32>>,
    eyes_a: VecDeque<Vec<f32>>,
}
impl EguiApp {
    fn new(_cc: &eframe::CreationContext<'_>) -> Self {
        EguiApp {}
    }
}

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.eye_receiver_a.try_recv() {
                self.eyes_a.push_back(eye);
            }
            while self.eyes_a.len() > max_eyes {
                self.eyes_a.pop_front();
            }
        });
    }
}

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
}