pile/crates/pile-value/src/value/value.rs

use mime::Mime;
use pile_config::objectpath::{ObjectPath, PathSegment};
use serde_json::{Map, Value};
use smartstring::{LazyCompact, SmartString};
use std::sync::Arc;

use crate::{
	extract::{
		item::ItemExtractor,
		misc::{ArrayExtractor, MapExtractor, VecExtractor},
		string::StringExtractor,
		traits::{ExtractState, ListExtractor, ObjectExtractor},
	},
	value::Item,
};

/// An immutable, cheaply-cloneable, lazily-computed value.
/// Very similar to [serde_json::Value].
pub enum PileValue {
	Null,
	U64(u64),
	I64(i64),

	/// A string
	String(Arc<SmartString<LazyCompact>>),

	/// An array of values
	Array(Arc<Vec<PileValue>>),

	/// A binary blob
	Blob {
		mime: Mime,
		bytes: Arc<Vec<u8>>,
	},

	/// A lazily-computed map of {label: value}
	ObjectExtractor(Arc<dyn ObjectExtractor>),

	/// A lazily-computed array
	ListExtractor(Arc<dyn ListExtractor>),

	/// An pointer to an item in this dataset
	Item(Item),
}

impl Clone for PileValue {
	fn clone(&self) -> Self {
		match self {
			Self::Null => Self::Null,
			Self::U64(x) => Self::U64(*x),
			Self::I64(x) => Self::I64(*x),
			Self::String(x) => Self::String(x.clone()),
			Self::Array(x) => Self::Array(x.clone()),
			Self::ObjectExtractor(x) => Self::ObjectExtractor(x.clone()),
			Self::ListExtractor(x) => Self::ListExtractor(x.clone()),
			Self::Blob { mime, bytes } => Self::Blob {
				mime: mime.clone(),
				bytes: bytes.clone(),
			},
			Self::Item(i) => Self::Item(i.clone()),
		}
	}
}

impl PileValue {
	pub fn object_extractor(&self) -> Arc<dyn ObjectExtractor> {
		match self {
			Self::Null => Arc::new(MapExtractor::default()),
			Self::U64(_) => Arc::new(MapExtractor::default()),
			Self::I64(_) => Arc::new(MapExtractor::default()),
			Self::Array(_) => Arc::new(MapExtractor::default()),
			Self::String(s) => Arc::new(StringExtractor::new(s)),
			Self::Blob { .. } => Arc::new(MapExtractor::default()),
			Self::ListExtractor(_) => Arc::new(MapExtractor::default()),
			Self::ObjectExtractor(e) => e.clone(),
			Self::Item(i) => Arc::new(ItemExtractor::new(i)),
		}
	}

	pub fn list_extractor(&self) -> Arc<dyn ListExtractor> {
		match self {
			Self::Null => Arc::new(VecExtractor::default()),
			Self::U64(_) => Arc::new(VecExtractor::default()),
			Self::I64(_) => Arc::new(VecExtractor::default()),
			Self::Array(a) => Arc::new(ArrayExtractor::new(a.clone())),
			Self::String(_) => Arc::new(VecExtractor::default()),
			Self::Blob { .. } => Arc::new(VecExtractor::default()),
			Self::ListExtractor(e) => e.clone(),
			Self::ObjectExtractor(e) => {
				e.as_list().unwrap_or_else(|| Arc::new(VecExtractor::default()))
			}
			Self::Item(_) => Arc::new(VecExtractor::default()),
		}
	}

	pub async fn query(
		&self,
		state: &ExtractState,
		query: &ObjectPath,
	) -> Result<Option<Self>, std::io::Error> {
		let mut out: Option<PileValue> = Some(self.clone());

		for s in &query.segments {
			match s {
				PathSegment::Root => out = Some(self.clone()),
				PathSegment::Field { name, args } => {
					let e = match out.map(|x| x.object_extractor()) {
						Some(e) => e,
						None => {
							out = None;
							continue;
						}
					};

					out = e.field(state, name, args.as_deref()).await?;
				}

				PathSegment::Index(idx) => {
					let e = match out.map(|x| x.list_extractor()) {
						Some(e) => e,
						None => {
							out = None;
							continue;
						}
					};

					let idx = if *idx >= 0 {
						usize::try_from(*idx).ok()
					} else {
						usize::try_from(e.len(state).await? as i64 - idx).ok()
					};

					let idx = match idx {
						Some(idx) => idx,
						None => {
							out = None;
							continue;
						}
					};

					out = e.get(state, idx).await?;
				}
			}
		}

		return Ok(out.clone());
	}

	/// Like `to_json`, but counts populated fields instead of collecting values.
	///
	/// - Leaf values (non-null scalars, arrays, blobs) contribute `Some(1)`.
	/// - `Null` contributes `None`.
	/// - `ObjectExtractor` is recursed into; returns `Some(Object(map))` with
	///   only the fields that had data, or `None` if all fields were absent.
	/// - `Array` / `ListExtractor` are treated as opaque leaf values (not descended into).
	pub async fn count_fields(
		&self,
		state: &ExtractState,
	) -> Result<Option<Value>, std::io::Error> {
		Ok(match self {
			Self::Null => None,

			Self::U64(_) | Self::I64(_) | Self::String(_) | Self::Blob { .. } => {
				Some(Value::Number(1u64.into()))
			}

			Self::Array(x) => (!x.is_empty()).then(|| Value::Number(1u64.into())),
			Self::ListExtractor(x) => (x.len(state).await? > 0).then(|| Value::Number(1u64.into())),

			Self::ObjectExtractor(_) | Self::Item(_) => {
				let e = self.object_extractor();
				let keys = e.fields().await?;
				let mut map = Map::new();
				for k in &keys {
					let v = match e.field(state, k, None).await? {
						Some(x) => x,
						None => continue,
					};
					if let Some(counted) = Box::pin(v.count_fields(state)).await? {
						map.insert(k.to_string(), counted);
					}
				}
				if map.is_empty() {
					None
				} else {
					Some(Value::Object(map))
				}
			}
		})
	}

	pub fn as_str(&self) -> Option<&str> {
		match self {
			Self::String(x) => Some(x),
			_ => None,
		}
	}

	pub async fn to_json(&self, state: &ExtractState) -> Result<Value, std::io::Error> {
		Ok(match self {
			Self::Null => Value::Null,
			Self::U64(x) => Value::Number((*x).into()),
			Self::I64(x) => Value::Number((*x).into()),
			Self::String(x) => Value::String(x.to_string()),

			// TODO: replace with something meaningful?
			Self::Blob { mime, bytes } => {
				Value::String(format!("<Blob ({mime}, {} bytes)>", bytes.len()))
			}

			Self::Array(_) | Self::ListExtractor(_) => {
				let e = self.list_extractor();
				return e.to_json(state).await;
			}

			Self::ObjectExtractor(_) | Self::Item(_) => {
				let e = self.object_extractor();
				return e.to_json(state).await;
			}
		})
	}
}