json-ld 0.22.0

use std::collections::HashMap;

use btree_indexmap::{BTreeIndexMultiSet, BTreeIndexSet};
use json_syntax::ryu_js;
use rdf_syntax::{BlankIdBuf, Generator, Id};

use crate::{
	JsonLdLenient,
	object::{JsonLdObjectMut, JsonLdObjectRef},
};

pub trait JsonLdVisit {
	fn visit_with(&self, f: &mut impl FnMut(JsonLdObjectRef));

	fn visit(&self, mut f: impl FnMut(JsonLdObjectRef)) {
		self.visit_with(&mut f);
	}

	fn visit_mut_with(&mut self, f: &mut impl FnMut(JsonLdObjectMut));

	fn visit_mut(&mut self, mut f: impl FnMut(JsonLdObjectMut)) {
		self.visit_mut_with(&mut f);
	}

	fn relabel_with(&mut self, relabeling: &mut JsonLdRelabeling<impl Generator>) {
		self.visit_mut(|object| {
			relabel_object(object, relabeling);
		})
	}

	fn relabel<G: Generator>(&mut self, generator: G) -> JsonLdRelabeling<G> {
		let mut relabeling = JsonLdRelabeling::new(generator);
		self.relabel_with(&mut relabeling);
		relabeling
	}

	fn canonicalize_with(&mut self, buffer: &mut ryu_js::Buffer) {
		self.visit_mut(|object| {
			canonicalize_object(object, buffer);
		});
	}

	fn canonicalize(&mut self) {
		let mut buffer = ryu_js::Buffer::new();
		self.canonicalize_with(&mut buffer);
	}

	fn relabel_and_canonicalize_with(
		&mut self,
		relabeling: &mut JsonLdRelabeling<impl Generator>,
		buffer: &mut ryu_js::Buffer,
	) {
		self.visit_mut(|mut object| {
			relabel_object(object.reborrow(), relabeling);
			canonicalize_object(object, buffer);
		});
	}

	fn relabel_and_canonicalize<G: Generator>(&mut self, generator: G) -> JsonLdRelabeling<G> {
		let mut relabeling = JsonLdRelabeling::new(generator);
		let mut buffer = ryu_js::Buffer::new();
		self.relabel_and_canonicalize_with(&mut relabeling, &mut buffer);
		relabeling
	}
}

fn relabel_object(object: JsonLdObjectMut, relabeling: &mut JsonLdRelabeling<impl Generator>) {
	if let JsonLdObjectMut::Node(node) = object {
		node.id = Some(relabeling.relabel(node.id.take()));

		if let Some(types) = node.types.take() {
			node.types = Some(
				types
					.into_iter()
					.map(|ty| relabeling.relabel(Some(ty)))
					.collect(),
			)
		}
	}
}

fn canonicalize_object(object: JsonLdObjectMut, buffer: &mut ryu_js::Buffer) {
	if let JsonLdObjectMut::Value(value) = object {
		value.canonicalize_with(buffer);
	}
}

impl<T> JsonLdVisit for [T]
where
	T: JsonLdVisit,
{
	fn visit_with(&self, f: &mut impl FnMut(JsonLdObjectRef)) {
		for t in self {
			t.visit_with(f);
		}
	}

	fn visit_mut_with(&mut self, f: &mut impl FnMut(JsonLdObjectMut)) {
		for t in self {
			t.visit_mut_with(f);
		}
	}
}

impl<T> JsonLdVisit for BTreeIndexSet<T>
where
	T: JsonLdVisit + Ord,
{
	fn visit_with(&self, f: &mut impl FnMut(JsonLdObjectRef)) {
		for t in self {
			t.visit_with(f);
		}
	}

	fn visit_mut_with(&mut self, f: &mut impl FnMut(JsonLdObjectMut)) {
		for mut t in std::mem::take(self) {
			t.visit_mut_with(f);
			self.insert(t);
		}
	}
}

impl<T> JsonLdVisit for BTreeIndexMultiSet<T>
where
	T: JsonLdVisit + Ord,
{
	fn visit_with(&self, f: &mut impl FnMut(JsonLdObjectRef)) {
		for t in self {
			t.visit_with(f);
		}
	}

	fn visit_mut_with(&mut self, f: &mut impl FnMut(JsonLdObjectMut)) {
		for mut t in std::mem::take(self) {
			t.visit_mut_with(f);
			self.insert(t);
		}
	}
}

impl<T> JsonLdVisit for Option<T>
where
	T: JsonLdVisit,
{
	fn visit_with(&self, f: &mut impl FnMut(JsonLdObjectRef)) {
		if let Some(t) = self {
			t.visit_with(f);
		}
	}

	fn visit_mut_with(&mut self, f: &mut impl FnMut(JsonLdObjectMut)) {
		if let Some(t) = self {
			t.visit_mut_with(f);
		}
	}
}

pub struct JsonLdRelabeling<G> {
	generator: G,
	map: HashMap<BlankIdBuf, Id>,
}

impl<G> JsonLdRelabeling<G> {
	pub fn new(generator: G) -> Self {
		Self {
			generator,
			map: HashMap::new(),
		}
	}
}

impl<G> JsonLdRelabeling<G>
where
	G: Generator,
{
	pub fn relabel(&mut self, id: Option<JsonLdLenient<Id>>) -> JsonLdLenient<Id> {
		match id {
			Some(JsonLdLenient::Valid(Id::BlankId(b))) => JsonLdLenient::Valid(
				self.map
					.entry(b)
					.or_insert_with(|| self.generator.next_id())
					.clone(),
			),
			Some(id) => id,
			None => JsonLdLenient::Valid(self.generator.next_id()),
		}
	}
}