use crate::{
	proto, proto::InputChunkIndex, ArcNodeExt as _, BoundedInputStreams, Builder, CoreScheduler,
	ExecCache, NodeExt as _, PackageError, ServiceMessages, StepOutcome, Timer,
};
use bytes::Bytes;
use codec::Encode;
use corevm_host::{
	fs as corevm_fs, CoreVmExtrinsics, CoreVmPayload, InputStreams, Outcome, PageNum, StorageKey,
	VmSpec, VmState,
};
use corevm_tooling::NodeExt;
use jam_std_common::{BlockDesc, Node, NodeError, NodeExt as _};
use jam_tooling::{format::concise_duration, NodeExt as _};
use jam_types::{slot_period_sec, CodeHash, HeaderHash, ServiceId, SignedGas, Slot};
use parking_lot::Mutex;
use std::{
	sync::{atomic::AtomicBool, Arc},
	time::{Duration, Instant},
};
use tokio::{sync::watch, time::timeout};

pub struct BuilderTask {
	pub(crate) service_id: ServiceId,
	provision_timeout: Slot,
	engine: Arc<polkavm::Engine>,
	exec_cache: Arc<ExecCache>,
	exec_env_hash: Option<corevm_fs::Hash>,
	last_processed_slot: Option<Slot>,
	core_scheduler: Arc<Mutex<CoreScheduler>>,
	stopped: Arc<AtomicBool>,
	/// The state of inter-service message queue.
	pub(crate) service_messages: ServiceMessages,
	/// The state of input streams.
	pub(crate) input_streams: Arc<Mutex<BoundedInputStreams>>,
	waiting_for_input: Option<InputStreams>,
	input_chunk_index: watch::Receiver<InputChunkIndex>,
	/// Wait until the last chunk of video/audio from the monitor arrives before running the guest.
	///
	/// This is automatically set to `true` if either video or audio input is enabled.
	wait_for_video_audio: bool,
}

impl BuilderTask {
	pub fn new(
		service_id: ServiceId,
		provision_timeout: Slot,
		engine: Arc<polkavm::Engine>,
		exec_cache: Arc<ExecCache>,
		core_scheduler: Arc<Mutex<CoreScheduler>>,
		input_streams: Arc<Mutex<BoundedInputStreams>>,
		stopped: Arc<AtomicBool>,
		input_chunk_index: watch::Receiver<InputChunkIndex>,
	) -> Self {
		Self {
			service_id,
			provision_timeout,
			exec_env_hash: None,
			last_processed_slot: None,
			engine,
			exec_cache,
			core_scheduler,
			service_messages: ServiceMessages::new(),
			input_streams,
			stopped,
			waiting_for_input: None,
			input_chunk_index,
			wait_for_video_audio: false,
		}
	}

	pub async fn run(mut self, mut best_block_rx: watch::Receiver<(BlockDesc, Arc<dyn Node>)>) {
		let mut builder: Option<Builder> = None;
		while matches!(
			self.do_run(&mut builder, &mut best_block_rx)
				.await
				.inspect(|(step, outcome)| log::debug!(
					service_id = self.service_id;
					"Builder of {:x} finished step {step} with outcome \"{outcome}\"",
					self.service_id,
				))
				.inspect_err(|e| log::debug!(
					service_id = self.service_id;
					"Builder of {:x} failed: {e}",
					self.service_id
				)),
			Ok((_, StepOutcome::Continue)) | Err(NodeError::NetworkFailure)
		) {}
	}

	async fn do_run(
		&mut self,
		builder: &mut Option<Builder>,
		best_block_rx: &mut watch::Receiver<(BlockDesc, Arc<dyn Node>)>,
	) -> Result<(u32, StepOutcome), NodeError> {
		let (block, node) = best_block_rx
			.wait_for(|(block, _node)| self.last_processed_slot.unwrap_or(0) < block.slot)
			.await
			.map_err(|_| NodeError::Other("General builder stopped".into()))?
			.clone();
		let (step, outcome) = self.step(block, node, builder).await?;
		self.last_processed_slot = Some(block.slot);
		Ok((step, outcome))
	}

	async fn step(
		&mut self,
		block: BlockDesc,
		node: Arc<dyn Node>,
		builder: &mut Option<Builder>,
	) -> Result<(u32, StepOutcome), NodeError> {
		if self.exec_cache.is_timed_out(self.service_id) {
			return Err(NodeError::Other("Guest timed out".into()));
		}
		let (code_hash, code_len) = self.check_is_corevm(block.header_hash, node.as_ref()).await?;
		let (exec_env_ref, vm_reset) =
			self.check_vm_reset(block.header_hash, node.as_ref()).await?;
		let builder = match builder {
			Some(ref mut builder) if !vm_reset => builder,
			old_builder => {
				// TODO @ivan Memorize for how long the preimages are provided and force prefetching
				// afterwards.
				*old_builder = None;
				let new_builder = self
					.create_builder(
						node.clone(),
						block.header_hash,
						code_hash,
						code_len,
						exec_env_ref,
					)
					.await?;
				old_builder.insert(new_builder)
			},
		};
		let max_step_time = Duration::from_secs(2 * slot_period_sec());
		match timeout(max_step_time, self.build(node, builder, block, code_hash)).await {
			Ok(ret) => ret,
			// This can also happen if the network is unreachable (or was briefly unreachable).
			// It's unclear how to detect that.
			Err(_) => {
				// Ensure that we will never execute this guest again unless its code/excution
				// parameters are updated (i.e. VM is reset).
				self.exec_cache.set_timed_out(self.service_id, exec_env_ref.hash.0);
				Err(NodeError::Other("Step took more time than two slots".into()))
			},
		}
	}

	/// Checks that the service is CoreVM.
	///
	/// Returns an error if not. On success returns code hash and code blob length.
	async fn check_is_corevm(
		&mut self,
		head: HeaderHash,
		node: &dyn Node,
	) -> Result<(CodeHash, u32), NodeError> {
		let service = node
			.service_data(head, self.service_id)
			.await?
			.ok_or_else(|| NodeError::Other("Service not found".into()))?;
		let code_hash = service.code_hash;
		if code_hash.0 != corevm_bin::HASH {
			return Err(NodeError::Other("Not a CoreVM service".into()));
		}
		let code_len = node
			.service_preimage_len(head, self.service_id, code_hash.0)
			.await?
			.ok_or_else(|| NodeError::Other("Service preimage not found".into()))?;
		Ok((code_hash, code_len))
	}

	/// Check that VM startup parameters were changed compared to the last processed slot.
	///
	/// Also returns `true` if this is the very first slot in which processing started.
	///
	/// Returns FS block reference of `ExecEnv` and `true` if the parameters were changed.
	async fn check_vm_reset(
		&mut self,
		head: HeaderHash,
		node: &dyn Node,
	) -> Result<(corevm_fs::BlockRef, bool), NodeError> {
		let exec_env_ref = node
			.typed_service_value::<corevm_fs::BlockRef>(
				head,
				self.service_id,
				&StorageKey::ExecEnvRef,
			)
			.await?
			.ok_or_else(|| NodeError::Other("No ExecEnvRef".into()))?;
		let is_reset = Some(exec_env_ref.hash) != self.exec_env_hash;
		if is_reset {
			self.exec_cache.remove(self.service_id);
			self.exec_env_hash = Some(exec_env_ref.hash);
		}
		Ok((exec_env_ref, is_reset))
	}

	async fn create_builder(
		&mut self,
		node: Arc<dyn Node>,
		head: HeaderHash,
		code_hash: CodeHash,
		code_len: u32,
		exec_env_ref: corevm_fs::BlockRef,
	) -> Result<Builder, NodeError> {
		node.wait_until_preimage_is_provided(
			self.service_id,
			code_hash.0,
			code_len,
			self.provision_timeout,
		)
		.await?;
		let (gas, vm_spec) = tokio::join!(
			node.typed_service_value::<SignedGas>(head, self.service_id, &StorageKey::Gas),
			node.typed_service_value::<VmSpec>(head, self.service_id, &StorageKey::VmSpec),
		);
		let vm_spec = vm_spec?;
		if let Some(ref vm_spec) = vm_spec {
			match vm_spec.output.outcome {
				Outcome::Halt => return Err(other_error("Can't build service that halted")),
				Outcome::Panic => return Err(other_error("Can't build service that panicked")),
				Outcome::PageFault { .. } |
				Outcome::OutOfGas |
				Outcome::OutputLimitReached |
				Outcome::TimeLimitReached |
				Outcome::InputLimitReached |
				Outcome::WaitingForInput(..) => {},
			}
			self.update_waiting_for_input(vm_spec.output.outcome);
		}
		let gas = gas?.ok_or_else(|| NodeError::Other("Gas not set".into()))?;
		let vm_state = vm_spec.map(|vm_spec| vm_spec.state).unwrap_or(VmState::initial());
		let exec = node.fetch_exec_env(exec_env_ref, Some(self.provision_timeout)).await?;
		let all_pages = node.fetch_pages(head, self.service_id).await?;
		node.prefetch_file_blocks_recursively(exec.program, Some(self.provision_timeout))
			.await
			.map_err(other_error)?;
		if exec.root_dir.hash != Default::default() {
			node.prefetch_file_blocks_recursively(exec.root_dir, Some(self.provision_timeout))
				.await
				.map_err(other_error)?;
		}
		let payload = CoreVmPayload { gas, exec_ref: exec_env_ref, vm_state };
		let input_key = exec
			.input_key
			.map(|key| {
				key.try_into()
					.map_err(|_| NodeError::Other("Failed to parse verification key".into()))
			})
			.transpose()?;
		self.input_streams.lock().reset(input_key);
		self.wait_for_video_audio =
			input_key.is_some() && (exec.video_input.is_some() || exec.audio_input.is_some());
		Ok(Builder::new(
			self.engine.clone(),
			payload,
			all_pages,
			self.provision_timeout,
			self.stopped.clone(),
		))
	}

	async fn build(
		&mut self,
		node: Arc<dyn Node>,
		builder: &mut Builder,
		block: BlockDesc,
		code_hash: CodeHash,
	) -> Result<(u32, StepOutcome), NodeError> {
		let mut timer = Timer::new();
		timer.add("input");
		let step = builder.step();
		let Some((mut extrinsics, mut raw_input_chunks)) =
			self.create_extrinsics(node.clone(), block).await?
		else {
			return Ok((step, StepOutcome::Continue));
		};
		// Simulate the program run to predict which memory pages will be accessed.
		let t = Instant::now();
		let (mut output, imported_pages, exports, input_stats) = builder
			.simulate(node.clone(), self.service_id, extrinsics.clone())
			.await
			.map_err(other_error)?;
		timer.add("simulation");
		let simulation_secs = t.elapsed().as_secs_f32();
		if simulation_secs > slot_period_sec() as f32 * 0.5 {
			log::warn!(
				service_id = self.service_id;
				"Service {:x}: Simulation took more than half a slot ({simulation_secs:.1} s)",
				self.service_id,
			);
		}
		raw_input_chunks.truncate(input_stats.num_input_chunks_read as usize);
		self.truncate_extrinsics(&mut extrinsics, input_stats);
		let raw_service_messages = Bytes::from(extrinsics.incoming_service_messages.encode());
		self.build_on_jam_node(
			node,
			builder,
			code_hash,
			&imported_pages,
			exports.len() as u16,
			raw_service_messages,
			raw_input_chunks,
		)
		.await?;
		timer.add("refining");
		log::trace!(
			service_id = self.service_id;
			"Refined work package in {}: {timer}",
			concise_duration(timer.total())
		);
		self.service_messages
			.push_refined(std::mem::take(&mut output.processed_service_messages))?;
		self.update_waiting_for_input(output.vm_output.outcome);
		let outcome = builder.advance(output, exports);
		Ok((step, outcome))
	}

	async fn build_on_jam_node(
		&mut self,
		node: Arc<dyn Node>,
		builder: &mut Builder,
		code_hash: CodeHash,
		imported_pages: &[PageNum],
		export_count: u16,
		raw_service_messages: Bytes,
		raw_input_chunks: Vec<Bytes>,
	) -> Result<(), NodeError> {
		match self
			.do_build_on_jam_node(
				node.clone(),
				builder,
				code_hash,
				imported_pages,
				export_count,
				raw_service_messages.clone(),
				raw_input_chunks.clone(),
			)
			.await
		{
			Ok(()) => Ok(()),
			Err(PackageError::Failed(message)) => {
				log::info!(
					service_id = self.service_id;
					"Work-package failed: {message}. Resubmitting..."
				);
				// TODO @ivan Currently the resubmission is a workaround for nodes not importing
				// the block (from which we get the work report *on the previous step*)
				// fast enough: the builder has this block already, but the node to which the
				// package is submitted doesn't.
				self.do_build_on_jam_node(
					node,
					builder,
					code_hash,
					imported_pages,
					export_count,
					raw_service_messages,
					raw_input_chunks,
				)
				.await?;
				Ok(())
			},
			Err(PackageError::Node(e)) => Err(e),
		}
	}

	async fn do_build_on_jam_node(
		&mut self,
		node: Arc<dyn Node>,
		builder: &mut Builder,
		code_hash: CodeHash,
		imported_pages: &[PageNum],
		export_count: u16,
		raw_service_messages: Bytes,
		raw_input_chunks: Vec<Bytes>,
	) -> Result<(), PackageError> {
		let (context, ..) = node.create_refine_context().await?;
		let anchor = context.anchor;
		let (package_hash, package_bundle) = builder
			.create_work_package_bundle(
				self.service_id,
				code_hash,
				context,
				imported_pages,
				export_count,
				raw_service_messages,
				raw_input_chunks,
			)
			.map_err(other_error)?;
		let core = self.core_scheduler.lock().next_core();
		node.submit_encoded_work_package_bundle(core, package_bundle).await?;
		node.wait_for_package(package_hash, anchor).await?;
		Ok(())
	}

	/// Creates extrinsics from the input provided by the monitor and from inter-service messages
	/// fetched from the service storage.
	///
	/// Returns extrinsics and raw (signed) input chunks.
	///
	/// Returns `None` if all of the input streams the guest is waiting for are empty.
	async fn create_extrinsics(
		&mut self,
		node: Arc<dyn Node>,
		block: BlockDesc,
	) -> Result<Option<(CoreVmExtrinsics, Vec<Bytes>)>, NodeError> {
		// N.B. Monitor records input in the previous slot.
		let input_slot = block.slot.saturating_sub(1);
		if self.wait_for_video_audio {
			self.wait_for_last_input_chunk_in_slot(input_slot).await?;
		}
		let (input_chunks, raw_input_chunks) = {
			let mut input_streams = self.input_streams.lock();
			input_streams.remove_older_than(input_slot);
			if let Some(streams) = self.waiting_for_input {
				if !input_streams.has_any(block.slot, streams) {
					log::debug!(service_id = self.service_id; "Waiting for input: {streams:?}");
					return Ok(None);
				}
			}
			input_streams.take_chunks(block.slot)
		};
		let extrinsics = CoreVmExtrinsics {
			// Will be set by the builder.
			imported_memory_pages: Default::default(),
			incoming_service_messages: self
				.unrefined_inter_service_messages(node.as_ref(), block.header_hash)
				.await?,
			input_chunks,
		};
		Ok(Some((extrinsics, raw_input_chunks)))
	}

	/// Waits for the last input chunk in the specified slot to arrive from the monitor.
	///
	/// This is only practical for video/audio streams that are supposed to be continuous.
	async fn wait_for_last_input_chunk_in_slot(&mut self, slot: Slot) -> Result<(), NodeError> {
		let last_chunk_index = InputChunkIndex::last_in_slot(slot);
		let wait_for_last_input_chunk = self.input_chunk_index.wait_for(|index| {
			let success = *index >= last_chunk_index;
			if !success {
				log::trace!(
					service_id = self.service_id;
					"Waiting for the last input chunk in slot {slot} to arrive..."
				);
			}
			success
		});
		let period = proto::input_chunk_period();
		match timeout(period, wait_for_last_input_chunk).await {
			Ok(ret) => {
				ret.map_err(|_| NodeError::Other("General builder stopped".into()))?;
			},
			Err(..) => {
				// Continue without the last chunk.
				log::debug!(
					service_id = self.service_id;
					"Failed to receive the last input chunk in slot {slot}"
				);
			},
		};
		Ok(())
	}

	fn update_waiting_for_input(&mut self, outcome: Outcome) {
		self.waiting_for_input = match outcome {
			Outcome::WaitingForInput(streams) => Some(streams),
			_ => None,
		};
	}
}

fn other_error(e: impl ToString) -> NodeError {
	NodeError::Other(e.to_string())
}
