use crate::{
	proto, proto::InputChunkIndex, BoundedInputStreams, BuilderTask, ConnectionMan, CoreScheduler,
	ExecCache,
};
use corevm_host::{fs as corevm_fs, Outcome, StorageKey, VmSpec};
use futures::{stream::FuturesUnordered, Stream, StreamExt};
use jam_std_common::{BlockDesc, Node, NodeError, NodeExt as _};
use jam_tooling::NodeExt as _;
use jam_types::{HeaderHash, ServiceId, Slot, VecSet};
use parking_lot::Mutex;
use std::{
	collections::{hash_map::Entry, HashMap, HashSet},
	num::NonZero,
	sync::{
		atomic::{AtomicBool, Ordering},
		Arc,
	},
};
use tokio::{
	sync::{mpsc, watch},
	task::{AbortHandle, JoinHandle},
};

/// Means of communication with a particular builder task.
struct BuilderComms {
	/// Aborting builder task doesn't stop the simulator.
	builder_task: AbortHandle,
	/// Setting this to `true` stops the simulator as soon as it calls `invoke`.
	simulator_stopped: Arc<AtomicBool>,
	input_streams: Arc<Mutex<BoundedInputStreams>>,
	/// Index of the latest input chunk that we received from the monitor.
	input_chunk_index: watch::Sender<InputChunkIndex>,
}

impl BuilderComms {
	fn abort(&self) {
		self.builder_task.abort();
		self.simulator_stopped.store(true, Ordering::Relaxed);
	}
}

/// A command that can be issued to a general builder.
pub enum BuilderCommand {
	/// Send input chunk to the service builder.
	SendInput(proto::InputChunkMessage),
}

/// General builder configuration.
#[derive(Debug)]
#[non_exhaustive]
pub struct BuilderConfig {
	/// How many CoreVM services can be built in parallel.
	pub max_builders: usize,
	/// Max. no. of guests in cache.
	///
	/// From statistical point of view should be equal to a typical max. number of _active_ CoreVM
	/// services (i.e. the number of services being built) multiplied by a typical max. number of
	/// guests a service executes during its lifetime.
	///
	/// Cache entry size is 36 bytes.
	pub exec_cache_size: u32,
	/// Preimage provision timeout.
	pub provision_timeout: Slot,
}

impl Default for BuilderConfig {
	fn default() -> Self {
		// If you change defaults here you should consider changing them in
		// `corevm-builder/src/main.rs` as well.
		Self {
			max_builders: std::thread::available_parallelism()
				.unwrap_or(NonZero::new(1).expect("Constant is non-zero"))
				.get(),
			exec_cache_size: 4 * 64 * 1024,
			provision_timeout: 10,
		}
	}
}

/// Manages invdividual CoreVM builders.
pub struct GeneralBuilder {
	connections: ConnectionMan,
	/// A subset of currently active JAM services.
	///
	/// Might contain non-CoreVM services.
	services: HashSet<ServiceId>,
	engine: Arc<polkavm::Engine>,
	exec_cache: Arc<ExecCache>,
	core_scheduler: Arc<Mutex<CoreScheduler>>,
	/// Current slot as reported by JAM node.
	current_slot: Option<Slot>,
	builder_comms: HashMap<ServiceId, BuilderComms>,
	builder_tasks: FuturesUnordered<JoinHandle<ServiceId>>,
	max_builder_tasks: usize,
	provision_timeout: Slot,
}

impl GeneralBuilder {
	/// Creates new builder from the provided configuration and JAM node URLs.
	pub fn new(
		connections: ConnectionMan,
		core_scheduler: Arc<Mutex<CoreScheduler>>,
		engine: Arc<polkavm::Engine>,
		config: BuilderConfig,
	) -> Self {
		Self {
			connections,
			services: HashSet::new(),
			engine,
			exec_cache: Arc::new(ExecCache::new(config.exec_cache_size)),
			core_scheduler,
			current_slot: None,
			builder_comms: HashMap::new(),
			builder_tasks: FuturesUnordered::new(),
			max_builder_tasks: config.max_builders,
			provision_timeout: config.provision_timeout,
		}
	}

	pub async fn run(mut self, mut commands: mpsc::Receiver<BuilderCommand>) {
		let node = loop {
			let node = self.connections.get().await;
			match node.wait_for_sync().await {
				Ok(()) => break node,
				Err(e) => log::debug!("Failed to syncrhonize with JAM node: {e}; retrying..."),
			}
		};
		// Initial channel value can be anything because it's not considered "seen" by new
		// receivers.
		let (mut best_block_tx, _) =
			watch::channel((BlockDesc { header_hash: Default::default(), slot: 0 }, node));
		// This loop automatically reconnects to JAM node on network failures.
		loop {
			let node = self.connections.get().await;
			if let Err(e) = self.do_run(node.clone(), &mut commands, &mut best_block_tx).await {
				log::debug!("General builder loop failed: {e}");
			}
		}
	}

	async fn do_run(
		&mut self,
		node: Arc<dyn Node>,
		commands: &mut mpsc::Receiver<BuilderCommand>,
		best_block_tx: &mut watch::Sender<(BlockDesc, Arc<dyn Node>)>,
	) -> Result<(), NodeError> {
		let mut best_block_sub = node.subscribe_best_block().await?;
		while let Some(block) = self.wait_until_next_block(&mut best_block_sub, commands).await? {
			let BlockDesc { header_hash, slot } = block;
			if let Some(current_slot) = self.current_slot {
				if current_slot >= slot {
					// This is a safeguard against node asynchrony: when network failure occurs, the
					// builder reconnects to the new node, but the new node might have
					// previous/next slot as the current one.
					continue;
				}
			}
			self.current_slot = Some(slot);
			let service_ids = node.list_services(header_hash).await?;
			if let Err(e) = self
				.update_services(header_hash, node.clone(), service_ids.into_iter(), best_block_tx)
				.await
			{
				log::debug!("Failed to update service list {header_hash}: {e}");
				if matches!(e, NodeError::NetworkFailure) {
					// Force reconnect.
					return Err(e);
				}
				// Try next block.
				continue;
			}
			while let Ok(command) = commands.try_recv() {
				if let Err(e) = self.process_command(command) {
					log::debug!("Failed to process command: {e}");
				}
			}
			best_block_tx.send_replace((BlockDesc { header_hash, slot }, node.clone()));
		}
		Err(NodeError::Other("Statistics subscription failed".into()))
	}

	/// Waits until the next best block and awaits builder tasks in parallel.
	async fn wait_until_next_block(
		&mut self,
		best_block_sub: &mut (dyn Stream<Item = Result<BlockDesc, NodeError>> + Unpin + Send),
		commands: &mut mpsc::Receiver<BuilderCommand>,
	) -> Result<Option<BlockDesc>, NodeError> {
		loop {
			tokio::select! {
				update = best_block_sub.next() => return update.transpose(),
				result = self.builder_tasks.next(), if !self.builder_tasks.is_empty() => {
					match result.expect("Checked that tasks are not empty") {
						Err(e) => log::debug!("Failed to join task: {e}"),
						Ok(service_id) => {
							log::trace!(service_id; "Builder for service {service_id:x} terminated");
							self.builder_comms.remove(&service_id);
						}
					}
				}
				Some(command) = commands.recv() => {
					if let Err(e) = self.process_command(command) {
						log::debug!("Failed to process command: {e}");
					}
				}
			}
		}
	}

	/// Update the list of CoreVM services from the specified block.
	async fn update_services(
		&mut self,
		head: HeaderHash,
		node: Arc<dyn Node>,
		service_ids: impl Iterator<Item = ServiceId>,
		best_block_tx: &mut watch::Sender<(BlockDesc, Arc<dyn Node>)>,
	) -> Result<(), NodeError> {
		let service_ids: VecSet<ServiceId> = service_ids.collect();
		self.services.retain(|service_id| {
			let retain = service_ids.contains(service_id);
			if !retain {
				log::trace!(service_id; "Removing service {service_id:x}");
				if let Some(comms) = self.builder_comms.remove(service_id) {
					comms.abort();
				}
			}
			retain
		});
		// TODO @ivan It might be beneficial to do this in parallel or in batches.
		for service_id in service_ids.into_iter() {
			if let Err(e) = self.update_service(node.clone(), head, service_id, best_block_tx).await
			{
				log::debug!(service_id; "Failed to process service {service_id:x}: {e}");
			}
		}
		Ok(())
	}

	/// Update single CoreVM service information from the specified block.
	async fn update_service(
		&mut self,
		node: Arc<dyn Node>,
		head: HeaderHash,
		service_id: ServiceId,
		best_block_tx: &mut watch::Sender<(BlockDesc, Arc<dyn Node>)>,
	) -> Result<(), NodeError> {
		if self.services.contains(&service_id) {
			return Ok(());
		}
		let service = node
			.service_data(head, service_id)
			.await?
			.ok_or_else(|| NodeError::Other("Service not found".into()))?;
		if service.code_hash.0 != corevm_bin::HASH {
			return Ok(());
		}
		if let Some(vm_spec) = node
			.typed_service_value::<VmSpec>(head, service_id, &StorageKey::VmSpec)
			.await?
		{
			match vm_spec.output.outcome {
				Outcome::Halt => {
					log::info!(service_id; "Can't build service that halted");
					return Ok(());
				},
				Outcome::Panic => {
					log::info!(service_id; "Can't build service that panicked");
					return Ok(());
				},
				Outcome::PageFault { .. } |
				Outcome::OutOfGas |
				Outcome::OutputLimitReached |
				Outcome::TimeLimitReached |
				Outcome::InputLimitReached |
				Outcome::WaitingForInput(..) => {},
			}
		}
		if let Some(exec_env_ref_hash) = self.exec_cache.get_timed_out(service_id) {
			let exec_env_ref = node
				.typed_service_value::<corevm_fs::BlockRef>(
					head,
					service_id,
					&StorageKey::ExecEnvRef,
				)
				.await?
				.ok_or_else(|| NodeError::Other("No ExecEnvRef".into()))?;
			if exec_env_ref.hash.0 == exec_env_ref_hash {
				log::info!(service_id; "Skipping CoreVM service {service_id:x}: Guest timed out");
				return Ok(());
			}
			self.exec_cache.remove(service_id);
		}
		if self.builder_tasks.len() == self.max_builder_tasks {
			log::warn!(
				service_id;
				"Failed to create builder for service {service_id:x}: reached configured limit ({})",
				self.max_builder_tasks
			);
			return Ok(());
		}
		self.services.insert(service_id);
		log::trace!(service_id; "Building service {service_id:x}");
		self.spawn_builder_task(service_id, best_block_tx);
		Ok(())
	}

	fn spawn_builder_task(
		&mut self,
		service_id: ServiceId,
		best_block_tx: &mut watch::Sender<(BlockDesc, Arc<dyn Node>)>,
	) {
		let Entry::Vacant(v) = self.builder_comms.entry(service_id) else {
			return;
		};
		let input_streams = Arc::new(Mutex::new(BoundedInputStreams::new()));
		let simulator_stopped = Arc::new(AtomicBool::new(false));
		let (input_chunk_index_tx, input_chunk_index_rx) = watch::channel(InputChunkIndex(0));
		let builder = BuilderTask::new(
			service_id,
			self.provision_timeout,
			self.engine.clone(),
			self.exec_cache.clone(),
			self.core_scheduler.clone(),
			input_streams.clone(),
			simulator_stopped.clone(),
			input_chunk_index_rx,
		);
		let best_block_rx = best_block_tx.subscribe();
		let task = tokio::spawn(async move {
			builder.run(best_block_rx).await;
			service_id
		});
		v.insert(BuilderComms {
			builder_task: task.abort_handle(),
			simulator_stopped,
			input_streams,
			input_chunk_index: input_chunk_index_tx,
		});
		self.builder_tasks.push(task);
	}

	fn process_command(&mut self, command: BuilderCommand) -> Result<(), CommandError> {
		let BuilderCommand::SendInput(chunk) = command;
		let comms = self
			.builder_comms
			.get(&chunk.service_id)
			.ok_or(CommandError::InvalidServiceId(chunk.service_id))?;
		let last_index = *comms.input_chunk_index.borrow();
		let Some(current_slot) = self.current_slot else {
			log::trace!(service_id = chunk.service_id; "Dropping input chunk: current chain slot is unknown");
			return Ok(());
		};
		let min_index = InputChunkIndex::first_in_slot(current_slot.saturating_sub(1));
		let max_index = InputChunkIndex::last_in_slot(current_slot.saturating_add(1));
		let accepted_indices = min_index.max(last_index)..=max_index;
		if !accepted_indices.contains(&chunk.index) {
			log::trace!(
				service_id = chunk.service_id;
				"Dropping input chunk: invalid index ({}); accepted indices {}-{}",
				chunk.index,
				accepted_indices.start(),
				accepted_indices.end(),
			);
			return Ok(());
		}
		if let Err(e) = comms.input_streams.lock().append_chunk(chunk.index, chunk.body) {
			log::info!(service_id = chunk.service_id; "Failed to append input chunk: {e}");
		}
		comms.input_chunk_index.send_replace(chunk.index);
		Ok(())
	}
}

#[derive(thiserror::Error, Debug)]
enum CommandError {
	#[error("Invalid service id: {0:x}")]
	InvalidServiceId(ServiceId),
}
