use crate::{
	proto::{InputChunkIndex, CHUNKS_PER_SLOT},
	BuilderTask,
};
use bytes::Bytes;
use corevm_engine::InputStats;
use corevm_host::{CoreVmExtrinsics, InputChunk, InputStreams};
use ed25519_consensus::VerificationKey;
use jam_types::{max_input, Slot};
use std::collections::VecDeque;

struct Entry {
	/// Input chunk index.
	///
	/// This is like JAM time-slot but more frequent.
	index: InputChunkIndex,
	/// Decoded and verified input chunk.
	chunk: InputChunk,
	/// Original encoded input chunk with signature.
	raw_chunk: Bytes,
}

/// Per-service input streams.
///
/// Max. size of the input is bounded by the max. bundle size plus some small amount to account for
/// asynchrony between the monitor and the builder.
pub struct BoundedInputStreams {
	/// Input chunks.
	chunks: VecDeque<Entry>,
	/// Total size of all streams in bytes.
	in_memory_size: u32,
	/// Maximum size of all streams in bytes.
	max_size: u32,
	/// Ed25519 public key that is used to verify input chunks' signatures.
	input_key: Option<VerificationKey>,
}

impl BoundedInputStreams {
	pub fn new() -> Self {
		let max_input = max_input();
		// We add one chunk to account for asynchrony between the monitor and the builder.
		let max_size = max_input + max_input / CHUNKS_PER_SLOT;
		Self { chunks: VecDeque::new(), in_memory_size: 0, max_size, input_key: None }
	}

	/// Clears all streams and sets input verification key.
	pub fn reset(&mut self, input_key: Option<VerificationKey>) {
		// Force memory deallocation via assignment.
		self.chunks = VecDeque::new();
		self.in_memory_size = 0;
		self.input_key = input_key;
	}

	/// Removes all data older than `min_slot`.
	pub fn remove_older_than(&mut self, min_slot: Slot) {
		self.chunks.retain(|Entry { index, raw_chunk, .. }| {
			let slot = index.slot();
			let retain = min_slot <= slot;
			if !retain {
				log::debug!("Removing input chunk at slot #{slot}; min. slot is {min_slot}");
				self.in_memory_size -= raw_chunk.len() as u32;
			}
			retain
		});
	}

	/// Append new input chunk with the specified index.
	///
	/// Returns an error if the memory was exhausted or if decoding fails.
	pub fn append_chunk(
		&mut self,
		index: InputChunkIndex,
		raw_chunk: Bytes,
	) -> Result<(), AppendError> {
		let key = self.input_key.as_ref().ok_or(AppendError::NoVerificationKey)?;
		let chunk = InputChunk::decode_signed(raw_chunk.clone(), key)?;
		let delta = raw_chunk.len() as u32;
		// InputChunk uses zero-copy decoding, so we should count chunk size only once even though
		// we store both raw and decoded form.
		if self.in_memory_size + delta > self.max_size {
			return Err(AppendError::TooLargeInput);
		}
		self.chunks.push_back(Entry { index, chunk, raw_chunk });
		self.in_memory_size += delta;
		log::debug!("Appended input chunk #{index} at slot #{}", index.slot());
		Ok(())
	}

	/// Removes all input chunks up to `max_slot` and returns them in both decoded and raw form.
	///
	/// Raw form includes signatures, and this is the form that should be used to supply input
	/// chunks to the guest.
	pub fn take_chunks(&mut self, max_slot: Slot) -> (Vec<InputChunk>, Vec<Bytes>) {
		let max_index = InputChunkIndex::last_in_slot(max_slot);
		let mut chunks = Vec::new();
		let mut raw_chunks = Vec::new();
		while let Some(entry) = self.chunks.pop_front_if(|entry| entry.index <= max_index) {
			let Entry { chunk, raw_chunk, .. } = entry;
			self.in_memory_size -= raw_chunk.len() as u32;
			chunks.push(chunk);
			raw_chunks.push(raw_chunk);
		}
		(chunks, raw_chunks)
	}

	/// Returns `true` if any of the specified streams are non-empty in slots up to `max_slot`.
	pub fn has_any(&self, max_slot: Slot, streams: InputStreams) -> bool {
		let max_index = InputChunkIndex::last_in_slot(max_slot);
		for Entry { index, chunk, .. } in self.chunks.iter() {
			if max_index < *index {
				break;
			}
			if chunk.has_any(streams) {
				return true;
			}
		}
		false
	}
}

#[derive(thiserror::Error, Debug)]
pub enum AppendError {
	#[error("No verification key")]
	NoVerificationKey,
	#[error("Invalid chunk: {0}")]
	Codec(#[from] codec::Error),
	#[error("Guest input size exceeded max. input size")]
	TooLargeInput,
}

impl BuilderTask {
	/// Truncates extrinsics' components to only include portions of input streams and message
	/// queues that were actually consumed during the simulation.
	///
	/// The unprocessed input is removed from the extrinsics and destroyed. The unprocessed
	/// inter-service messages, however, are retained by the builder and will be included in the
	/// extrinsics in the next guest invocation.
	pub(crate) fn truncate_extrinsics(
		&mut self,
		extrinsics: &mut CoreVmExtrinsics,
		bundle: InputStats,
	) {
		if bundle.num_service_messages_read != extrinsics.incoming_service_messages.len() as u32 {
			log::debug!(
				service_id = self.service_id;
				"Guest processed only {} inter-service messages out of {}",
				bundle.num_service_messages_read, extrinsics.incoming_service_messages.len()
			);
			extrinsics
				.incoming_service_messages
				.drain(bundle.num_service_messages_read as usize..);
		}
		if bundle.num_input_chunks_read != extrinsics.input_chunks.len() as u32 {
			log::debug!(
				service_id = self.service_id; "Guest processed only {} input chunks out of {}",
				bundle.num_input_chunks_read, extrinsics.input_chunks.len()
			);
			extrinsics.input_chunks.truncate(bundle.num_input_chunks_read as usize);
		}
	}
}

trait InputChunkExt {
	/// Returns `true` if any of the specified streams are present in the input chunk.
	fn has_any(&self, streams: InputStreams) -> bool;
}

impl InputChunkExt for InputChunk {
	fn has_any(&self, streams: InputStreams) -> bool {
		for stream in streams.iter() {
			match stream {
				InputStreams::HOST_MESSAGES if !self.host_messages.is_empty() => return true,
				InputStreams::CONSOLE if !self.console.is_empty() => return true,
				InputStreams::VIDEO if !self.video.is_empty() => return true,
				InputStreams::AUDIO if !self.audio.is_empty() => return true,
				// The following streams aren't part of the input chunk.
				InputStreams::SERVICE_MESSAGES => {},
				InputStreams::TRANSACTIONS => {},
				_ => {},
			}
		}
		false
	}
}
