use alloc::vec::Vec;
use crate::elements;
use super::{
	import,
	export,
	global,
	data,
	invoke::{Invoke, Identity},
	code::{self, SignaturesBuilder, FunctionBuilder},
	memory::{self, MemoryBuilder},
	table::{self, TableBuilder},
};

/// Module builder
pub struct ModuleBuilder<F=Identity> {
	callback: F,
	module: ModuleScaffold,
}

/// Location of the internal module function
pub struct CodeLocation {
	/// Location (index in 'functions' section) of the signature
	pub signature: u32,
	/// Location (index in the 'code' section) of the body
	pub body: u32,
}

#[derive(Default, PartialEq)]
struct ModuleScaffold {
	pub types: elements::TypeSection,
	pub import: elements::ImportSection,
	pub functions: elements::FunctionSection,
	pub table: elements::TableSection,
	pub memory: elements::MemorySection,
	pub global: elements::GlobalSection,
	pub export: elements::ExportSection,
	pub start: Option<u32>,
	pub element: elements::ElementSection,
	pub code: elements::CodeSection,
	pub data: elements::DataSection,
	pub other: Vec<elements::Section>,
}

impl From<elements::Module> for ModuleScaffold {
	fn from(module: elements::Module) -> Self {
		let mut types: Option<elements::TypeSection> = None;
		let mut import: Option<elements::ImportSection> = None;
		let mut funcs: Option<elements::FunctionSection> = None;
		let mut table: Option<elements::TableSection> = None;
		let mut memory: Option<elements::MemorySection> = None;
		let mut global: Option<elements::GlobalSection> = None;
		let mut export: Option<elements::ExportSection> = None;
		let mut start: Option<u32> = None;
		let mut element: Option<elements::ElementSection> = None;
		let mut code: Option<elements::CodeSection> = None;
		let mut data: Option<elements::DataSection> = None;

		let mut other = Vec::new();
		let mut sections = module.into_sections();
		while let Some(section) = sections.pop() {
			match section {
				elements::Section::Type(sect) => { types = Some(sect); }
				elements::Section::Import(sect) => { import = Some(sect); }
				elements::Section::Function(sect) => { funcs = Some(sect); }
				elements::Section::Table(sect) => { table = Some(sect); }
				elements::Section::Memory(sect) => { memory = Some(sect); }
				elements::Section::Global(sect) => { global = Some(sect); }
				elements::Section::Export(sect) => { export = Some(sect); }
				elements::Section::Start(index) => { start = Some(index); }
				elements::Section::Element(sect) => { element = Some(sect); }
				elements::Section::Code(sect) => { code = Some(sect); }
				elements::Section::Data(sect) => { data = Some(sect); }
				section => other.push(section)
			}
		}

		ModuleScaffold {
			types: types.unwrap_or_default(),
			import: import.unwrap_or_default(),
			functions: funcs.unwrap_or_default(),
			table: table.unwrap_or_default(),
			memory: memory.unwrap_or_default(),
			global: global.unwrap_or_default(),
			export: export.unwrap_or_default(),
			start: start,
			element: element.unwrap_or_default(),
			code: code.unwrap_or_default(),
			data: data.unwrap_or_default(),
			other,
		}
	}
}

impl From<ModuleScaffold> for elements::Module {
	fn from(module: ModuleScaffold) -> Self {
		let mut sections = Vec::new();

		let types = module.types;
		if types.types().len() > 0 {
			sections.push(elements::Section::Type(types));
		}
		let import = module.import;
		if import.entries().len() > 0 {
			sections.push(elements::Section::Import(import));
		}
		let functions = module.functions;
		if functions.entries().len() > 0 {
			sections.push(elements::Section::Function(functions));
		}
		let table = module.table;
		if table.entries().len() > 0 {
			sections.push(elements::Section::Table(table));
		}
		let memory = module.memory;
		if memory.entries().len() > 0 {
			sections.push(elements::Section::Memory(memory));
		}
		let global = module.global;
		if global.entries().len() > 0 {
			sections.push(elements::Section::Global(global));
		}
		let export = module.export;
		if export.entries().len() > 0 {
			sections.push(elements::Section::Export(export));
		}
		if let Some(start) = module.start {
			sections.push(elements::Section::Start(start));
		}
		let element = module.element;
		if element.entries().len() > 0 {
			sections.push(elements::Section::Element(element));
		}
		let code = module.code;
		if code.bodies().len() > 0 {
			sections.push(elements::Section::Code(code));
		}
		let data = module.data;
		if data.entries().len() > 0 {
			sections.push(elements::Section::Data(data));
		}
		sections.extend(module.other);
		elements::Module::new(sections)
	}
}

impl ModuleBuilder {
	/// New empty module builder
	pub fn new() -> Self {
		ModuleBuilder::with_callback(Identity)
	}
}

impl<F> ModuleBuilder<F> where F: Invoke<elements::Module> {
	/// New module builder with bound callback
	pub fn with_callback(callback: F) -> Self {
		ModuleBuilder {
			callback: callback,
			module: Default::default(),
		}
	}

	/// Builder from raw module
	pub fn with_module(mut self, module: elements::Module) -> Self {
		self.module = module.into();
		self
	}

	/// Fill module with sections from iterator
	pub fn with_sections<I>(mut self, sections: I) -> Self
		where I: IntoIterator<Item=elements::Section>
	{
		self.module.other.extend(sections);
		self
	}

	/// Add additional section
	pub fn with_section(mut self, section: elements::Section) -> Self {
		self.module.other.push(section);
		self
	}

	/// Binds to the type section, creates additional types when required
	pub fn with_signatures(mut self, bindings: code::SignatureBindings) -> Self {
		self.push_signatures(bindings);
		self
	}

	/// Push stand-alone function definition, creating sections, signature and code blocks
	/// in corresponding sections.
	/// `FunctionDefinition` can be build using `builder::function` builder
	pub fn push_function(&mut self, func: code::FunctionDefinition) -> CodeLocation {
		let signature = func.signature;
		let body = func.code;

		let type_ref = self.resolve_type_ref(signature);

		self.module.functions.entries_mut().push(elements::Func::new(type_ref));
		let signature_index = self.module.functions.entries_mut().len() as u32 - 1;
		self.module.code.bodies_mut().push(body);
		let body_index = self.module.code.bodies_mut().len() as u32 - 1;

		if func.is_main {
			self.module.start = Some(body_index);
		}

		CodeLocation {
			signature: signature_index,
			body: body_index,
		}
	}

	/// Push linear memory region
	pub fn push_memory(&mut self, mut memory: memory::MemoryDefinition) -> u32 {
		let entries = self.module.memory.entries_mut();
		entries.push(elements::MemoryType::new(memory.min, memory.max));
		let memory_index = (entries.len() - 1) as u32;
		for data in memory.data.drain(..) {
			self.module.data.entries_mut()
				.push(elements::DataSegment::new(memory_index, Some(data.offset), data.values))
		}
		memory_index
	}

	/// Push table
	pub fn push_table(&mut self, mut table: table::TableDefinition) -> u32 {
		let entries = self.module.table.entries_mut();
		entries.push(elements::TableType::new(table.min, table.max));
		let table_index = (entries.len() - 1) as u32;
		for entry in table.elements.drain(..) {
			self.module.element.entries_mut()
				.push(elements::ElementSegment::new(table_index, Some(entry.offset), entry.values))
		}
		table_index
	}

	/// Push global.
	pub fn push_global(&mut self, global: elements::GlobalEntry) -> u32 {
		let entries = self.module.global.entries_mut();
		entries.push(global);
		entries.len() as u32 - 1
	}

	fn resolve_type_ref(&mut self, signature: code::Signature) -> u32 {
		match signature {
			code::Signature::Inline(func_type) => {
				if let Some(existing_entry) = self.module.types.types().iter().enumerate().find(|(_idx, t)| {
					let elements::Type::Function(ref existing) = t;
					*existing == func_type
				}) {
					return existing_entry.0 as u32
				}
				self.module.types.types_mut().push(elements::Type::Function(func_type));
				self.module.types.types().len() as u32 - 1
			}
			code::Signature::TypeReference(type_ref) => {
				type_ref
			}
		}
	}

	/// Push one function signature, returning it's calling index.
	/// Can create corresponding type in type section.
	pub fn push_signature(&mut self, signature: code::Signature) -> u32 {
		self.resolve_type_ref(signature)
	}

	/// Push signatures in the module, returning corresponding indices of pushed signatures
	pub fn push_signatures(&mut self, signatures: code::SignatureBindings) -> Vec<u32> {
		signatures.into_iter().map(|binding|
			self.resolve_type_ref(binding)
		).collect()
	}

	/// Push import entry to module. Note that this does not update calling indices in
	/// function bodies.
	pub fn push_import(&mut self, import: elements::ImportEntry) -> u32 {
		self.module.import.entries_mut().push(import);
		// todo: actually update calling addresses in function bodies
		// todo: also batch push

		self.module.import.entries_mut().len() as u32 - 1
	}

	/// Push export entry to module.
	pub fn push_export(&mut self, export: elements::ExportEntry) -> u32 {
		self.module.export.entries_mut().push(export);
		self.module.export.entries_mut().len() as u32 - 1
	}

	/// Add new function using dedicated builder
	pub fn function(self) -> FunctionBuilder<Self> {
		FunctionBuilder::with_callback(self)
	}

	/// Add new linear memory using dedicated builder
	pub fn memory(self) -> MemoryBuilder<Self> {
		MemoryBuilder::with_callback(self)
	}

	/// Add new table using dedicated builder
	pub fn table(self) -> TableBuilder<Self> {
		TableBuilder::with_callback(self)
	}

	/// Define functions section
	pub fn functions(self) -> SignaturesBuilder<Self> {
		SignaturesBuilder::with_callback(self)
	}

	/// With inserted export entry
	pub fn with_export(mut self, entry: elements::ExportEntry) -> Self {
		self.module.export.entries_mut().push(entry);
		self
	}

	/// With inserted import entry
	pub fn with_import(mut self, entry: elements::ImportEntry) -> Self {
		self.module.import.entries_mut().push(entry);
		self
	}

	/// Import entry builder
	/// # Examples
	/// ```
	/// use parity_wasm::builder::module;
	///
	/// let module = module()
	///    .import()
	///        .module("env")
	///        .field("memory")
	///        .external().memory(256, Some(256))
	///        .build()
	///    .build();
	///
	/// assert_eq!(module.import_section().expect("import section to exist").entries().len(), 1);
	/// ```
	pub fn import(self) -> import::ImportBuilder<Self> {
		import::ImportBuilder::with_callback(self)
	}

	/// With global variable
	pub fn with_global(mut self, global: elements::GlobalEntry) -> Self {
		self.module.global.entries_mut().push(global);
		self
	}

	/// With table
	pub fn with_table(mut self, table: elements::TableType) -> Self {
		self.module.table.entries_mut().push(table);
		self
	}

	/// Export entry builder
	/// # Examples
	/// ```
	/// use parity_wasm::builder::module;
	/// use parity_wasm::elements::Instruction::*;
	///
	/// let module = module()
	///    .global()
	///         .value_type().i32()
	///         .init_expr(I32Const(0))
	///         .build()
	///    .export()
	///        .field("_zero")
	///        .internal().global(0)
	///        .build()
	///    .build();
	///
	/// assert_eq!(module.export_section().expect("export section to exist").entries().len(), 1);
	/// ```
	pub fn export(self) -> export::ExportBuilder<Self> {
		export::ExportBuilder::with_callback(self)
	}

	/// Glboal entry builder
	/// # Examples
	/// ```
	/// use parity_wasm::builder::module;
	/// use parity_wasm::elements::Instruction::*;
	///
	/// let module = module()
	///    .global()
	///         .value_type().i32()
	///         .init_expr(I32Const(0))
	///         .build()
	///    .build();
	///
	/// assert_eq!(module.global_section().expect("global section to exist").entries().len(), 1);
	/// ```
	pub fn global(self) -> global::GlobalBuilder<Self> {
		global::GlobalBuilder::with_callback(self)
	}

	/// Add data segment to the builder
	pub fn with_data_segment(mut self, segment: elements::DataSegment) -> Self {
		self.module.data.entries_mut().push(segment);
		self
	}

	/// Data entry builder
	pub fn data(self) -> data::DataSegmentBuilder<Self> {
		data::DataSegmentBuilder::with_callback(self)
	}

	/// Build module (final step)
	pub fn build(self) -> F::Result {
		self.callback.invoke(self.module.into())
	}
}

impl<F> Invoke<elements::FunctionSection> for ModuleBuilder<F>
	where F: Invoke<elements::Module>
{
	type Result = Self;

	fn invoke(self, section: elements::FunctionSection) -> Self {
		self.with_section(elements::Section::Function(section))
	}
}

impl<F> Invoke<code::SignatureBindings> for ModuleBuilder<F>
	where F: Invoke<elements::Module>
{
	type Result = Self;

	fn invoke(self, bindings: code::SignatureBindings) -> Self {
		self.with_signatures(bindings)
	}
}


impl<F> Invoke<code::FunctionDefinition> for ModuleBuilder<F>
	where F: Invoke<elements::Module>
{
	type Result = Self;

	fn invoke(self, def: code::FunctionDefinition) -> Self {
		let mut b = self;
		b.push_function(def);
		b
	}
}

impl<F> Invoke<memory::MemoryDefinition> for ModuleBuilder<F>
	where F: Invoke<elements::Module>
{
	type Result = Self;

	fn invoke(self, def: memory::MemoryDefinition) -> Self {
		let mut b = self;
		b.push_memory(def);
		b
	}
}

impl<F> Invoke<table::TableDefinition> for ModuleBuilder<F>
	where F: Invoke<elements::Module>
{
	type Result = Self;

	fn invoke(self, def: table::TableDefinition) -> Self {
		let mut b = self;
		b.push_table(def);
		b
	}
}

impl<F> Invoke<elements::ImportEntry> for ModuleBuilder<F>
	where F: Invoke<elements::Module>
{
	type Result = Self;

	fn invoke(self, entry: elements::ImportEntry) -> Self::Result {
		self.with_import(entry)
	}
}

impl<F> Invoke<elements::ExportEntry> for ModuleBuilder<F>
	where F: Invoke<elements::Module>
{
	type Result = Self;

	fn invoke(self, entry: elements::ExportEntry) -> Self::Result {
		self.with_export(entry)
	}
}

impl<F> Invoke<elements::GlobalEntry> for ModuleBuilder<F>
	where F: Invoke<elements::Module>
{
	type Result = Self;

	fn invoke(self, entry: elements::GlobalEntry) -> Self::Result {
		self.with_global(entry)
	}
}

impl<F> Invoke<elements::DataSegment> for ModuleBuilder<F>
	where F: Invoke<elements::Module>
{
	type Result = Self;

	fn invoke(self, segment: elements::DataSegment) -> Self {
		self.with_data_segment(segment)
	}
}

/// Start new module builder
/// # Examples
///
/// ```
/// use parity_wasm::builder;
///
/// let module = builder::module()
///     .function()
///         .signature().param().i32().build()
///         .body().build()
///         .build()
///     .build();
///
/// assert_eq!(module.type_section().expect("type section to exist").types().len(), 1);
/// assert_eq!(module.function_section().expect("function section to exist").entries().len(), 1);
/// assert_eq!(module.code_section().expect("code section to exist").bodies().len(), 1);
/// ```
pub fn module() -> ModuleBuilder {
	ModuleBuilder::new()
}

/// Start builder to extend existing module
pub fn from_module(module: elements::Module) -> ModuleBuilder {
	ModuleBuilder::new().with_module(module)
}

#[cfg(test)]
mod tests {

	use crate::elements;
	use super::module;

	#[test]
	fn smoky() {
		let module = module().build();
		assert_eq!(module.sections().len(), 0);
	}

	#[test]
	fn functions() {
		let module = module()
			.function()
				.signature().param().i32().build()
				.body().build()
				.build()
			.build();

		assert_eq!(module.type_section().expect("type section to exist").types().len(), 1);
		assert_eq!(module.function_section().expect("function section to exist").entries().len(), 1);
		assert_eq!(module.code_section().expect("code section to exist").bodies().len(), 1);
	}

	#[test]
	fn export() {
		let module = module()
			.export().field("call").internal().func(0).build()
			.build();

		assert_eq!(module.export_section().expect("export section to exist").entries().len(), 1);
	}

	#[test]
	fn global() {
		let module = module()
			.global().value_type().i64().mutable().init_expr(elements::Instruction::I64Const(5)).build()
			.build();

		assert_eq!(module.global_section().expect("global section to exist").entries().len(), 1);
	}

	#[test]
	fn data() {
		let module = module()
			.data()
				.offset(elements::Instruction::I32Const(16))
				.value(vec![0u8, 15, 10, 5, 25])
				.build()
			.build();

		assert_eq!(module.data_section().expect("data section to exist").entries().len(), 1);
	}

	#[test]
	fn reuse_types() {
		let module = module()
			.function()
				.signature().param().i32().build()
				.body().build()
				.build()
			.function()
				.signature().param().i32().build()
				.body().build()
				.build()
			.build();

		assert_eq!(module.type_section().expect("type section failed").types().len(), 1);
	}
 }