// This file is part of Substrate.

// Copyright (C) 2020-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.

//! Proc macro for a npos solution type.

use proc_macro::TokenStream;
use proc_macro2::{Ident, Span, TokenStream as TokenStream2};
use proc_macro_crate::{crate_name, FoundCrate};
use quote::quote;
use syn::parse::{Parse, ParseStream, Result};

mod codec;
mod from_assignment_helpers;
mod index_assignment;
mod single_page;

/// Get the name of a filed based on voter count.
pub(crate) fn vote_field(n: usize) -> Ident {
	quote::format_ident!("votes{}", n)
}

/// Generate a `syn::Error`.
pub(crate) fn syn_err(message: &'static str) -> syn::Error {
	syn::Error::new(Span::call_site(), message)
}

/// Generates a struct to store the election result in a small/compact way. This can encode a
/// structure which is the equivalent of a `sp_npos_elections::Assignment<_>`.
///
/// The following data types can be configured by the macro.
///
/// - The identifier of the voter. This can be any type that supports `parity-scale-codec`'s compact
///   encoding.
/// - The identifier of the target. This can be any type that supports `parity-scale-codec`'s
///   compact encoding.
/// - The accuracy of the ratios. This must be one of the `PerThing` types defined in
///   `sp-arithmetic`.
///
/// Moreover, the maximum number of edges per voter (distribution per assignment) also need to be
/// specified. Attempting to convert from/to an assignment with more distributions will fail.
///
/// For example, the following generates a public struct with name `TestSolution` with `u16` voter
/// type, `u8` target type and `Perbill` accuracy with maximum of 4 edges per voter.
///
/// ```
/// # use sp_npos_elections_solution_type::generate_solution_type;
/// # use sp_arithmetic::per_things::Perbill;
/// generate_solution_type!(pub struct TestSolution::<
///     VoterIndex = u16,
///     TargetIndex = u8,
///     Accuracy = Perbill,
/// >(4));
/// ```
///
/// The output of this macro will roughly look like:
///
/// ```ignore
/// struct TestSolution {
/// 	voters1: vec![(u16 /* voter */, u8 /* target */)]
/// 	voters2: vec![
/// 		(u16 /* voter */, [u8 /* first target*/, Perbill /* proportion for first target */], u8 /* last target */)
/// 	]
/// 	voters3: vec![
/// 		(u16 /* voter */,  [
/// 			(u8 /* first target*/, Perbill /* proportion for first target */ ),
/// 			(u8 /* second target */, Perbill /* proportion for second target*/)
/// 		], u8 /* last target */)
///		],
/// 	voters4: ...,
/// }
///
/// impl NposSolution for TestSolution {};
/// impl Solution for TestSolution {};
/// ```
///
/// The given struct provides function to convert from/to `Assignment` as part of
/// [`sp_npos_elections::Solution`] trait:
///
/// - `fn from_assignment<..>(..)`
/// - `fn into_assignment<..>(..)`
///
/// ## Compact Encoding
///
/// The generated struct is by default deriving both `Encode` and `Decode`. This is okay but could
/// lead to many `0`s in the solution. If prefixed with `#[compact]`, then a custom compact encoding
/// for numbers will be used, similar to how `parity-scale-codec`'s `Compact` works.
///
/// ```
/// # use sp_npos_elections_solution_type::generate_solution_type;
/// # use sp_npos_elections::NposSolution;
/// # use sp_arithmetic::per_things::Perbill;
/// generate_solution_type!(
///     #[compact]
///     pub struct TestSolutionCompact::<VoterIndex = u16, TargetIndex = u8, Accuracy = Perbill>(8)
/// );
/// ```
#[proc_macro]
pub fn generate_solution_type(item: TokenStream) -> TokenStream {
	let solution_def = syn::parse_macro_input!(item as SolutionDef);

	let imports = imports().unwrap_or_else(|e| e.to_compile_error());

	let def = single_page::generate(solution_def).unwrap_or_else(|e| e.to_compile_error());

	quote!(
		#imports
		#def
	)
	.into()
}

struct SolutionDef {
	vis: syn::Visibility,
	ident: syn::Ident,
	voter_type: syn::Type,
	target_type: syn::Type,
	weight_type: syn::Type,
	count: usize,
	compact_encoding: bool,
}

fn check_attributes(input: ParseStream) -> syn::Result<bool> {
	let attrs = input.call(syn::Attribute::parse_outer).unwrap_or_default();
	if attrs.len() > 1 {
		return Err(syn_err("compact solution can accept only #[compact]"))
	}

	Ok(attrs.iter().any(|attr| {
		if attr.path.segments.len() == 1 {
			let segment = attr.path.segments.first().expect("Vec with len 1 can be popped.");
			if segment.ident == Ident::new("compact", Span::call_site()) {
				return true
			}
		}
		false
	}))
}

impl Parse for SolutionDef {
	fn parse(input: ParseStream) -> syn::Result<Self> {
		// optional #[compact]
		let compact_encoding = check_attributes(input)?;

		// <vis> struct <name>
		let vis: syn::Visibility = input.parse()?;
		let _ = <syn::Token![struct]>::parse(input)?;
		let ident: syn::Ident = input.parse()?;

		// ::<V, T, W>
		let _ = <syn::Token![::]>::parse(input)?;
		let generics: syn::AngleBracketedGenericArguments = input.parse()?;

		if generics.args.len() != 3 {
			return Err(syn_err("Must provide 3 generic args."))
		}

		let expected_types = ["VoterIndex", "TargetIndex", "Accuracy"];

		let mut types: Vec<syn::Type> = generics
			.args
			.iter()
			.zip(expected_types.iter())
			.map(|(t, expected)| match t {
				syn::GenericArgument::Type(ty) => {
					// this is now an error
					Err(syn::Error::new_spanned(
						ty,
						format!("Expected binding: `{} = ...`", expected),
					))
				},
				syn::GenericArgument::Binding(syn::Binding { ident, ty, .. }) => {
					// check that we have the right keyword for this position in the argument list
					if ident == expected {
						Ok(ty.clone())
					} else {
						Err(syn::Error::new_spanned(ident, format!("Expected `{}`", expected)))
					}
				},
				_ => Err(syn_err("Wrong type of generic provided. Must be a `type`.")),
			})
			.collect::<Result<_>>()?;

		let weight_type = types.pop().expect("Vector of length 3 can be popped; qed");
		let target_type = types.pop().expect("Vector of length 2 can be popped; qed");
		let voter_type = types.pop().expect("Vector of length 1 can be popped; qed");

		// (<count>)
		let count_expr: syn::ExprParen = input.parse()?;
		let count = parse_parenthesized_number::<usize>(count_expr)?;

		Ok(Self { vis, ident, voter_type, target_type, weight_type, count, compact_encoding })
	}
}

fn parse_parenthesized_number<N: std::str::FromStr>(input_expr: syn::ExprParen) -> syn::Result<N>
where
	<N as std::str::FromStr>::Err: std::fmt::Display,
{
	let expr = input_expr.expr;
	let expr_lit = match *expr {
		syn::Expr::Lit(count_lit) => count_lit.lit,
		_ => return Err(syn_err("Count must be literal.")),
	};
	let int_lit = match expr_lit {
		syn::Lit::Int(int_lit) => int_lit,
		_ => return Err(syn_err("Count must be int literal.")),
	};
	int_lit.base10_parse::<N>()
}

fn imports() -> Result<TokenStream2> {
	match crate_name("sp-npos-elections") {
		Ok(FoundCrate::Itself) => Ok(quote! { use crate as _npos; }),
		Ok(FoundCrate::Name(sp_npos_elections)) => {
			let ident = syn::Ident::new(&sp_npos_elections, Span::call_site());
			Ok(quote!( extern crate #ident as _npos; ))
		},
		Err(e) => Err(syn::Error::new(Span::call_site(), e)),
	}
}

#[cfg(test)]
mod tests {
	#[test]
	fn ui_fail() {
		let cases = trybuild::TestCases::new();
		cases.compile_fail("tests/ui/fail/*.rs");
	}
}