// This file is part of Substrate.

// Copyright (C) 2019-2021 Parity Technologies (UK) Ltd.
// SPDX-License-Identifier: Apache-2.0

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// 	http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::{from_assignment_helpers::*, syn_err, vote_field};
use proc_macro2::TokenStream as TokenStream2;
use quote::quote;
use syn::parse::Result;

pub(crate) fn generate(def: crate::SolutionDef) -> Result<TokenStream2> {
	let crate::SolutionDef {
		vis,
		ident,
		count,
		voter_type,
		target_type,
		weight_type,
		compact_encoding,
	} = def;

	if count <= 2 {
		Err(syn_err("cannot build solution struct with capacity less than 3."))?
	}

	let single = {
		let name = vote_field(1);
		// NOTE: we use the visibility of the struct for the fields as well.. could be made better.
		quote!(
			#vis #name: _npos::sp_std::prelude::Vec<(#voter_type, #target_type)>,
		)
	};

	let rest = (2..=count)
		.map(|c| {
			let field_name = vote_field(c);
			let array_len = c - 1;
			quote!(
				#vis #field_name: _npos::sp_std::prelude::Vec<(
					#voter_type,
					[(#target_type, #weight_type); #array_len],
					#target_type
				)>,
			)
		})
		.collect::<TokenStream2>();

	let len_impl = len_impl(count);
	let edge_count_impl = edge_count_impl(count);
	let unique_targets_impl = unique_targets_impl(count);
	let remove_voter_impl = remove_voter_impl(count);

	let derives_and_maybe_compact_encoding = if compact_encoding {
		// custom compact encoding.
		let compact_impl = crate::codec::codec_impl(
			ident.clone(),
			voter_type.clone(),
			target_type.clone(),
			weight_type.clone(),
			count,
		);
		quote! {
			#compact_impl
			#[derive(Default, PartialEq, Eq, Clone, Debug, PartialOrd, Ord)]
		}
	} else {
		// automatically derived.
		quote!(#[derive(Default, PartialEq, Eq, Clone, Debug, _npos::codec::Encode, _npos::codec::Decode)])
	};

	let struct_name = syn::Ident::new("solution", proc_macro2::Span::call_site());
	let assignment_name = syn::Ident::new("all_assignments", proc_macro2::Span::call_site());

	let from_impl = from_impl(&struct_name, count);
	let into_impl = into_impl(&assignment_name, count, weight_type.clone());
	let from_index_impl = crate::index_assignment::from_impl(&struct_name, count);

	Ok(quote! (
		/// A struct to encode a election assignment in a compact way.
		#derives_and_maybe_compact_encoding
		#vis struct #ident { #single #rest }

		use _npos::__OrInvalidIndex;
		impl _npos::NposSolution for #ident {
			const LIMIT: usize = #count;
			type VoterIndex = #voter_type;
			type TargetIndex = #target_type;
			type Accuracy = #weight_type;

			fn remove_voter(&mut self, to_remove: Self::VoterIndex) -> bool {
				#remove_voter_impl
				return false
			}

			fn from_assignment<FV, FT, A>(
				assignments: &[_npos::Assignment<A, #weight_type>],
				voter_index: FV,
				target_index: FT,
			) -> Result<Self, _npos::Error>
				where
					A: _npos::IdentifierT,
					for<'r> FV: Fn(&'r A) -> Option<Self::VoterIndex>,
					for<'r> FT: Fn(&'r A) -> Option<Self::TargetIndex>,
			{
				let mut #struct_name: #ident = Default::default();
				for _npos::Assignment { who, distribution } in assignments {
					match distribution.len() {
						0 => continue,
						#from_impl
						_ => {
							return Err(_npos::Error::SolutionTargetOverflow);
						}
					}
				};
				Ok(#struct_name)
			}

			fn into_assignment<A: _npos::IdentifierT>(
				self,
				voter_at: impl Fn(Self::VoterIndex) -> Option<A>,
				target_at: impl Fn(Self::TargetIndex) -> Option<A>,
			) -> Result<_npos::sp_std::prelude::Vec<_npos::Assignment<A, #weight_type>>, _npos::Error> {
				let mut #assignment_name: _npos::sp_std::prelude::Vec<_npos::Assignment<A, #weight_type>> = Default::default();
				#into_impl
				Ok(#assignment_name)
			}

			fn voter_count(&self) -> usize {
				let mut all_len = 0usize;
				#len_impl
				all_len
			}

			fn edge_count(&self) -> usize {
				let mut all_edges = 0usize;
				#edge_count_impl
				all_edges
			}

			fn unique_targets(&self) -> _npos::sp_std::prelude::Vec<Self::TargetIndex> {
				// NOTE: this implementation returns the targets sorted, but we don't use it yet per
				// se, nor is the API enforcing it.
				use _npos::sp_std::collections::btree_set::BTreeSet;
				let mut all_targets: BTreeSet<Self::TargetIndex> = BTreeSet::new();
				let mut maybe_insert_target = |t: Self::TargetIndex| {
					all_targets.insert(t);
				};

				#unique_targets_impl

				all_targets.into_iter().collect()
			}
		}

		type __IndexAssignment = _npos::IndexAssignment<
			<#ident as _npos::NposSolution>::VoterIndex,
			<#ident as _npos::NposSolution>::TargetIndex,
			<#ident as _npos::NposSolution>::Accuracy,
		>;
		impl<'a> _npos::sp_std::convert::TryFrom<&'a [__IndexAssignment]> for #ident {
			type Error = _npos::Error;
			fn try_from(index_assignments: &'a [__IndexAssignment]) -> Result<Self, Self::Error> {
				let mut #struct_name =  #ident::default();

				for _npos::IndexAssignment { who, distribution } in index_assignments {
					match distribution.len() {
						0 => {}
						#from_index_impl
						_ => {
							return Err(_npos::Error::SolutionTargetOverflow);
						}
					}
				};

				Ok(#struct_name)
			}
		}
	))
}

fn remove_voter_impl(count: usize) -> TokenStream2 {
	let field_name = vote_field(1);
	let single = quote! {
		if let Some(idx) = self.#field_name.iter().position(|(x, _)| *x == to_remove) {
			self.#field_name.remove(idx);
			return true
		}
	};

	let rest = (2..=count)
		.map(|c| {
			let field_name = vote_field(c);
			quote! {
				if let Some(idx) = self.#field_name.iter().position(|(x, _, _)| *x == to_remove) {
					self.#field_name.remove(idx);
					return true
				}
			}
		})
		.collect::<TokenStream2>();

	quote! {
		#single
		#rest
	}
}

fn len_impl(count: usize) -> TokenStream2 {
	(1..=count)
		.map(|c| {
			let field_name = vote_field(c);
			quote!(
				all_len = all_len.saturating_add(self.#field_name.len());
			)
		})
		.collect::<TokenStream2>()
}

fn edge_count_impl(count: usize) -> TokenStream2 {
	(1..=count)
		.map(|c| {
			let field_name = vote_field(c);
			quote!(
				all_edges = all_edges.saturating_add(
					self.#field_name.len().saturating_mul(#c as usize)
				);
			)
		})
		.collect::<TokenStream2>()
}

fn unique_targets_impl(count: usize) -> TokenStream2 {
	let unique_targets_impl_single = {
		let field_name = vote_field(1);
		quote! {
			self.#field_name.iter().for_each(|(_, t)| {
				maybe_insert_target(*t);
			});
		}
	};

	let unique_targets_impl_rest = (2..=count)
		.map(|c| {
			let field_name = vote_field(c);
			quote! {
				self.#field_name.iter().for_each(|(_, inners, t_last)| {
					inners.iter().for_each(|(t, _)| {
						maybe_insert_target(*t);
					});
					maybe_insert_target(*t_last);
				});
			}
		})
		.collect::<TokenStream2>();

	quote! {
		#unique_targets_impl_single
		#unique_targets_impl_rest
	}
}

pub(crate) fn from_impl(struct_name: &syn::Ident, count: usize) -> TokenStream2 {
	let from_impl_single = {
		let field = vote_field(1);
		let push_code = from_impl_single_push_code();
		quote!(1 => #struct_name.#field.#push_code,)
	};

	let from_impl_rest = (2..=count)
		.map(|c| {
			let field = vote_field(c);
			let push_code = from_impl_rest_push_code(c);
			quote!(#c => #struct_name.#field.#push_code,)
		})
		.collect::<TokenStream2>();

	quote!(
		#from_impl_single
		#from_impl_rest
	)
}

pub(crate) fn into_impl(
	assignments: &syn::Ident,
	count: usize,
	per_thing: syn::Type,
) -> TokenStream2 {
	let into_impl_single = {
		let name = vote_field(1);
		quote!(
			for (voter_index, target_index) in self.#name {
				#assignments.push(_npos::Assignment {
					who: voter_at(voter_index).or_invalid_index()?,
					distribution: vec![
						(target_at(target_index).or_invalid_index()?, #per_thing::one())
					],
				})
			}
		)
	};

	let into_impl_rest = (2..=count)
		.map(|c| {
			let name = vote_field(c);
			quote!(
				for (voter_index, inners, t_last_idx) in self.#name {
					let mut sum = #per_thing::zero();
					let mut inners_parsed = inners
						.iter()
						.map(|(ref t_idx, p)| {
							sum = _npos::sp_arithmetic::traits::Saturating::saturating_add(sum, *p);
							let target = target_at(*t_idx).or_invalid_index()?;
							Ok((target, *p))
						})
						.collect::<Result<_npos::sp_std::prelude::Vec<(A, #per_thing)>, _npos::Error>>()?;

					if sum >= #per_thing::one() {
						return Err(_npos::Error::SolutionWeightOverflow);
					}

					// defensive only. Since Percent doesn't have `Sub`.
					let p_last = _npos::sp_arithmetic::traits::Saturating::saturating_sub(
						#per_thing::one(),
						sum,
					);

					inners_parsed.push((target_at(t_last_idx).or_invalid_index()?, p_last));

					#assignments.push(_npos::Assignment {
						who: voter_at(voter_index).or_invalid_index()?,
						distribution: inners_parsed,
					});
				}
			)
		})
		.collect::<TokenStream2>();

	quote!(
		#into_impl_single
		#into_impl_rest
	)
}