use std::{
    collections::BTreeMap,
    net::UdpSocket,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc, Mutex,
    },
    thread,
    time::{Duration, Instant},
};

use chacha20poly1305::{aead::KeyInit, ChaCha20Poly1305, Key};
use cpal::traits::{DeviceTrait, HostTrait, StreamTrait};
use ringbuf::{
    storage::Heap,
    traits::{Consumer, Observer, Producer, Split},
    CachingCons, CachingProd, HeapRb, SharedRb,
};
use tauri::State;

use crate::{commands::config::ConfigState, crypto::SessionCipher};

const TARGET_SAMPLE_RATE: u32 = 48_000;
const PACKET_SAMPLES: usize = 960;
const MAX_PACKET_SIZE: usize = 4096;

const INITIAL_PACKET_CUSHION: usize = 3;

const VAD_THRESHOLD: f32 = 0.01;
const VAD_HANGOVER_FRAMES: usize = 10;

type AudioProducer = Arc<Mutex<CachingProd<Arc<SharedRb<Heap<f32>>>>>>;
type AudioConsumer = Arc<Mutex<CachingCons<Arc<SharedRb<Heap<f32>>>>>>;

#[derive(Clone, Default)]
pub struct VoiceState {
    pub inner: Arc<VoiceStateInner>,
}

#[derive(Default)]
pub struct VoiceStateInner {
    pub session: Mutex<Option<VoiceSession>>,
}

pub struct VoiceSession {
    pub input_stream: cpal::Stream,
    pub output_stream: Arc<Mutex<cpal::Stream>>,
    pub socket: Arc<UdpSocket>,
    pub pin: u64,
    pub shutdown: Arc<AtomicBool>,

    pub current_input_device: Option<String>,
    pub current_output_device: Option<String>,

    pub producer_in: AudioProducer,
    pub consumer_out: AudioConsumer,
}

// Simple Linear Resampler for real-time audio conversion
struct LinearResampler {
    phase: f64,
}

impl LinearResampler {
    fn new() -> Self {
        Self { phase: 0.0 }
    }

    /// Resamples dynamic buffers from `src_rate` to `dst_rate`
    fn process(&mut self, input: &[f32], src_rate: u32, dst_rate: u32, output: &mut Vec<f32>) {
        if src_rate == dst_rate {
            output.extend_from_slice(input);
            return;
        }

        let ratio = src_rate as f64 / dst_rate as f64;
        while self.phase < input.len() as f64 {
            let idx = self.phase as usize;
            let frac = (self.phase - idx as f64) as f32;
            let next_idx = (idx + 1).min(input.len() - 1);

            let sample = input[idx] * (1.0 - frac) + input[next_idx] * frac;
            output.push(sample);

            self.phase += ratio;
        }

        self.phase -= input.len() as f64;
    }
}

impl VoiceSession {
    pub fn update_input_device(
        &mut self,
        device_name: Option<String>,
        state_inner: Arc<VoiceStateInner>,
    ) -> Result<(), String> {
        let host = cpal::default_host();
        let device = match &device_name {
            Some(name) => host
                .input_devices()
                .map_err(|e| e.to_string())?
                .find(|d| {
                    d.description()
                        .ok()
                        .map(|x| x.name() == *name)
                        .unwrap_or(false)
                })
                .ok_or_else(|| format!("Input device '{name}' not found"))?,
            None => host
                .default_input_device()
                .ok_or("No default input device")?,
        };

        let input_config = device
            .default_input_config()
            .map_err(|e| format!("Failed to get default input config: {e}"))?
            .config();

        let producer = Arc::clone(&self.producer_in);
        let new_stream = build_input_stream(&device, input_config, producer, state_inner)?;

        new_stream
            .play()
            .map_err(|e| format!("Failed to play input stream: {e}"))?;

        self.input_stream = new_stream;
        self.current_input_device = device_name;
        Ok(())
    }

    pub fn update_output_device(
        &mut self,
        device_name: Option<String>,
        state_inner: Arc<VoiceStateInner>,
    ) -> Result<(), String> {
        let host = cpal::default_host();
        let device = match &device_name {
            Some(name) => host
                .output_devices()
                .map_err(|e| e.to_string())?
                .find(|d| {
                    d.description()
                        .ok()
                        .map(|x| x.name() == *name)
                        .unwrap_or(false)
                })
                .ok_or_else(|| format!("Output device '{name}' not found"))?,
            None => host
                .default_output_device()
                .ok_or("No default output device")?,
        };

        let output_config = device
            .default_output_config()
            .map_err(|e| format!("Failed to get default output config: {e}"))?
            .config();

        let consumer = Arc::clone(&self.consumer_out);
        let new_stream = build_output_stream(&device, output_config, consumer, state_inner)?;
        new_stream
            .play()
            .map_err(|e| format!("Failed to play output stream: {e}"))?;

        if let Ok(mut active_stream) = self.output_stream.lock() {
            *active_stream = new_stream;
        }

        self.current_output_device = device_name;
        Ok(())
    }
}

// Helper to build normalized Input Stream (Resampled & Downmixed to 48kHz Mono)
fn build_input_stream(
    device: &cpal::Device,
    config: cpal::StreamConfig,
    producer: AudioProducer,
    state_inner: Arc<VoiceStateInner>,
) -> Result<cpal::Stream, String> {
    let native_sample_rate = config.sample_rate;
    let channels = config.channels as usize;
    let mut resampler = LinearResampler::new();
    let mut mono_buffer = Vec::with_capacity(2048);
    let mut resampled_buffer = Vec::with_capacity(2048);

    let inner_input_err = Arc::clone(&state_inner);

    device
        .build_input_stream(
            config,
            move |data: &[f32], _| {
                mono_buffer.clear();
                resampled_buffer.clear();

                // Downmix channels to mono
                for chunk in data.chunks_exact(channels) {
                    let sum: f32 = chunk.iter().sum();
                    mono_buffer.push(sum / channels as f32);
                }

                // Resample to 48kHz standard target
                resampler.process(
                    &mono_buffer,
                    native_sample_rate,
                    TARGET_SAMPLE_RATE,
                    &mut resampled_buffer,
                );

                if let Ok(mut prod) = producer.lock() {
                    let _ = prod.push_slice(&resampled_buffer);
                }
            },
            move |err| {
                eprintln!("[vc] Input error: {err}. Attempting recovery...");
                if let Ok(mut lock) = inner_input_err.session.lock() {
                    if let Some(session) = lock.as_mut() {
                        let target_device = session.current_input_device.clone();
                        let _ = session
                            .update_input_device(target_device, Arc::clone(&inner_input_err));
                    }
                }
            },
            None,
        )
        .map_err(|e| e.to_string())
}

// Helper to build normalized Output Stream (48kHz Mono -> Device Native Channels & Rate)
fn build_output_stream(
    device: &cpal::Device,
    config: cpal::StreamConfig,
    consumer: AudioConsumer,
    state_inner: Arc<VoiceStateInner>,
) -> Result<cpal::Stream, String> {
    let native_sample_rate = config.sample_rate;
    let channels = config.channels as usize;
    let mut resampler = LinearResampler::new();
    let mut raw_mono_samples = Vec::with_capacity(2048);
    let mut resampled_mono = Vec::with_capacity(2048);
    let mut last_sample = 0.0f32;

    let inner_output_err = Arc::clone(&state_inner);

    device
        .build_output_stream(
            config,
            move |data: &mut [f32], _| {
                let required_mono_samples = (data.len() / channels) * TARGET_SAMPLE_RATE as usize
                    / native_sample_rate as usize;

                raw_mono_samples.clear();
                resampled_mono.clear();

                if let Ok(mut cons) = consumer.lock() {
                    for _ in 0..required_mono_samples {
                        if let Some(s) = cons.try_pop() {
                            last_sample = s;
                            raw_mono_samples.push(s);
                        } else {
                            // Exponential decay to prevent clicking when underflowing
                            last_sample *= 0.92;
                            raw_mono_samples.push(last_sample);
                        }
                    }
                }

                // Resample from 48kHz mono to target native output rate
                resampler.process(
                    &raw_mono_samples,
                    TARGET_SAMPLE_RATE,
                    native_sample_rate,
                    &mut resampled_mono,
                );

                // Interleave mono into hardware channels
                let mut res_idx = 0;
                let mut out_idx = 0;
                while out_idx < data.len() && res_idx < resampled_mono.len() {
                    let mono_val = resampled_mono[res_idx];
                    for ch in 0..channels {
                        if out_idx + ch < data.len() {
                            data[out_idx + ch] = mono_val;
                        }
                    }
                    out_idx += channels;
                    res_idx += 1;
                }
            },
            move |err| {
                eprintln!("[vc] Output error: {err}. Attempting recovery...");
                if let Ok(mut lock) = inner_output_err.session.lock() {
                    if let Some(session) = lock.as_mut() {
                        let target_device = session.current_output_device.clone();
                        let _ = session
                            .update_output_device(target_device, Arc::clone(&inner_output_err));
                    }
                }
            },
            None,
        )
        .map_err(|e| e.to_string())
}

#[tauri::command]
pub fn list_input_devices() -> Result<Vec<String>, String> {
    let host = cpal::default_host();
    Ok(host
        .input_devices()
        .map_err(|e| e.to_string())?
        .filter_map(|d| d.description().ok().map(|x| x.name().to_string()))
        .collect())
}

#[tauri::command]
pub fn list_output_devices() -> Result<Vec<String>, String> {
    let host = cpal::default_host();
    Ok(host
        .output_devices()
        .map_err(|e| e.to_string())?
        .filter_map(|d| d.description().ok().map(|x| x.name().to_string()))
        .collect())
}

#[tauri::command]
pub fn disconnect_from_vc(voice_state: State<'_, VoiceState>) -> Result<(), String> {
    let mut lock = voice_state
        .inner
        .session
        .lock()
        .map_err(|e| e.to_string())?;

    if let Some(session) = lock.take() {
        session.shutdown.store(true, Ordering::SeqCst);
        let _ = session.input_stream.pause();
        if let Ok(output) = session.output_stream.lock() {
            let _ = output.pause();
        }
    }

    Ok(())
}

#[tauri::command]
pub fn connect_to_vc(
    hostname: String,
    pin: u64,
    shared_secret: Vec<u8>, // output of js `x25519.getSharedSecret`
    config_state: State<'_, ConfigState>,
    voice_state: State<'_, VoiceState>,
) -> Result<(), String> {
    disconnect_from_vc(voice_state.clone())?;

    let key = Key::try_from(shared_secret.as_slice()).map_err(|v| v.to_string())?;
    let cipher = Arc::new(Mutex::new(SessionCipher::new(ChaCha20Poly1305::new(&key))));

    let config = config_state.0.lock().unwrap().clone();
    let state_inner = Arc::clone(&voice_state.inner);

    let socket = Arc::new(
        UdpSocket::bind("0.0.0.0:0").map_err(|e| format!("Failed to bind UDP socket: {e}"))?,
    );
    socket
        .connect(&hostname)
        .map_err(|e| format!("Failed to connect UDP socket: {e}"))?;
    socket
        .set_read_timeout(Some(Duration::from_millis(5)))
        .map_err(|e| e.to_string())?;
    socket
        .send(&pin.to_be_bytes())
        .map_err(|e| format!("Failed to send pin: {e}"))?;

    let host = cpal::default_host();

    let input_device = match &config.input_device_name {
        Some(name) => host
            .input_devices()
            .map_err(|e| e.to_string())?
            .find(|d| {
                d.description()
                    .ok()
                    .map(|x| x.name() == *name)
                    .unwrap_or(false)
            })
            .ok_or("Input device not found")?,
        None => host
            .default_input_device()
            .ok_or("No default input device")?,
    };

    let output_device = match &config.output_device_name {
        Some(name) => host
            .output_devices()
            .map_err(|e| e.to_string())?
            .find(|d| {
                d.description()
                    .ok()
                    .map(|x| x.name() == *name)
                    .unwrap_or(false)
            })
            .ok_or("Output device not found")?,
        None => host
            .default_output_device()
            .ok_or("No default output device")?,
    };

    let shutdown = Arc::new(AtomicBool::new(false));

    // Ring Buffer Setup
    let rb_out = HeapRb::<f32>::new(19200);
    let (producer_out, consumer_out) = rb_out.split();
    let mut producer_out = producer_out;
    let shared_consumer_out = Arc::new(Mutex::new(consumer_out));

    let output_config = output_device
        .default_output_config()
        .map_err(|e| format!("Failed to get default output config: {e}"))?
        .config();

    let output_stream = build_output_stream(
        &output_device,
        output_config,
        Arc::clone(&shared_consumer_out),
        Arc::clone(&state_inner),
    )?;
    output_stream
        .play()
        .map_err(|e| format!("Failed to start output: {e}"))?;
    let output_stream = Arc::new(Mutex::new(output_stream));

    let rb_in = HeapRb::<f32>::new(19200);
    let (producer_in, mut consumer_in) = rb_in.split();
    let shared_producer_in = Arc::new(Mutex::new(producer_in));

    let input_config = input_device
        .default_input_config()
        .map_err(|e| format!("Failed to get default input config: {e}"))?
        .config();

    let input_stream = build_input_stream(
        &input_device,
        input_config,
        Arc::clone(&shared_producer_in),
        Arc::clone(&state_inner),
    )?;

    // Sender Thread
    {
        let input_socket = socket.clone();
        let shutdown = shutdown.clone();
        let cipher = cipher.clone();

        thread::spawn(move || {
            let mut sequence = 0u32;
            let mut frame_buf = vec![0.0f32; PACKET_SAMPLES];
            let mut hangover_counter = 0;

            while !shutdown.load(Ordering::Relaxed) {
                if consumer_in.occupied_len() >= PACKET_SAMPLES {
                    let _ = consumer_in.pop_slice(&mut frame_buf);

                    let sum_squares: f32 = frame_buf.iter().map(|&s| s * s).sum();
                    let rms = (sum_squares / PACKET_SAMPLES as f32).sqrt();

                    let is_speaking = if rms >= VAD_THRESHOLD {
                        hangover_counter = VAD_HANGOVER_FRAMES;
                        true
                    } else if hangover_counter > 0 {
                        hangover_counter -= 1;
                        true
                    } else {
                        false
                    };

                    if is_speaking {
                        // sequence goes INSIDE the plaintext now, prefixed before the PCM
                        let mut plaintext = Vec::with_capacity(4 + PACKET_SAMPLES * 2);
                        plaintext.extend_from_slice(&sequence.to_be_bytes());

                        for sample in frame_buf.iter() {
                            let pcm = (sample.clamp(-1.0, 1.0) * 32767.0) as i16;
                            plaintext.extend_from_slice(&pcm.to_be_bytes());
                        }

                        let Ok(net_packet) = cipher.lock().unwrap().encrypt(&plaintext) else {
                            eprintln!("[vc] failed to encrypt outgoing packet");
                            sequence = sequence.wrapping_add(1);
                            continue;
                        };

                        let _ = input_socket.send(&net_packet);
                        sequence = sequence.wrapping_add(1);
                    }
                } else {
                    thread::sleep(Duration::from_millis(2));
                }
            }
        });
    }

    // Receiver Thread
    {
        let socket = socket.clone();
        let shutdown = shutdown.clone();
        let cipher = cipher.clone();

        thread::spawn(move || {
            let mut packets: BTreeMap<u32, Vec<f32>> = BTreeMap::new();
            let mut expected: Option<u32> = None;
            let mut is_prebuffering = true;
            let mut last_good_frame = vec![0.0f32; PACKET_SAMPLES];
            let mut udp_buffer = [0u8; MAX_PACKET_SIZE];
            let mut next_frame_time = Instant::now();

            while !shutdown.load(Ordering::Relaxed) {
                if let Ok(len) = socket.recv(&mut udp_buffer) {
                    match cipher.lock().unwrap().decrypt(&udp_buffer[..len]) {
                        Ok(plaintext) => {
                            if plaintext.len() < 4 {
                                eprintln!("[vc] dropped packet: too short after decrypt");
                            } else {
                                let seq = u32::from_be_bytes(plaintext[..4].try_into().unwrap());
                                let pcm = &plaintext[4..];

                                let samples: Vec<f32> = pcm
                                    .chunks_exact(2)
                                    .map(|c| i16::from_be_bytes([c[0], c[1]]) as f32 / 32768.0)
                                    .collect();

                                packets.insert(seq, samples);
                            }
                        }
                        Err(_) => {
                            eprintln!("[vc] dropped packet: decryption failed");
                        }
                    }
                }

                if is_prebuffering {
                    if packets.len() >= INITIAL_PACKET_CUSHION {
                        expected = packets.keys().next().copied();
                        is_prebuffering = false;
                        next_frame_time = Instant::now();
                    } else {
                        thread::sleep(Duration::from_millis(1));
                        continue;
                    }
                }

                if Instant::now() >= next_frame_time {
                    if let Some(seq) = expected {
                        if let Some(samples) = packets.remove(&seq) {
                            if samples.len() == PACKET_SAMPLES {
                                last_good_frame.copy_from_slice(&samples);
                            } else {
                                last_good_frame.clear();
                                last_good_frame
                                    .extend(samples.iter().take(PACKET_SAMPLES).copied());
                                if last_good_frame.len() < PACKET_SAMPLES {
                                    last_good_frame.resize(PACKET_SAMPLES, 0.0);
                                }
                            }

                            let _ = producer_out.push_slice(&last_good_frame);
                            expected = Some(seq.wrapping_add(1));
                        } else if packets.keys().any(|&x| x > seq) {
                            for sample in last_good_frame.iter_mut() {
                                *sample *= 0.65;
                            }
                            let _ = producer_out.push_slice(&last_good_frame);
                            expected = Some(seq.wrapping_add(1));
                        } else if packets.is_empty() {
                            is_prebuffering = true;
                        }
                    }
                    next_frame_time += Duration::from_millis(20);
                }

                thread::sleep(Duration::from_millis(1));
            }
        });
    }

    input_stream
        .play()
        .map_err(|e| format!("Failed to start input: {e}"))?;

    *state_inner.session.lock().unwrap() = Some(VoiceSession {
        input_stream,
        output_stream,
        socket,
        pin,
        shutdown,
        current_input_device: config.input_device_name,
        current_output_device: config.output_device_name,
        producer_in: shared_producer_in,
        consumer_out: shared_consumer_out,
    });

    Ok(())
}

impl Drop for VoiceSession {
    fn drop(&mut self) {
        self.shutdown.store(true, Ordering::SeqCst);

        let _ = self.input_stream.pause();
        if let Ok(output) = self.output_stream.lock() {
            let _ = output.pause();
        }

        eprintln!("[vc] VoiceSession dropped and audio streams paused.");
    }
}