use bytes::Bytes;
use futures::stream::BoxStream;
use jam_std_common::{
	BlockDesc, ChainSub, Node, NodeResult, SyncState, SyncStatus, VersionedParameters,
	WorkPackageStatus, WorkReport,
};
use jam_types::{
	CoreIndex, Hash, HeaderHash, MmrPeakHash, SegmentBytes, SegmentTreeRoot, ServiceId, Slot,
	StateRootHash, WorkPackageHash, WorkReportHash, MAX_PREIMAGE_BLOB_LEN,
};
use parking_lot::Mutex;
use schnellru::{ByLength, LruMap};
use std::sync::Arc;

/// A [`Node`] implementation that caches preimages.
///
/// Non-preimage calls are simply forwarded to the underlying [`Node`] implementation.
pub struct NodeWithCache<N: Node> {
	node: N,
	cached_preimages: Arc<Mutex<PreimageCache>>,
}

impl<N: Node> NodeWithCache<N> {
	pub fn new(node: N, cached_preimages: Arc<Mutex<PreimageCache>>) -> Self {
		Self { node, cached_preimages }
	}

	pub fn get(&self) -> &N {
		&self.node
	}
}

#[async_trait::async_trait]
impl<N: Node> Node for NodeWithCache<N> {
	async fn parameters(&self) -> NodeResult<VersionedParameters> {
		self.node.parameters().await
	}

	async fn best_block(&self) -> NodeResult<BlockDesc> {
		self.node.best_block().await
	}

	async fn subscribe_best_block(&self) -> NodeResult<BoxStream<NodeResult<BlockDesc>>> {
		self.node.subscribe_best_block().await
	}

	async fn finalized_block(&self) -> NodeResult<BlockDesc> {
		self.node.finalized_block().await
	}

	async fn subscribe_finalized_block(&self) -> NodeResult<BoxStream<NodeResult<BlockDesc>>> {
		self.node.subscribe_finalized_block().await
	}

	async fn parent(&self, header_hash: HeaderHash) -> NodeResult<BlockDesc> {
		self.node.parent(header_hash).await
	}

	async fn state_root(&self, header_hash: HeaderHash) -> NodeResult<StateRootHash> {
		self.node.state_root(header_hash).await
	}

	async fn beefy_root(&self, header_hash: HeaderHash) -> NodeResult<MmrPeakHash> {
		self.node.beefy_root(header_hash).await
	}

	async fn encoded_statistics(&self, header_hash: HeaderHash) -> NodeResult<Bytes> {
		self.node.encoded_statistics(header_hash).await
	}

	async fn subscribe_encoded_statistics(&self, finalized: bool) -> NodeResult<ChainSub<Bytes>> {
		self.node.subscribe_encoded_statistics(finalized).await
	}

	async fn encoded_service_data(
		&self,
		header_hash: HeaderHash,
		id: ServiceId,
	) -> NodeResult<Option<Bytes>> {
		self.node.encoded_service_data(header_hash, id).await
	}

	async fn subscribe_encoded_service_data(
		&self,
		id: ServiceId,
		finalized: bool,
	) -> NodeResult<ChainSub<Option<Bytes>>> {
		self.node.subscribe_encoded_service_data(id, finalized).await
	}

	async fn service_value(
		&self,
		header_hash: HeaderHash,
		id: ServiceId,
		key: &[u8],
	) -> NodeResult<Option<Bytes>> {
		self.node.service_value(header_hash, id, key).await
	}

	async fn subscribe_service_value(
		&self,
		id: ServiceId,
		key: &[u8],
		finalized: bool,
	) -> NodeResult<ChainSub<Option<Bytes>>> {
		self.node.subscribe_service_value(id, key, finalized).await
	}

	async fn service_preimage(
		&self,
		header_hash: HeaderHash,
		id: ServiceId,
		hash: Hash,
	) -> NodeResult<Option<Bytes>> {
		{
			let mut cache = self.cached_preimages.lock();
			if let Some(preimage) = cache.load_preimage(&hash) {
				return Ok(Some(preimage));
			};
		}
		let preimage: Option<Bytes> = self.node.service_preimage(header_hash, id, hash).await?;
		if let Some(ref preimage) = preimage {
			self.cached_preimages.lock().store_preimage(hash, preimage.clone());
		}
		Ok(preimage)
	}

	async fn subscribe_service_preimage(
		&self,
		id: ServiceId,
		hash: Hash,
		finalized: bool,
	) -> NodeResult<ChainSub<Option<Bytes>>> {
		self.node.subscribe_service_preimage(id, hash, finalized).await
	}

	async fn service_preimage_len(
		&self,
		header_hash: HeaderHash,
		id: ServiceId,
		hash: Hash,
	) -> NodeResult<Option<u32>> {
		{
			let mut cache = self.cached_preimages.lock();
			if let Some(preimage) = cache.load_preimage(&hash) {
				return Ok(Some(preimage.len() as u32));
			};
		}
		self.node.service_preimage_len(header_hash, id, hash).await
	}

	async fn service_request(
		&self,
		header_hash: HeaderHash,
		id: ServiceId,
		hash: Hash,
		len: u32,
	) -> NodeResult<Option<Vec<Slot>>> {
		self.node.service_request(header_hash, id, hash, len).await
	}

	async fn subscribe_service_request(
		&self,
		id: ServiceId,
		hash: Hash,
		len: u32,
		finalized: bool,
	) -> NodeResult<ChainSub<Option<Vec<Slot>>>> {
		self.node.subscribe_service_request(id, hash, len, finalized).await
	}

	async fn work_report(&self, hash: WorkReportHash) -> NodeResult<WorkReport> {
		self.node.work_report(hash).await
	}

	async fn submit_encoded_work_package(
		&self,
		core: CoreIndex,
		package: Bytes,
		extrinsics: &[Bytes],
	) -> NodeResult<()> {
		self.node.submit_encoded_work_package(core, package, extrinsics).await
	}

	async fn submit_encoded_work_package_bundle(
		&self,
		core: CoreIndex,
		bundle: Bytes,
	) -> NodeResult<()> {
		self.node.submit_encoded_work_package_bundle(core, bundle).await
	}

	async fn work_package_status(
		&self,
		header_hash: HeaderHash,
		hash: WorkPackageHash,
		anchor: HeaderHash,
	) -> NodeResult<WorkPackageStatus> {
		self.node.work_package_status(header_hash, hash, anchor).await
	}

	async fn subscribe_work_package_status(
		&self,
		hash: WorkPackageHash,
		anchor: HeaderHash,
		finalized: bool,
	) -> NodeResult<ChainSub<WorkPackageStatus>> {
		self.node.subscribe_work_package_status(hash, anchor, finalized).await
	}

	async fn submit_preimage(&self, requester: ServiceId, preimage: Bytes) -> NodeResult<()> {
		self.node.submit_preimage(requester, preimage).await
	}

	async fn list_services(&self, header_hash: HeaderHash) -> NodeResult<Vec<ServiceId>> {
		self.node.list_services(header_hash).await
	}

	async fn fetch_work_package_segments(
		&self,
		wp_hash: WorkPackageHash,
		indices: Vec<u16>,
	) -> NodeResult<Vec<SegmentBytes>> {
		self.node.fetch_work_package_segments(wp_hash, indices).await
	}

	async fn fetch_segments(
		&self,
		segment_root: SegmentTreeRoot,
		indices: Vec<u16>,
	) -> NodeResult<Vec<SegmentBytes>> {
		self.node.fetch_segments(segment_root, indices).await
	}

	async fn sync_state(&self) -> NodeResult<SyncState> {
		self.node.sync_state().await
	}

	async fn subscribe_sync_status(&self) -> NodeResult<BoxStream<NodeResult<SyncStatus>>> {
		self.node.subscribe_sync_status().await
	}
}

/// Preimage cache backed by an LRU map.
pub struct PreimageCache {
	map: LruMap<Hash, Bytes>,
}

impl PreimageCache {
	/// Create a cache for preimages that occupy at most `total_size_in_bytes` bytes in total.
	pub fn new(total_size_in_bytes: u64) -> Self {
		let max_preimages = total_size_in_bytes / MAX_PREIMAGE_BLOB_LEN as u64;
		let max_preimages = max_preimages.try_into().unwrap_or(u32::MAX);
		Self { map: LruMap::new(ByLength::new(max_preimages)) }
	}

	pub fn load_preimage(&mut self, hash: &Hash) -> Option<Bytes> {
		self.map.get(hash).cloned()
	}

	pub fn store_preimage(&mut self, hash: Hash, preimage: Bytes) {
		self.map.insert(hash, preimage);
	}
}
