pizzicato 0.1.0

use std::ffi::CString;

use anyhow::Result;

use rsmpeg::avcodec::{ AVCodec, AVCodecContext };
use rsmpeg::avformat::{ AVFormatContextInput, AVFormatContextOutput };
use rsmpeg::avutil::{ AVAudioFifo, AVFrame, AVSamples, ra };
use rsmpeg::error::RsmpegError;
use rsmpeg::ffi;
use rsmpeg::swresample::SwrContext;

pub fn transcode(
	output_ctx: &mut AVFormatContextOutput,
	input_file: &str,
	codec: ffi::AVCodecID,
) -> Result<()> {
	let mut input_ctx = AVFormatContextInput::open(&*CString::new(input_file)?)?;
	let (audio_index, decode_codec) = input_ctx
		.find_best_stream(ffi::AVMediaType_AVMEDIA_TYPE_AUDIO)?
		.ok_or(anyhow!("No audio stream in audio file"))?;
	
	let mut decode_ctx = AVCodecContext::new(&decode_codec);

	decode_ctx.apply_codecpar(
		input_ctx
			.streams()
			.get(audio_index)
			.unwrap()
			.codecpar()
	)?;

	decode_ctx.open(None)?;

	let encode_codec = AVCodec::find_encoder(codec)
		.ok_or(anyhow!("Failed to find encoder"))?;
	
	let mut encode_ctx = AVCodecContext::new(&encode_codec);

	encode_ctx.set_channels(decode_ctx.channels);
	encode_ctx.set_channel_layout(decode_ctx.channel_layout);
	encode_ctx.set_bit_rate(decode_ctx.bit_rate);
	encode_ctx.set_sample_rate(decode_ctx.sample_rate);
	encode_ctx.set_sample_fmt(encode_codec.sample_fmts().unwrap()[0]);

	encode_ctx.open(None)?;

	let mut stream = output_ctx.new_stream();
	stream.set_codecpar(encode_ctx.extract_codecpar());
	stream.set_time_base(ra(decode_ctx.sample_rate, 1));
	drop(stream);

	let mut resample_ctx = SwrContext::new(
		encode_ctx.channel_layout,
		encode_ctx.sample_fmt,
		encode_ctx.sample_rate,
		decode_ctx.channel_layout,
		decode_ctx.sample_fmt,
		decode_ctx.sample_rate,
	).ok_or(anyhow!("Failed to init swr"))?;
	resample_ctx.init()?;

	let mut fifo = AVAudioFifo::new(encode_ctx.sample_fmt, encode_ctx.channels, 1);

	output_ctx.write_header()?;

	let mut pts = 0;

	loop {
		loop {
			if fifo.size() >= encode_ctx.frame_size {
				break
			}

			let packet = match input_ctx.read_packet()? {
				Some(packet) => packet,
				None => break,
			};

			if packet.stream_index as usize != audio_index {
				continue
			}

			decode_ctx.send_packet(Some(&packet))?;

			loop {
				let frame = match decode_ctx.receive_frame() {
					Ok(frame) => frame,
					Err(
						RsmpegError::DecoderDrainError |
						RsmpegError::DecoderFlushedError
					) => break,
					Err(err) => bail!(err),
				};

				let mut samples = AVSamples::new(
					encode_ctx.channels,
					frame.nb_samples,
					encode_ctx.sample_fmt,
					0,
				).ok_or(anyhow!("Failed to create samples buffer"))?;

				unsafe {
					resample_ctx.convert(
						&mut samples,
						frame.extended_data as *const _,
						0,
					)?;
				}

				fifo.realloc(fifo.size() + frame.nb_samples);

				let written = unsafe {
					fifo.write(samples.audio_data.as_ptr(), frame.nb_samples)?
				};
				if written < frame.nb_samples {
					bail!("Failed to write all samples");
				}
			}
		}

		if fifo.size() < encode_ctx.frame_size {
			break
		}

		while fifo.size() >= encode_ctx.frame_size {
			let nb_samples = fifo.size()
				.min(encode_ctx.frame_size);
			
			let mut frame = AVFrame::new();
			frame.set_nb_samples(nb_samples);
			frame.set_channel_layout(encode_ctx.channel_layout);
			frame.set_format(encode_ctx.sample_fmt);
			frame.set_sample_rate(encode_ctx.sample_rate);

			frame.alloc_buffer()?;

			let read = unsafe {
				fifo.read(frame.data_mut().as_mut_ptr(), nb_samples)?
			};
			if read < nb_samples {
				bail!("Failed to read all samples");
			}

			pts = encode_audio_frame(Some(frame), output_ctx, &mut encode_ctx, pts)?;
		}
	}

	encode_audio_frame(None, output_ctx, &mut encode_ctx, pts)?;

	output_ctx.write_trailer()?;

	Ok(())
}

fn encode_audio_frame(
	mut frame: Option<AVFrame>,
	output_ctx: &mut AVFormatContextOutput,
	encode_ctx: &mut AVCodecContext,
	mut pts: i64,
) -> Result<i64> {
	if let Some(frame) = frame.as_mut() {
		frame.set_pts(pts);
		pts += frame.nb_samples as i64;
	}

	encode_ctx.send_frame(frame.as_ref())?;

	loop {
		let mut packet = match encode_ctx.receive_packet() {
			Ok(packet) => packet,
			Err(
				RsmpegError::EncoderDrainError |
				RsmpegError::EncoderFlushedError
			) => break,
			Err(err) => bail!(err),
		};

		output_ctx.write_frame(&mut packet)?;
	}

	Ok(pts)
}