use std::{collections::VecDeque, path::PathBuf};

use cap_fail::fail;
use ffmpeg::{
    codec::{context, encoder},
    format::{self, sample::Type, Sample},
    threading::Config,
};

use crate::{
    data::{cast_bytes_to_f32_slice, AudioInfo, FFAudio, FFRational},
    pipeline::task::PipelineSinkTask,
    MediaError,
};

use super::audio::AudioEncoder;

pub struct AACEncoder {
    tag: &'static str,
    encoder: encoder::Audio,
    packet: ffmpeg::Packet,
    resampler: Option<ffmpeg::software::resampling::Context>,
    resampled_frame: FFAudio,
    buffer: Vec<VecDeque<u8>>,
    stream_index: usize,
}

impl AACEncoder {
    const OUTPUT_BITRATE: usize = 320 * 1000; // 128k
    const SAMPLE_FORMAT: Sample = Sample::F32(Type::Planar);

    pub fn factory(
        tag: &'static str,
        input_config: AudioInfo,
    ) -> impl FnOnce(&mut format::context::Output) -> Result<Self, MediaError> {
        move |o| Self::init(tag, input_config, o)
    }

    pub fn init(
        tag: &'static str,
        input_config: AudioInfo,
        output: &mut format::context::Output,
    ) -> Result<Self, MediaError> {
        let codec = encoder::find_by_name("aac")
            .ok_or(MediaError::TaskLaunch("Could not find AAC codec".into()))?;
        let mut encoder_ctx = context::Context::new_with_codec(codec);
        encoder_ctx.set_threading(Config::count(4));
        let mut encoder = encoder_ctx.encoder().audio()?;

        let rate = {
            let mut rates = codec
                .audio()
                .unwrap()
                .rates()
                .into_iter()
                .flatten()
                .collect::<Vec<_>>();
            rates.sort();

            let Some(&rate) = rates
                .iter()
                .find(|r| **r >= input_config.rate())
                .or(rates.first())
            else {
                return Err(MediaError::TaskLaunch(format!(
                    "Opus Codec does not support sample rate {}",
                    input_config.rate()
                )));
            };
            rate
        };

        let mut output_config = input_config.clone();
        output_config.sample_format = Self::SAMPLE_FORMAT;
        output_config.sample_rate = rate as u32;

        let resampler = if (
            input_config.sample_format,
            input_config.channel_layout(),
            input_config.sample_rate,
        ) != (
            output_config.sample_format,
            output_config.channel_layout(),
            output_config.sample_rate,
        ) {
            Some(
                ffmpeg::software::resampler(
                    (
                        input_config.sample_format,
                        input_config.channel_layout(),
                        input_config.sample_rate,
                    ),
                    (
                        output_config.sample_format,
                        output_config.channel_layout(),
                        output_config.sample_rate,
                    ),
                )
                .unwrap(),
            )
        } else {
            None
        };

        encoder.set_bit_rate(Self::OUTPUT_BITRATE);
        encoder.set_rate(rate);
        encoder.set_format(output_config.sample_format);
        encoder.set_channel_layout(output_config.channel_layout());
        encoder.set_time_base(output_config.time_base);

        let encoder = encoder.open()?;

        let mut output_stream = output.add_stream(codec)?;
        let stream_index = output_stream.index();
        output_stream.set_time_base(FFRational(1, output_config.rate()));
        output_stream.set_parameters(&encoder);

        Ok(Self {
            tag,
            buffer: vec![VecDeque::new(); 2],
            encoder,
            stream_index,
            packet: ffmpeg::Packet::empty(),
            resampled_frame: FFAudio::empty(),
            resampler,
        })
    }

    pub fn queue_frame(&mut self, frame: FFAudio, output: &mut format::context::Output) {
        let frame = if let Some(resampler) = &mut self.resampler {
            resampler.run(&frame, &mut self.resampled_frame).unwrap();
            &self.resampled_frame
        } else {
            &frame
        };

        for i in 0..frame.planes() {
            self.buffer[i]
                .extend(&frame.data(i)[0..frame_size_bytes(&frame) / frame.channels() as usize]);
        }

        let channel_size_bytes = self.encoder.frame_size() as usize * self.encoder.format().bytes();

        loop {
            if self.buffer[0].len() < channel_size_bytes {
                break;
            }

            let mut frame = FFAudio::new(
                self.encoder.format(),
                self.encoder.frame_size() as usize,
                self.encoder.channel_layout(),
            );

            for i in 0..frame.planes() {
                let bytes = self.buffer[i]
                    .drain(0..channel_size_bytes)
                    .collect::<Vec<_>>();

                frame.data_mut(i)[0..channel_size_bytes]
                    .copy_from_slice(&bytes[0..channel_size_bytes]);
            }

            self.encoder.send_frame(&frame).unwrap();

            self.process_packets(output);
        }
    }

    fn process_packets(&mut self, output: &mut format::context::Output) {
        while self.encoder.receive_packet(&mut self.packet).is_ok() {
            self.packet.set_stream(self.stream_index);
            self.packet.rescale_ts(
                self.encoder.time_base(),
                output.stream(self.stream_index).unwrap().time_base(),
            );
            self.packet.write_interleaved(output).unwrap();
        }
    }

    pub fn finish(&mut self, output: &mut format::context::Output) {
        let frame_size_bytes = self.encoder.frame_size() as usize
            * self.encoder.channels() as usize
            * self.encoder.format().bytes();

        if let Some(mut resampler) = self.resampler.take() {
            while resampler.delay().is_some() {
                resampler.flush(&mut self.resampled_frame).unwrap();
                if self.resampled_frame.samples() == 0 {
                    break;
                }

                for i in 0..self.resampled_frame.planes() {
                    self.buffer[i].extend(
                        &self.resampled_frame.data(0)[0..self.resampled_frame.samples()
                            * self.resampled_frame.format().bytes()],
                    );
                }

                while self.buffer.len() >= frame_size_bytes {
                    let mut frame = FFAudio::new(
                        self.encoder.format(),
                        self.encoder.frame_size() as usize,
                        self.encoder.channel_layout(),
                    );

                    for i in 0..frame.planes() {
                        let bytes = self.buffer[i]
                            .drain(0..frame_size_bytes)
                            .collect::<Vec<_>>();

                        frame.data_mut(0)[0..frame_size_bytes].copy_from_slice(&bytes);
                    }

                    self.encoder.send_frame(&frame).unwrap();

                    self.process_packets(output);
                }
            }

            while self.buffer[0].len() > 0 {
                let channel_size_bytes =
                    (frame_size_bytes / self.encoder.channels() as usize).min(self.buffer[0].len());
                let frame_size = channel_size_bytes / self.encoder.format().bytes();

                let mut frame = FFAudio::new(
                    self.encoder.format(),
                    frame_size,
                    self.encoder.channel_layout(),
                );

                for i in 0..frame.planes() {
                    let bytes = self.buffer[i]
                        .drain(0..channel_size_bytes)
                        .collect::<Vec<_>>();

                    frame.data_mut(i)[0..channel_size_bytes]
                        .copy_from_slice(&bytes[0..channel_size_bytes]);
                }

                self.encoder.send_frame(&frame).unwrap();

                self.process_packets(output);
            }
        }

        self.encoder.send_eof().unwrap();

        self.process_packets(output);
    }
}

impl AudioEncoder for AACEncoder {
    fn queue_frame(&mut self, frame: FFAudio, output: &mut format::context::Output) {
        self.queue_frame(frame, output);
    }

    fn finish(&mut self, output: &mut format::context::Output) {
        self.finish(output);
    }
}

fn frame_size_bytes(frame: &FFAudio) -> usize {
    frame.samples() * frame.format().bytes() * frame.channels() as usize
}
