use crate::{CodeInfo, Error};
use codec::{Compact, Decode};
use futures::StreamExt;
use jam_program_blob_common::{ConventionalMetadata as Metadata, CrateInfo};
use jam_std_common::{BlockDesc, Node, NodeError, NodeExt as _, Service, SyncStatus};
use jam_types::{HeaderHash, RefineContext, ServiceId, Slot};
use std::future::Future;

pub trait NodeExt: Node {
	fn query_service(
		&self,
		head: HeaderHash,
		id: ServiceId,
	) -> impl Future<Output = Result<(Service, CodeInfo<CrateInfo>), anyhow::Error>> + Send {
		async move {
			use CodeInfo::*;
			let info = self.service_data(head, id).await?.ok_or(Error::ServiceNotFound)?;
			let maybe_code = self.service_preimage(head, id, info.code_hash.0).await?;
			let meta = match maybe_code {
				None => CodeNotProvided(info.code_hash.0),
				Some(code) =>
					match <(Compact<u32>, Metadata)>::decode(&mut &code[..]).ok().map(|x| x.1) {
						None => Undefined(info.code_hash.0),
						Some(Metadata::Info(crate_info)) => Known(crate_info),
					},
			};
			Ok((info, meta))
		}
	}

	fn wait_for_sync(&self) -> impl Future<Output = Result<(), anyhow::Error>> + Send {
		async move {
			let mut sync_state = self.subscribe_sync_status().await?;
			loop {
				let status = sync_state.next().await;
				match status {
					Some(Ok(SyncStatus::Completed)) => break,
					Some(Ok(_)) => continue,
					Some(Err(e)) => return Err(e.into()),
					None => break,
				}
			}
			// Wait for any non-genesis block to be finalized
			let mut final_block = self.subscribe_finalized_block().await?;
			loop {
				let status = final_block.next().await;
				match status {
					Some(Ok(block)) if block.slot > 0 => {
						match self.parent(block.header_hash).await {
							Ok(parent) => {
								if parent.slot > 0 {
									// Break only if the parent of the finalized block is
									// non-genesis.
									break;
								}
								eprintln!("Waiting for two non-genesis blocks to be finalized...");
							},
							Err(NodeError::BlockUnavailable(_)) => {
								// Just wait for another block
								eprintln!("Waiting for two non-genesis blocks to be finalized and available...");
							},
							Err(err) => return Err(err.into()),
						}
					},
					Some(Ok(_)) => {
						eprintln!("Waiting for two non-genesis blocks to be finalized...");
					},
					Some(Err(e)) => return Err(e.into()),
					None => break,
				}
			}
			Ok(())
		}
	}

	/// Returns lookup anchor block description.
	///
	/// This is the parent of the finalized block.
	fn lookup_anchor_block(&self) -> impl Future<Output = Result<BlockDesc, NodeError>> + Send {
		async move {
			let finalized = self.finalized_block().await?;
			let parent = self.parent(finalized.header_hash).await?;
			Ok(parent)
		}
	}

	/// Create refine context from the current chain's best block.
	///
	/// - Best block's parent is used as an anchor.
	/// - Finalized block is used as a lookup anchor.
	///
	/// Returns the resulting refine context and anchor slot.
	fn create_refine_context(
		&self,
	) -> impl Future<Output = Result<(RefineContext, Slot), NodeError>> + Send {
		async move {
			// Use the parent of the latest finalized block as the lookup anchor. The parent is used
			// as some nodes might have a bit of finality lag.
			let BlockDesc { header_hash: lookup_anchor, slot: lookup_anchor_slot } =
				self.lookup_anchor_block().await?;
			// Use the parent of the best block as the anchor. The parent is used as some nodes
			// might not have seen the best block yet.
			let best_block = self.best_block().await?;
			let parent = self.parent(best_block.header_hash).await?;
			let anchor = parent.header_hash;
			let state_root = self.state_root(anchor).await?;
			let beefy_root = self.beefy_root(anchor).await?;
			let context = RefineContext {
				anchor,
				state_root,
				beefy_root,
				lookup_anchor,
				lookup_anchor_slot,
				prerequisites: Default::default(),
			};
			Ok((context, parent.slot))
		}
	}
}

impl<T: Node + ?Sized> NodeExt for T {}
