use crate::{capture_logs_as, CreateWorkPackageError, SimulateError};
use bytes::Bytes;
use codec::Compact;
use corevm_engine::{InputStats, OuterVmSimulator, Simulator, WorkPackageParams};
use corevm_host::{
	CoreVmExtrinsics, CoreVmOutput, CoreVmPayload, Outcome, PageNum, PageSegmentOps,
};
use corevm_tooling::PreimageReader;
use jam_std_common::{build_encoded_bundle, hash_raw, hash_raw_concat, Node};
use jam_types::{
	max_accumulate_gas, max_exports, max_imports, max_input, max_refine_gas, Authorization,
	Authorizer, BoundedVec, CodeHash, Encode, ExtrinsicSpec, MaxExtrinsics, RefineContext,
	SegmentBytes, ServiceId, Slot, WorkItem, WorkPackage, WorkPackageHash, WorkPayload,
	SEGMENT_LEN,
};
use parking_lot::Mutex;
use std::{
	collections::HashMap,
	sync::{atomic::AtomicBool, Arc},
};
use tokio::task::spawn_blocking;

pub struct Builder {
	engine: Arc<polkavm::Engine>,
	/// All pages that are currently allocated by the program.
	all_pages: Arc<Mutex<HashMap<PageNum, SegmentBytes>>>,
	payload: CoreVmPayload,
	provision_timeout: Slot,
	stopped: Arc<AtomicBool>,
}

impl Builder {
	pub fn new(
		engine: Arc<polkavm::Engine>,
		initial_payload: CoreVmPayload,
		all_pages: HashMap<PageNum, SegmentBytes>,
		provision_timeout: Slot,
		stopped: Arc<AtomicBool>,
	) -> Self {
		Self {
			engine,
			all_pages: Arc::new(Mutex::new(all_pages)),
			payload: initial_payload,
			provision_timeout,
			stopped,
		}
	}

	pub fn step(&self) -> u32 {
		self.payload.vm_state.step
	}

	/// Simulate program run locally.
	pub async fn simulate(
		&mut self,
		node: Arc<dyn Node>,
		corevm_service_id: ServiceId,
		mut extrinsics: CoreVmExtrinsics,
	) -> Result<(CoreVmOutput, Vec<PageNum>, Vec<SegmentBytes>, InputStats), SimulateError> {
		use Outcome::*;
		extrinsics.imported_memory_pages = self.all_pages.lock().values().cloned().collect();
		let preimage_reader = PreimageReader::with_provision(node, Some(self.provision_timeout));
		let package_params = WorkPackageParams {
			encoded_size: (WORK_PACKAGE_WITHOUT_PAYLOAD_MAX_SIZE + self.payload.encoded_size())
				as u32,
			auth_output_size: 0,
			export_count: max_exports() as u16,
		};
		let outer_vm =
			OuterVmSimulator::new(self.engine.clone(), preimage_reader, self.stopped.clone())?;
		let payload = self.payload.clone();
		let (work_output, outer_vm, input_stats) = spawn_blocking(move || {
			let _guard = capture_logs_as(corevm_service_id);
			let engine = Simulator::new(payload, extrinsics, package_params, outer_vm)?;
			let ret = engine.run()?;
			Ok::<_, SimulateError>(ret)
		})
		.await??;
		match work_output.vm_output.outcome {
			Panic => return Err(SimulateError::Panic),
			Halt |
			OutOfGas |
			OutputLimitReached |
			TimeLimitReached |
			InputLimitReached |
			WaitingForInput(..) |
			PageFault { .. } => {
				// Finish the simulation.
			},
		}
		assert!(
			work_output.updated_pages.len() <= max_imports() as usize,
			"No. of updates pages = {}, maximum = {}",
			work_output.updated_pages.len(),
			max_imports()
		);
		let all_pages = self.all_pages.lock();
		let exports = outer_vm.into_exports();
		let imports_vec = work_output
			.updated_pages
			.iter()
			.map(|(page, _hash)| *page)
			// Remove duplicates.
			.filter(|page| !input_stats.accessed_imported_pages.contains(page))
			.filter(|page| {
				// Take only those memory pages that we already know.
				all_pages.contains_key(page)
			})
			.chain(input_stats.accessed_imported_pages.iter().copied())
			.collect();
		Ok((work_output, imports_vec, exports, input_stats))
	}

	pub fn create_work_package_bundle(
		&self,
		corevm_service_id: ServiceId,
		corevm_code_hash: CodeHash,
		context: RefineContext,
		imported_pages: &[PageNum],
		export_count: u16,
		raw_service_messages: Bytes,
		raw_input_chunks: Vec<Bytes>,
	) -> Result<(WorkPackageHash, Bytes), CreateWorkPackageError> {
		let raw_memory_pages = {
			let mut segments = Vec::with_capacity(imported_pages.len());
			let all_pages = self.all_pages.lock();
			for page in imported_pages {
				let page = all_pages.get(page).ok_or(CreateWorkPackageError::NoSuchPage(*page))?;
				segments.push(page.as_bytes().clone());
			}
			segments
		};
		// Prepend encoded length to each input chunk.
		let raw_input_chunks_with_lengths: Vec<Bytes> = raw_input_chunks
			.into_iter()
			.flat_map(|chunk| {
				let len_bytes = Bytes::from(Compact(chunk.len() as u32).encode());
				[len_bytes, chunk].into_iter()
			})
			.collect();
		let extrinsics_specs = vec![
			ExtrinsicSpec {
				hash: hash_raw_concat(raw_memory_pages.iter()).into(),
				len: raw_memory_pages.len() as u32 * SEGMENT_LEN as u32,
			},
			ExtrinsicSpec {
				hash: hash_raw(&raw_service_messages).into(),
				len: raw_service_messages.len() as u32,
			},
			ExtrinsicSpec {
				hash: hash_raw_concat(raw_input_chunks_with_lengths.iter()).into(),
				len: raw_input_chunks_with_lengths.iter().map(|chunk| chunk.len() as u32).sum(),
			},
		];
		let extrinsics_len: u32 = extrinsics_specs.iter().map(|xt| xt.len).sum();
		assert!(
			extrinsics_len <= max_input(),
			"extrinsics length = {extrinsics_len}, extrinsics = {:?}, max = {}",
			extrinsics_specs,
			max_input()
		);
		let work_package = create_work_package(
			context,
			corevm_service_id,
			corevm_code_hash,
			self.payload.encode().into(),
			extrinsics_specs.try_into().expect("No. of extrinsics is fixed"),
			export_count,
		);
		let chunked_extrinsics = raw_memory_pages
			.into_iter()
			.chain(std::iter::once(raw_service_messages))
			.chain(raw_input_chunks_with_lengths);
		let (wp_hash, bundle) = build_encoded_bundle(&work_package, chunked_extrinsics, &[]);
		Ok((wp_hash, bundle.into()))
	}

	pub fn advance(&mut self, output: CoreVmOutput, exports: Vec<SegmentBytes>) -> StepOutcome {
		use Outcome::*;
		self.payload.vm_state = output.vm_state;
		self.store_pages(exports, output.vm_output.num_memory_pages);
		match output.vm_output.outcome {
			Halt => StepOutcome::Halt,
			Panic => StepOutcome::Panic,
			PageFault { .. } => {
				// CoreVM was unable to handle page fault using the available pages. This might
				// happen even if we used simulator: the simplest example is a program that
				// allocates pages and fills them with some data in a loop. Such a program
				// will reach `MAX_EXPORTS` during a page fault, and the next program run will begin
				// with the same page fault.
				StepOutcome::Continue
			},
			OutOfGas | OutputLimitReached | TimeLimitReached | InputLimitReached |
			WaitingForInput(..) => StepOutcome::Continue,
		}
	}

	fn store_pages(&mut self, exports: Vec<SegmentBytes>, num_memory_pages: u32) {
		let mut all_pages = self.all_pages.lock();
		for segment in exports.into_iter().take(num_memory_pages as usize) {
			let page = segment.as_array().page_number();
			all_pages.insert(page, segment);
		}
	}
}

/// How to advance guest program execution.
pub enum StepOutcome {
	/// Program exited.
	Halt,
	/// Program panicked.
	Panic,
	/// Program either ran out of gas or encountered a page fault.
	Continue,
}

impl StepOutcome {
	const fn as_str(&self) -> &'static str {
		match self {
			Self::Halt => "halt",
			Self::Panic => "panic",
			Self::Continue => "continue",
		}
	}
}

impl std::fmt::Display for StepOutcome {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		f.write_str(self.as_str())
	}
}

/// Creates CoreVM-specific work package.
fn create_work_package(
	context: RefineContext,
	service: ServiceId,
	code_hash: CodeHash,
	payload: WorkPayload,
	extrinsics: BoundedVec<ExtrinsicSpec, MaxExtrinsics>,
	export_count: u16,
) -> WorkPackage {
	WorkPackage {
		authorization: Authorization::new(),
		auth_code_host: 0,
		authorizer: Authorizer {
			code_hash: jam_null_authorizer_bin::HASH.into(),
			config: Default::default(),
		},
		context,
		items: vec![WorkItem {
			service,
			code_hash,
			payload,
			refine_gas_limit: max_refine_gas(),
			accumulate_gas_limit: max_accumulate_gas(),
			import_segments: Default::default(),
			extrinsics,
			export_count,
		}]
		.try_into()
		.expect("Only one work item"),
	}
}

/// Maximum encoded size of a CoreVM-specific work package without a payload.
///
/// Computed in `work_package_without_payload_max_size_works` test.
const WORK_PACKAGE_WITHOUT_PAYLOAD_MAX_SIZE: usize = 336;

#[cfg(test)]
mod tests {
	use super::*;
	use jam_types::WorkPayload;

	#[test]
	fn work_package_without_payload_max_size_works() {
		// Let's create the largest possible CoreVM-specific work package without a payload.
		let largest_corevm_package = create_work_package(
			RefineContext {
				anchor: Default::default(),
				lookup_anchor: Default::default(),
				lookup_anchor_slot: Default::default(),
				beefy_root: Default::default(),
				state_root: Default::default(),
				prerequisites: Default::default(),
			},
			0,
			Default::default(),
			Default::default(),
			vec![ExtrinsicSpec { hash: Default::default(), len: 0 }; CoreVmExtrinsics::MAX_COUNT]
				.try_into()
				.expect("The no. of extrinsics is fixed"),
			0,
		);
		// Subtract encoded size of an empty payload.
		let expected_size =
			largest_corevm_package.encoded_size() - WorkPayload::default().encoded_size();
		assert_eq!(WORK_PACKAGE_WITHOUT_PAYLOAD_MAX_SIZE, expected_size);
	}
}
