pile/crates/pile-value/src/extract/item/flac.rs

use mime::Mime;
use pile_config::Label;
use pile_flac::{FlacBlock, FlacDecodeError, FlacReader};
use std::{
	collections::HashMap,
	io::BufReader,
	sync::{Arc, OnceLock},
};
use tracing::trace;

use crate::{
	extract::traits::{ExtractState, ListExtractor, ObjectExtractor},
	value::{Item, PileValue, SyncReadBridge},
};

pub struct FlacImagesExtractor {
	item: Item,
	cached_count: OnceLock<usize>,
}

impl FlacImagesExtractor {
	pub fn new(item: &Item) -> Self {
		Self {
			item: item.clone(),
			cached_count: OnceLock::new(),
		}
	}

	async fn get_count(&self) -> Result<usize, std::io::Error> {
		let reader = SyncReadBridge::new_current(self.item.read().await?);
		let count = tokio::task::spawn_blocking(move || {
			let reader = FlacReader::new(BufReader::new(reader));
			let mut count = 0usize;
			for block in reader {
				match block {
					Ok(FlacBlock::AudioFrame(_)) => break,
					Ok(FlacBlock::Picture(_)) => count += 1,
					Err(FlacDecodeError::IoError(err)) => return Err(err),
					Err(_) => return Ok(0),
					_ => {}
				}
			}
			Ok::<_, std::io::Error>(count)
		})
		.await??;

		return Ok(count);
	}

	fn mime_ok(&self, state: &ExtractState) -> bool {
		if state.ignore_mime {
			return true;
		}
		let essence = self.item.mime().essence_str();
		essence == "audio/flac" || essence == "audio/x-flac"
	}
}

#[async_trait::async_trait]
impl ListExtractor for FlacImagesExtractor {
	async fn get(
		&self,
		state: &ExtractState,
		mut idx: usize,
	) -> Result<Option<PileValue>, std::io::Error> {
		trace!(
			key = self.item.key().as_str(),
			"Getting index {idx} from FlacImagesExtractor",
		);

		if !self.mime_ok(state) {
			return Ok(None);
		}

		let key = self.item.key();
		let reader = SyncReadBridge::new_current(self.item.read().await?);
		let image = tokio::task::spawn_blocking(move || {
			let reader = FlacReader::new(BufReader::new(reader));
			let mut out: Option<(Mime, Vec<u8>)> = None;
			'blocks: for block in reader {
				match block {
					Ok(FlacBlock::AudioFrame(_)) => break,
					Ok(FlacBlock::Picture(picture)) => {
						if idx > 0 {
							idx -= 1;
							continue;
						}

						out = Some((picture.mime, picture.img_data));
						break 'blocks;
					}

					Err(FlacDecodeError::IoError(err)) => return Err(err),
					Err(error) => {
						trace!(
							message = "Could not parse FLAC images",
							key = key.as_str(),
							?error
						);
						return Ok(None);
					}

					_ => {}
				}
			}
			Ok::<_, std::io::Error>(out)
		})
		.await
		.map_err(std::io::Error::other)??;

		Ok(image.map(|(mime, data)| PileValue::Blob {
			mime,
			bytes: Arc::new(data),
		}))
	}

	async fn len(&self, state: &ExtractState) -> Result<usize, std::io::Error> {
		if !self.mime_ok(state) {
			return Ok(0);
		}

		if let Some(x) = self.cached_count.get() {
			return Ok(*x);
		}

		let count = self.get_count().await?;
		return Ok(*self.cached_count.get_or_init(|| count));
	}
}

pub struct FlacExtractor {
	item: Item,
	output: OnceLock<HashMap<Label, PileValue>>,
	images: PileValue,
}

impl FlacExtractor {
	pub fn new(item: &Item) -> Self {
		Self {
			item: item.clone(),
			output: OnceLock::new(),
			images: PileValue::ListExtractor(Arc::new(FlacImagesExtractor::new(item))),
		}
	}

	async fn get_inner(&self) -> Result<&HashMap<Label, PileValue>, std::io::Error> {
		if let Some(x) = self.output.get() {
			return Ok(x);
		}

		trace!(
			message = "Reading FLAC tags",
			key = self.item.key().as_str()
		);

		let key = self.item.key();
		let reader = SyncReadBridge::new_current(self.item.read().await?);
		let output = tokio::task::spawn_blocking(move || {
			let reader = FlacReader::new(BufReader::new(reader));
			let mut output: HashMap<Label, Vec<PileValue>> = HashMap::new();

			for block in reader {
				match block {
					Ok(FlacBlock::AudioFrame(_)) => break,
					Ok(FlacBlock::VorbisComment(comment)) => {
						for (k, v) in comment.comment.comments {
							if let Some(label) = Label::new(k.to_string().to_lowercase()) {
								output
									.entry(label)
									.or_default()
									.push(PileValue::String(Arc::new(v)));
							}
						}
					}

					Err(FlacDecodeError::IoError(err)) => return Err(err),
					Err(error) => {
						trace!(
							message = "Could not parse FLAC metadata",
							key = key.as_str(),
							?error
						);
						return Ok(HashMap::new());
					}

					_ => {}
				}
			}

			let output: HashMap<Label, PileValue> = output
				.into_iter()
				.map(|(k, v)| (k, PileValue::Array(Arc::new(v))))
				.collect();

			Ok::<HashMap<Label, PileValue>, std::io::Error>(output)
		})
		.await??;

		return Ok(self.output.get_or_init(|| output));
	}

	fn mime_ok(&self, state: &ExtractState) -> bool {
		if state.ignore_mime {
			return true;
		}
		let essence = self.item.mime().essence_str();
		essence == "audio/flac" || essence == "audio/x-flac"
	}
}

#[async_trait::async_trait]
impl ObjectExtractor for FlacExtractor {
	async fn field(
		&self,
		state: &ExtractState,
		name: &Label,
		args: Option<&str>,
	) -> Result<Option<PileValue>, std::io::Error> {
		if args.is_some() {
			return Ok(None);
		}

		if !self.mime_ok(state) {
			return Ok(None);
		}

		if name.as_str() == "images" {
			return Ok(Some(self.images.clone()));
		}

		Ok(self.get_inner().await?.get(name).cloned())
	}

	#[expect(clippy::unwrap_used)]
	async fn fields(&self) -> Result<Vec<Label>, std::io::Error> {
		Ok(self
			.get_inner()
			.await?
			.keys()
			.cloned()
			.chain([Label::new("images").unwrap()])
			.collect::<Vec<_>>())
	}
}