use crate::circular_buffer_source::CircularBufferSource;
use aws_lc_rs::aead::{AES_256_GCM, Aad, NONCE_LEN, Nonce, RandomizedNonceKey};
use aws_lc_rs::error::Unspecified;
use log::{debug, error, info, trace};
use prost::{DecodeError, Message};
use std::io::Write;
use tap_interface::tap::message::TapMessage;
use thiserror::Error;
use tokio::net::UdpSocket;
use tokio::select;
use tokio_util::sync::CancellationToken;

pub const SAMPLE_RATE: u32 = 48_000;

#[derive(Debug, Error)]
#[allow(clippy::enum_variant_names)]
pub enum Error {
    #[error("IO Error {0}")]
    IOError(#[from] std::io::Error),
    #[error("Failed to decode audio in packet: {0}")]
    OpusError(#[from] opus::Error),
    #[error("Failed to decode packet: {0}")]
    DecodeError(#[from] DecodeError),
}

pub struct UdpSource {
    udp: UdpSocket,
    source: CircularBufferSource,
    decoder: opus::Decoder,
    expected_seq: u32,
    key: RandomizedNonceKey,
}

impl UdpSource {
    pub fn new(udp_socket: UdpSocket, source: CircularBufferSource, key: &[u8]) -> Self {
        Self {
            udp: udp_socket,
            source,
            decoder: opus::Decoder::new(SAMPLE_RATE, opus::Channels::Stereo).unwrap(),
            expected_seq: 0,
            key: RandomizedNonceKey::new(&AES_256_GCM, key).unwrap(),
        }
    }

    fn decode_data(
        &mut self,
        input_buffer: &[u8],
        output_buffer: &mut [f32],
    ) -> Result<usize, Error> {
        let sample_size = self
            .decoder
            .decode_float(input_buffer, output_buffer, false)?;

        Ok(sample_size * 2)
    }

    fn write_audio_data_to_source(&mut self, samples: &[f32]) -> Result<(), Error> {
        let mut sample_bytes = Vec::with_capacity(size_of_val(samples));

        for sample in samples {
            let bytes = sample.to_le_bytes();

            sample_bytes.extend_from_slice(bytes.as_slice());
        }

        self.source.write_all(&sample_bytes)?;
        Ok(())
    }

    async fn recv_message(&mut self, buffer: &mut [u8]) -> Result<usize, std::io::Error> {
        self.udp.recv(buffer).await
    }

    fn decrypt<'a>(&'a mut self, buffer: &'a mut [u8]) -> Result<&'a mut [u8], Unspecified> {
        let nonce = &buffer[0..NONCE_LEN];
        let nonce = Nonce::try_assume_unique_for_key(nonce)?;

        self.key
            .open_in_place(nonce, Aad::empty(), &mut buffer[NONCE_LEN..])
    }

    pub async fn worker(&mut self, cancellation_token: CancellationToken) {
        info!("Starting TAP UDP Endpoint");
        let mut buffer = vec![0; 1024 * 32];
        let mut decode_data = vec![0.0; tap_interface::STEREO_FRAME_SIZE];
        loop {
            let msg_size = select! {
                res = self.recv_message(&mut buffer) => {


                    match res {
                        Ok(msg_size) => msg_size,
                        Err(err) => {
                            error!("Failed to recv udp packet: {err:?}");
                            continue;
                        }
                    }
                }
                _ = cancellation_token.cancelled() => {
                    info!("Closing TAP UDP Endpoint...");
                    return;
                }
            };

            let decrypted_msg = match self.decrypt(&mut buffer[0..msg_size]) {
                Ok(msg) => msg,
                Err(_err) => {
                    trace!("Failed to decrypt message");
                    continue;
                }
            };

            let tap_message = match TapMessage::decode(&*decrypted_msg) {
                Ok(tap_msg) => tap_msg,
                Err(err) => {
                    error!("Failed decode proto message: {err:?}");
                    continue;
                }
            };

            if tap_message.seq != self.expected_seq {
                debug!(
                    "Mismatch sequence count found expected '{}': {tap_message:?}",
                    tap_message.seq
                )
            }

            (self.expected_seq, _) = tap_message.seq.overflowing_add(1);

            let decoded_samples = match self.decode_data(&tap_message.audio_data, &mut decode_data)
            {
                Ok(decoded_samples) => decoded_samples,
                Err(err) => {
                    error!("Failed to decode audio data: {err:?}");
                    continue;
                }
            };

            if let Err(err) = self.write_audio_data_to_source(&decode_data[..decoded_samples]) {
                error!("Failed to write data to source: {err:?}");
            }
        }
    }
}