use crate::{fetch_page, PackageError};
use codec::DecodeAll;
use corevm_host::{fs as corevm_fs, fs::ReadBlock as _, ExecEnv, PageNum, RangeSet, StorageKey};
use futures::StreamExt;
use jam_std_common::{Node, NodeError, NodeExt as _, WorkPackageStatus};
use jam_types::{HeaderHash, SegmentBytes, ServiceId, Slot, WorkPackageHash};
use std::{collections::HashMap, sync::Arc};
use tokio::task::spawn_blocking;

pub trait NodeExt: Node {
	async fn wait_for_package(
		&self,
		package_hash: WorkPackageHash,
		anchor: HeaderHash,
	) -> Result<(), PackageError> {
		let mut sub = self.subscribe_work_package_status(package_hash, anchor, false).await?;
		loop {
			match sub.next().await.ok_or(NodeError::NetworkFailure)??.value {
				WorkPackageStatus::Reportable { .. } => (), // Wait
				WorkPackageStatus::Reported { .. } | WorkPackageStatus::Ready { .. } =>
					break Ok(()),
				WorkPackageStatus::Failed(message) => break Err(PackageError::Failed(message)),
			}
		}
	}
}

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

pub trait ArcNodeExt {
	/// Fetch resident memory pages of a service from the DA.
	async fn fetch_pages(
		&self,
		head: HeaderHash,
		service_id: ServiceId,
	) -> Result<HashMap<PageNum, SegmentBytes>, NodeError>;

	/// Read file from CoreVM FS.
	async fn read_file(
		&self,
		block_ref: corevm_fs::BlockRef,
		provision_timeout: Option<Slot>,
	) -> Result<Vec<u8>, corevm_fs::Error>;

	/// Caches CoreVM FS blocks and optionally waits until they're provided.
	async fn prefetch_file_blocks_recursively(
		&self,
		root: corevm_fs::BlockRef,
		provision_timeout: Option<Slot>,
	) -> Result<(), corevm_fs::Error>;

	async fn fetch_exec_env(
		&self,
		block_ref: corevm_fs::BlockRef,
		provision_timeout: Option<Slot>,
	) -> Result<ExecEnv, NodeError>;
}

impl ArcNodeExt for Arc<dyn Node> {
	async fn fetch_pages(
		&self,
		head: HeaderHash,
		service_id: ServiceId,
	) -> Result<HashMap<PageNum, SegmentBytes>, NodeError> {
		let resident_pages = self
			.typed_service_value::<RangeSet>(head, service_id, &StorageKey::StoredPages)
			.await?
			.unwrap_or_default();
		let mut all_pages = HashMap::new();
		for range in resident_pages.as_slice() {
			for page in range.start..range.end {
				let page = PageNum(page);
				let info = self
					.typed_service_value::<corevm_host::PageInfo>(
						head,
						service_id,
						&StorageKey::PageInfo(page),
					)
					.await?
					.ok_or_else(|| {
						NodeError::Other(format!("Memory page {page} is missing in the storage"))
					})?;
				let segment =
					fetch_page(self.as_ref(), info.exports_root, info.export_index, info.hash)
						.await?;
				all_pages.insert(page, segment);
				log::trace!("Fetched page {page}");
			}
		}
		Ok(all_pages)
	}

	async fn read_file(
		&self,
		block_ref: corevm_fs::BlockRef,
		provision_timeout: Option<Slot>,
	) -> Result<Vec<u8>, corevm_fs::Error> {
		let mut reader =
			corevm_tooling::PreimageReader::with_provision(self.clone(), provision_timeout);
		spawn_blocking(move || corevm_fs::read(&block_ref, &mut reader))
			.await
			.map_err(|_| corevm_fs::Error::Io)?
	}

	async fn prefetch_file_blocks_recursively(
		&self,
		root: corevm_fs::BlockRef,
		provision_timeout: Option<Slot>,
	) -> Result<(), corevm_fs::Error> {
		let mut reader =
			corevm_tooling::PreimageReader::with_provision(self.clone(), provision_timeout);
		spawn_blocking(move || {
			let nodes = corevm_fs::NodeIter::new(root, &mut reader);
			let mut blocks = Vec::new();
			for entry in nodes {
				let (_hash, file) = entry?;
				let main_block = file.main_block();
				if main_block.kind() != corevm_fs::NodeKind::File {
					continue;
				}
				blocks.extend_from_slice(main_block.block_refs());
			}
			for block in blocks.iter() {
				reader.read_block(block)?;
			}
			Ok(())
		})
		.await
		.map_err(|_| corevm_fs::Error::Io)?
	}

	async fn fetch_exec_env(
		&self,
		block_ref: corevm_fs::BlockRef,
		provision_timeout: Option<Slot>,
	) -> Result<ExecEnv, NodeError> {
		let bytes = self
			.read_file(block_ref, provision_timeout)
			.await
			.map_err(|e| NodeError::Other(format!("CoreVM FS i/o error: {e}")))?;
		let exec = ExecEnv::decode_all(&mut &bytes[..])
			.map_err(|e| NodeError::Other(format!("Failed to decode ExecEnv: {e}")))?;
		Ok(exec)
	}
}
