use crate::{
	nibble::NibbleSlice,
	node::{decode_hash, Node, NodeHandle, NodeHandleOwned, NodeOwned, Value, ValueOwned},
	node_codec::NodeCodec,
	rstd::{boxed::Box, vec::Vec},
	Bytes, CError, CachedValue, DBValue, MerkleValue, Query, RecordedForKey, Result, TrieAccess,
	TrieCache, TrieError, TrieHash, TrieLayout, TrieRecorder,
};
use hash_db::{HashDBRef, Hasher, Prefix};
pub struct Lookup<'a, 'cache, L: TrieLayout, Q: Query<L::Hash>> {
	pub db: &'a dyn HashDBRef<L::Hash, DBValue>,
	pub query: Q,
	pub hash: TrieHash<L>,
	pub cache: Option<&'cache mut dyn TrieCache<L::Codec>>,
	pub recorder: Option<&'cache mut dyn TrieRecorder<TrieHash<L>>>,
}
impl<'a, 'cache, L, Q> Lookup<'a, 'cache, L, Q>
where
	L: TrieLayout,
	Q: Query<L::Hash>,
{
	fn load_value(
		v: Value,
		prefix: Prefix,
		full_key: &[u8],
		db: &dyn HashDBRef<L::Hash, DBValue>,
		recorder: &mut Option<&mut dyn TrieRecorder<TrieHash<L>>>,
		query: Q,
	) -> Result<Q::Item, TrieHash<L>, CError<L>> {
		match v {
			Value::Inline(value) => {
				if let Some(recorder) = recorder {
					recorder.record(TrieAccess::InlineValue { full_key });
				}
				Ok(query.decode(&value))
			},
			Value::Node(hash) => {
				let mut res = TrieHash::<L>::default();
				res.as_mut().copy_from_slice(hash);
				if let Some(value) = db.get(&res, prefix) {
					if let Some(recorder) = recorder {
						recorder.record(TrieAccess::Value {
							hash: res,
							value: value.as_slice().into(),
							full_key,
						});
					}
					Ok(query.decode(&value))
				} else {
					Err(Box::new(TrieError::IncompleteDatabase(res)))
				}
			},
		}
	}
	fn load_owned_value(
		v: ValueOwned<TrieHash<L>>,
		prefix: Prefix,
		full_key: &[u8],
		cache: &mut dyn crate::TrieCache<L::Codec>,
		db: &dyn HashDBRef<L::Hash, DBValue>,
		recorder: &mut Option<&mut dyn TrieRecorder<TrieHash<L>>>,
	) -> Result<(Bytes, TrieHash<L>), TrieHash<L>, CError<L>> {
		match v {
			ValueOwned::Inline(value, hash) => {
				if let Some(recorder) = recorder {
					recorder.record(TrieAccess::InlineValue { full_key });
				}
				Ok((value.clone(), hash))
			},
			ValueOwned::Node(hash) => {
				let node = cache.get_or_insert_node(hash, &mut || {
					let value = db
						.get(&hash, prefix)
						.ok_or_else(|| Box::new(TrieError::IncompleteDatabase(hash)))?;
					Ok(NodeOwned::Value(value.into(), hash))
				})?;
				let value = node
					.data()
					.expect(
						"We are caching a `NodeOwned::Value` for a value node \
						hash and this cached node has always data attached; qed",
					)
					.clone();
				if let Some(recorder) = recorder {
					recorder.record(TrieAccess::Value {
						hash,
						value: value.as_ref().into(),
						full_key,
					});
				}
				Ok((value, hash))
			},
		}
	}
	fn record<'b>(&mut self, get_access: impl FnOnce() -> TrieAccess<'b, TrieHash<L>>)
	where
		TrieHash<L>: 'b,
	{
		if let Some(recorder) = self.recorder.as_mut() {
			recorder.record(get_access());
		}
	}
	pub fn lookup_first_descendant(
		mut self,
		full_key: &[u8],
		nibble_key: NibbleSlice,
	) -> Result<Option<MerkleValue<TrieHash<L>>>, TrieHash<L>, CError<L>> {
		let mut partial = nibble_key;
		let mut hash = self.hash;
		let mut key_nibbles = 0;
		let mut cache = self.cache.take();
		for depth in 0.. {
			let mut _owned_node = NodeOwned::Empty;
			let mut node_data = Vec::new();
			let mut get_owned_node = |depth: i32| {
				let data = match self.db.get(&hash, nibble_key.mid(key_nibbles).left()) {
					Some(value) => value,
					None =>
						return Err(Box::new(match depth {
							0 => TrieError::InvalidStateRoot(hash),
							_ => TrieError::IncompleteDatabase(hash),
						})),
				};
				let decoded = match L::Codec::decode(&data[..]) {
					Ok(node) => node,
					Err(e) => return Err(Box::new(TrieError::DecoderError(hash, e))),
				};
				let owned = decoded.to_owned_node::<L>()?;
				node_data = data;
				Ok(owned)
			};
			let mut node = if let Some(cache) = &mut cache {
				let node = cache.get_or_insert_node(hash, &mut || get_owned_node(depth))?;
				self.record(|| TrieAccess::NodeOwned { hash, node_owned: node });
				node
			} else {
				_owned_node = get_owned_node(depth)?;
				self.record(|| TrieAccess::EncodedNode {
					hash,
					encoded_node: node_data.as_slice().into(),
				});
				&_owned_node
			};
			let mut is_inline = false;
			loop {
				let next_node = match node {
					NodeOwned::Leaf(slice, _) => {
						if !slice.starts_with_slice(&partial) {
							self.record(|| TrieAccess::NonExisting { full_key });
							return Ok(None)
						}
						if partial.len() != slice.len() {
							self.record(|| TrieAccess::NonExisting { full_key });
						}
						let res = is_inline
							.then(|| MerkleValue::Node(node_data))
							.unwrap_or_else(|| MerkleValue::Hash(hash));
						return Ok(Some(res))
					},
					NodeOwned::Extension(slice, item) => {
						if partial.len() < slice.len() {
							self.record(|| TrieAccess::NonExisting { full_key });
							return if slice.starts_with_slice(&partial) {
								let res = is_inline
									.then(|| MerkleValue::Node(node_data))
									.unwrap_or_else(|| MerkleValue::Hash(hash));
								Ok(Some(res))
							} else {
								Ok(None)
							}
						}
						if partial.starts_with_vec(&slice) {
							partial = partial.mid(slice.len());
							key_nibbles += slice.len();
							item
						} else {
							self.record(|| TrieAccess::NonExisting { full_key });
							return Ok(None)
						}
					},
					NodeOwned::Branch(children, value) =>
						if partial.is_empty() {
							if value.is_none() {
								self.record(|| TrieAccess::NonExisting { full_key });
							}
							let res = is_inline
								.then(|| MerkleValue::Node(node_data))
								.unwrap_or_else(|| MerkleValue::Hash(hash));
							return Ok(Some(res))
						} else {
							match &children[partial.at(0) as usize] {
								Some(x) => {
									partial = partial.mid(1);
									key_nibbles += 1;
									x
								},
								None => {
									self.record(|| TrieAccess::NonExisting { full_key });
									return Ok(None)
								},
							}
						},
					NodeOwned::NibbledBranch(slice, children, value) => {
						if partial.len() < slice.len() {
							self.record(|| TrieAccess::NonExisting { full_key });
							return if slice.starts_with_slice(&partial) {
								let res = is_inline
									.then(|| MerkleValue::Node(node_data))
									.unwrap_or_else(|| MerkleValue::Hash(hash));
								Ok(Some(res))
							} else {
								Ok(None)
							}
						}
						if !partial.starts_with_vec(&slice) {
							self.record(|| TrieAccess::NonExisting { full_key });
							return Ok(None)
						}
						if partial.len() == slice.len() {
							if value.is_none() {
								self.record(|| TrieAccess::NonExisting { full_key });
							}
							let res = is_inline
								.then(|| MerkleValue::Node(node_data))
								.unwrap_or_else(|| MerkleValue::Hash(hash));
							return Ok(Some(res))
						} else {
							match &children[partial.at(slice.len()) as usize] {
								Some(x) => {
									partial = partial.mid(slice.len() + 1);
									key_nibbles += slice.len() + 1;
									x
								},
								None => {
									self.record(|| TrieAccess::NonExisting { full_key });
									return Ok(None)
								},
							}
						}
					},
					NodeOwned::Empty => {
						self.record(|| TrieAccess::NonExisting { full_key });
						return Ok(None)
					},
					NodeOwned::Value(_, _) => {
						unreachable!(
							"`NodeOwned::Value` can not be reached by using the hash of a node. \
							 `NodeOwned::Value` is only constructed when loading a value into memory, \
							 which needs to have a different hash than any node; qed",
						)
					},
				};
				match next_node {
					NodeHandleOwned::Hash(new_hash) => {
						hash = *new_hash;
						break
					},
					NodeHandleOwned::Inline(inline_node) => {
						node = &inline_node;
						is_inline = true;
					},
				}
			}
		}
		Ok(None)
	}
	pub fn look_up(
		mut self,
		full_key: &[u8],
		nibble_key: NibbleSlice,
	) -> Result<Option<Q::Item>, TrieHash<L>, CError<L>> {
		match self.cache.take() {
			Some(cache) => self.look_up_with_cache(full_key, nibble_key, cache),
			None => self.look_up_without_cache(nibble_key, full_key, Self::load_value),
		}
	}
	pub fn look_up_hash(
		mut self,
		full_key: &[u8],
		nibble_key: NibbleSlice,
	) -> Result<Option<TrieHash<L>>, TrieHash<L>, CError<L>> {
		match self.cache.take() {
			Some(cache) => self.look_up_hash_with_cache(full_key, nibble_key, cache),
			None => self.look_up_without_cache(
				nibble_key,
				full_key,
				|v, _, full_key, _, recorder, _| {
					Ok(match v {
						Value::Inline(v) => {
							if let Some(recoder) = recorder.as_mut() {
								recoder.record(TrieAccess::InlineValue { full_key });
							}
							L::Hash::hash(&v)
						},
						Value::Node(hash_bytes) => {
							if let Some(recoder) = recorder.as_mut() {
								recoder.record(TrieAccess::Hash { full_key });
							}
							let mut hash = TrieHash::<L>::default();
							hash.as_mut().copy_from_slice(hash_bytes);
							hash
						},
					})
				},
			),
		}
	}
	fn look_up_hash_with_cache(
		mut self,
		full_key: &[u8],
		nibble_key: NibbleSlice,
		cache: &mut dyn crate::TrieCache<L::Codec>,
	) -> Result<Option<TrieHash<L>>, TrieHash<L>, CError<L>> {
		let value_cache_allowed = self
			.recorder
			.as_ref()
			.map(|r| !r.trie_nodes_recorded_for_key(full_key).is_none())
			.unwrap_or(true);
		let res = if let Some(hash) = value_cache_allowed
			.then(|| cache.lookup_value_for_key(full_key).map(|v| v.hash()))
			.flatten()
		{
			hash
		} else {
			let hash_and_value = self.look_up_with_cache_internal(
				nibble_key,
				full_key,
				cache,
				|value, _, full_key, _, _, recorder| match value {
					ValueOwned::Inline(value, hash) => {
						if let Some(recoder) = recorder.as_mut() {
							recoder.record(TrieAccess::InlineValue { full_key });
						}
						Ok((hash, Some(value.clone())))
					},
					ValueOwned::Node(hash) => {
						if let Some(recoder) = recorder.as_mut() {
							recoder.record(TrieAccess::Hash { full_key });
						}
						Ok((hash, None))
					},
				},
			)?;
			match &hash_and_value {
				Some((hash, Some(value))) =>
					cache.cache_value_for_key(full_key, (value.clone(), *hash).into()),
				Some((hash, None)) => cache.cache_value_for_key(full_key, (*hash).into()),
				None => cache.cache_value_for_key(full_key, CachedValue::NonExisting),
			}
			hash_and_value.map(|v| v.0)
		};
		Ok(res)
	}
	fn look_up_with_cache(
		mut self,
		full_key: &[u8],
		nibble_key: NibbleSlice,
		cache: &mut dyn crate::TrieCache<L::Codec>,
	) -> Result<Option<Q::Item>, TrieHash<L>, CError<L>> {
		let trie_nodes_recorded =
			self.recorder.as_ref().map(|r| r.trie_nodes_recorded_for_key(full_key));
		let (value_cache_allowed, value_recording_required) = match trie_nodes_recorded {
			Some(RecordedForKey::Value) | None => (true, false),
			Some(RecordedForKey::Hash) => (true, true),
			Some(RecordedForKey::None) => (false, true),
		};
		let lookup_data = |lookup: &mut Self,
		                   cache: &mut dyn crate::TrieCache<L::Codec>|
		 -> Result<Option<Bytes>, TrieHash<L>, CError<L>> {
			let data = lookup.look_up_with_cache_internal(
				nibble_key,
				full_key,
				cache,
				Self::load_owned_value,
			)?;
			cache.cache_value_for_key(full_key, data.clone().into());
			Ok(data.map(|d| d.0))
		};
		let res = match value_cache_allowed.then(|| cache.lookup_value_for_key(full_key)).flatten()
		{
			Some(CachedValue::NonExisting) => None,
			Some(CachedValue::ExistingHash(hash)) => {
				let data = Self::load_owned_value(
					ValueOwned::Node(*hash),
					nibble_key.original_data_as_prefix(),
					full_key,
					cache,
					self.db,
					&mut self.recorder,
				)?;
				cache.cache_value_for_key(full_key, data.clone().into());
				Some(data.0)
			},
			Some(CachedValue::Existing { data, hash, .. }) =>
				if let Some(data) = data.upgrade() {
					let is_inline =
						L::MAX_INLINE_VALUE.map_or(true, |max| max as usize > data.as_ref().len());
					if value_recording_required && !is_inline {
						self.record(|| TrieAccess::Value {
							hash: *hash,
							value: data.as_ref().into(),
							full_key,
						});
					}
					Some(data)
				} else {
					lookup_data(&mut self, cache)?
				},
			None => lookup_data(&mut self, cache)?,
		};
		Ok(res.map(|v| self.query.decode(&v)))
	}
	fn look_up_with_cache_internal<R>(
		&mut self,
		nibble_key: NibbleSlice,
		full_key: &[u8],
		cache: &mut dyn crate::TrieCache<L::Codec>,
		load_value_owned: impl Fn(
			ValueOwned<TrieHash<L>>,
			Prefix,
			&[u8],
			&mut dyn crate::TrieCache<L::Codec>,
			&dyn HashDBRef<L::Hash, DBValue>,
			&mut Option<&mut dyn TrieRecorder<TrieHash<L>>>,
		) -> Result<R, TrieHash<L>, CError<L>>,
	) -> Result<Option<R>, TrieHash<L>, CError<L>> {
		let mut partial = nibble_key;
		let mut hash = self.hash;
		let mut key_nibbles = 0;
		for depth in 0.. {
			let mut node = cache.get_or_insert_node(hash, &mut || {
				let node_data = match self.db.get(&hash, nibble_key.mid(key_nibbles).left()) {
					Some(value) => value,
					None =>
						return Err(Box::new(match depth {
							0 => TrieError::InvalidStateRoot(hash),
							_ => TrieError::IncompleteDatabase(hash),
						})),
				};
				let decoded = match L::Codec::decode(&node_data[..]) {
					Ok(node) => node,
					Err(e) => return Err(Box::new(TrieError::DecoderError(hash, e))),
				};
				decoded.to_owned_node::<L>()
			})?;
			self.record(|| TrieAccess::NodeOwned { hash, node_owned: node });
			loop {
				let next_node = match node {
					NodeOwned::Leaf(slice, value) =>
						return if partial == *slice {
							let value = (*value).clone();
							load_value_owned(
								value,
								nibble_key.original_data_as_prefix(),
								full_key,
								cache,
								self.db,
								&mut self.recorder,
							)
							.map(Some)
						} else {
							self.record(|| TrieAccess::NonExisting { full_key });
							Ok(None)
						},
					NodeOwned::Extension(slice, item) =>
						if partial.starts_with_vec(&slice) {
							partial = partial.mid(slice.len());
							key_nibbles += slice.len();
							item
						} else {
							self.record(|| TrieAccess::NonExisting { full_key });
							return Ok(None)
						},
					NodeOwned::Branch(children, value) =>
						if partial.is_empty() {
							return if let Some(value) = value.clone() {
								load_value_owned(
									value,
									nibble_key.original_data_as_prefix(),
									full_key,
									cache,
									self.db,
									&mut self.recorder,
								)
								.map(Some)
							} else {
								self.record(|| TrieAccess::NonExisting { full_key });
								Ok(None)
							}
						} else {
							match &children[partial.at(0) as usize] {
								Some(x) => {
									partial = partial.mid(1);
									key_nibbles += 1;
									x
								},
								None => {
									self.record(|| TrieAccess::NonExisting { full_key });
									return Ok(None)
								},
							}
						},
					NodeOwned::NibbledBranch(slice, children, value) => {
						if !partial.starts_with_vec(&slice) {
							self.record(|| TrieAccess::NonExisting { full_key });
							return Ok(None)
						}
						if partial.len() == slice.len() {
							return if let Some(value) = value.clone() {
								load_value_owned(
									value,
									nibble_key.original_data_as_prefix(),
									full_key,
									cache,
									self.db,
									&mut self.recorder,
								)
								.map(Some)
							} else {
								self.record(|| TrieAccess::NonExisting { full_key });
								Ok(None)
							}
						} else {
							match &children[partial.at(slice.len()) as usize] {
								Some(x) => {
									partial = partial.mid(slice.len() + 1);
									key_nibbles += slice.len() + 1;
									x
								},
								None => {
									self.record(|| TrieAccess::NonExisting { full_key });
									return Ok(None)
								},
							}
						}
					},
					NodeOwned::Empty => {
						self.record(|| TrieAccess::NonExisting { full_key });
						return Ok(None)
					},
					NodeOwned::Value(_, _) => {
						unreachable!(
							"`NodeOwned::Value` can not be reached by using the hash of a node. \
							 `NodeOwned::Value` is only constructed when loading a value into memory, \
							 which needs to have a different hash than any node; qed",
						)
					},
				};
				match next_node {
					NodeHandleOwned::Hash(new_hash) => {
						hash = *new_hash;
						break
					},
					NodeHandleOwned::Inline(inline_node) => {
						node = &inline_node;
					},
				}
			}
		}
		Ok(None)
	}
	fn look_up_without_cache<R>(
		mut self,
		nibble_key: NibbleSlice,
		full_key: &[u8],
		load_value: impl Fn(
			Value,
			Prefix,
			&[u8],
			&dyn HashDBRef<L::Hash, DBValue>,
			&mut Option<&mut dyn TrieRecorder<TrieHash<L>>>,
			Q,
		) -> Result<R, TrieHash<L>, CError<L>>,
	) -> Result<Option<R>, TrieHash<L>, CError<L>> {
		let mut partial = nibble_key;
		let mut hash = self.hash;
		let mut key_nibbles = 0;
		for depth in 0.. {
			let node_data = match self.db.get(&hash, nibble_key.mid(key_nibbles).left()) {
				Some(value) => value,
				None =>
					return Err(Box::new(match depth {
						0 => TrieError::InvalidStateRoot(hash),
						_ => TrieError::IncompleteDatabase(hash),
					})),
			};
			self.record(|| TrieAccess::EncodedNode {
				hash,
				encoded_node: node_data.as_slice().into(),
			});
			let mut node_data = &node_data[..];
			loop {
				let decoded = match L::Codec::decode(node_data) {
					Ok(node) => node,
					Err(e) => return Err(Box::new(TrieError::DecoderError(hash, e))),
				};
				let next_node = match decoded {
					Node::Leaf(slice, value) =>
						return if slice == partial {
							load_value(
								value,
								nibble_key.original_data_as_prefix(),
								full_key,
								self.db,
								&mut self.recorder,
								self.query,
							)
							.map(Some)
						} else {
							self.record(|| TrieAccess::NonExisting { full_key });
							Ok(None)
						},
					Node::Extension(slice, item) =>
						if partial.starts_with(&slice) {
							partial = partial.mid(slice.len());
							key_nibbles += slice.len();
							item
						} else {
							self.record(|| TrieAccess::NonExisting { full_key });
							return Ok(None)
						},
					Node::Branch(children, value) =>
						if partial.is_empty() {
							return if let Some(val) = value {
								load_value(
									val,
									nibble_key.original_data_as_prefix(),
									full_key,
									self.db,
									&mut self.recorder,
									self.query,
								)
								.map(Some)
							} else {
								self.record(|| TrieAccess::NonExisting { full_key });
								Ok(None)
							}
						} else {
							match children[partial.at(0) as usize] {
								Some(x) => {
									partial = partial.mid(1);
									key_nibbles += 1;
									x
								},
								None => {
									self.record(|| TrieAccess::NonExisting { full_key });
									return Ok(None)
								},
							}
						},
					Node::NibbledBranch(slice, children, value) => {
						if !partial.starts_with(&slice) {
							self.record(|| TrieAccess::NonExisting { full_key });
							return Ok(None)
						}
						if partial.len() == slice.len() {
							return if let Some(val) = value {
								load_value(
									val,
									nibble_key.original_data_as_prefix(),
									full_key,
									self.db,
									&mut self.recorder,
									self.query,
								)
								.map(Some)
							} else {
								self.record(|| TrieAccess::NonExisting { full_key });
								Ok(None)
							}
						} else {
							match children[partial.at(slice.len()) as usize] {
								Some(x) => {
									partial = partial.mid(slice.len() + 1);
									key_nibbles += slice.len() + 1;
									x
								},
								None => {
									self.record(|| TrieAccess::NonExisting { full_key });
									return Ok(None)
								},
							}
						}
					},
					Node::Empty => {
						self.record(|| TrieAccess::NonExisting { full_key });
						return Ok(None)
					},
				};
				match next_node {
					NodeHandle::Hash(data) => {
						hash = decode_hash::<L::Hash>(data)
							.ok_or_else(|| Box::new(TrieError::InvalidHash(hash, data.to_vec())))?;
						break
					},
					NodeHandle::Inline(data) => {
						node_data = data;
					},
				}
			}
		}
		Ok(None)
	}
}