//! A module for generating the code produced by this macro.  This code
//! will interact with the underlying `sigma` macro.

use super::sigma::codegen::{StructField, StructFieldList};
use super::syntax::*;
use proc_macro2::TokenStream;
use quote::{format_ident, quote};
#[cfg(test)]
use syn::parse_quote;
use syn::Ident;

/// The main struct to handle code generation for this macro.
///
/// Initialize a [`CodeGen`] with the [`SigmaCompSpec`] you get by
/// parsing the macro input.  Pass it to the various transformations and
/// statement handlers, which will both update the code it will
/// generate, and modify the [`SigmaCompSpec`].  Then at the end, call
/// [`CodeGen::generate`] with the modified [`SigmaCompSpec`] to generate the
/// code output by this macro.
pub struct CodeGen {
    /// The protocol name specified in the `sigma_compiler` macro
    /// invocation
    proto_name: Ident,
    /// The group name specified in the `sigma_compiler` macro
    /// invocation
    group_name: Ident,
    /// The variables that were explicitly listed in the
    /// `sigma_compiler` macro invocation
    vars: TaggedVarDict,
    /// A prefix that does not appear at the beginning of any variable
    /// name in `vars`
    unique_prefix: String,
    /// Variables (not necessarily appearing in `vars`, since they may
    /// be generated by the sigma_compiler itself) that the prover needs
    /// to send to the verifier along with the proof.  These could
    /// include commitments to bits in range proofs, for example.
    sent_instance: StructFieldList,
    /// Extra code that will be emitted in the `prove` function
    prove_code: TokenStream,
    /// Extra code that will be emitted in the `verify` function
    verify_code: TokenStream,
    /// Extra code that will be emitted in the `verify` function before
    /// the `sent_instance` are deserialized.  This is where the verifier
    /// sets the lengths of vector variables in the `sent_instance`.
    verify_pre_instance_code: TokenStream,
}

impl CodeGen {
    /// Find a prefix that does not appear at the beginning of any
    /// variable name in `vars`
    fn unique_prefix(vars: &TaggedVarDict) -> String {
        'outer: for tag in 0usize.. {
            let try_prefix = if tag == 0 {
                "gen__".to_string()
            } else {
                format!("gen{}__", tag)
            };
            for v in vars.keys() {
                if v.starts_with(&try_prefix) {
                    continue 'outer;
                }
            }
            return try_prefix;
        }
        // The compiler complains if this isn't here, but it will only
        // get hit if vars contains at least usize::MAX entries, which
        // isn't going to happen.
        String::new()
    }

    /// Create a new [`CodeGen`] given the [`SigmaCompSpec`] you get by
    /// parsing the macro input.
    pub fn new(spec: &SigmaCompSpec) -> Self {
        Self {
            proto_name: spec.proto_name.clone(),
            group_name: spec.group_name.clone(),
            vars: spec.vars.clone(),
            unique_prefix: Self::unique_prefix(&spec.vars),
            sent_instance: StructFieldList::default(),
            prove_code: quote! {},
            verify_code: quote! {},
            verify_pre_instance_code: quote! {},
        }
    }

    #[cfg(test)]
    /// Create an empty [`CodeGen`].  Primarily useful in testing.
    pub fn new_empty() -> Self {
        Self {
            proto_name: parse_quote! { proto },
            group_name: parse_quote! { G },
            vars: TaggedVarDict::default(),
            unique_prefix: "gen__".into(),
            sent_instance: StructFieldList::default(),
            prove_code: quote! {},
            verify_code: quote! {},
            verify_pre_instance_code: quote! {},
        }
    }

    /// Create a new generated private Scalar variable to put in the
    /// Witness.
    ///
    /// If you call this, you should also call
    /// [`prove_append`](Self::prove_append) with code like `quote!{ let
    /// #id = ... }` where `id` is the [`struct@Ident`] returned from
    /// this function.
    pub fn gen_scalar(
        &self,
        vars: &mut TaggedVarDict,
        base: &Ident,
        is_rand: bool,
        is_vec: bool,
    ) -> Ident {
        let id = format_ident!("{}{}", self.unique_prefix, base);
        vars.insert(
            id.to_string(),
            TaggedIdent::Scalar(TaggedScalar {
                id: id.clone(),
                is_pub: false,
                is_rand,
                is_vec,
            }),
        );
        id
    }

    /// Create a new public Point variable to put in the Instance,
    /// optionally marking it as needing to be sent from the prover to
    /// the verifier along with the proof.
    ///
    /// If you call this function, you should also call
    /// [`prove_append`](Self::prove_append) with code like `quote!{ let
    /// #id = ... }` where `id` is the [`struct@Ident`] returned from
    /// this function.  If `is_vec` is `true`, then you should also call
    /// [`verify_pre_instance_append`](Self::verify_pre_instance_append)
    /// with code like `quote!{ let mut #id = Vec::<Point>::new();
    /// #id.resize(#len, Point::default()); }` where `len` is the number
    /// of elements you expect to have in the vector (computed at
    /// runtime, perhaps based on the values of public parameters).
    pub fn gen_point(
        &mut self,
        vars: &mut TaggedVarDict,
        base: &Ident,
        is_vec: bool,
        send_to_verifier: bool,
    ) -> Ident {
        let id = format_ident!("{}{}", self.unique_prefix, base);
        vars.insert(
            id.to_string(),
            TaggedIdent::Point(TaggedPoint {
                id: id.clone(),
                is_cind: false,
                is_const: false,
                is_vec,
            }),
        );
        if send_to_verifier {
            if is_vec {
                self.sent_instance.push_vecpoint(&id);
            } else {
                self.sent_instance.push_point(&id);
            }
        }
        id
    }

    /// Create a new identifier, using the unique prefix
    pub fn gen_ident(&self, base: &Ident) -> Ident {
        format_ident!("{}{}", self.unique_prefix, base)
    }

    /// Append some code to the generated `prove` function
    pub fn prove_append(&mut self, code: TokenStream) {
        let prove_code = &self.prove_code;
        self.prove_code = quote! {
            #prove_code
            #code
        };
    }

    /// Append some code to the generated `verify` function
    pub fn verify_append(&mut self, code: TokenStream) {
        let verify_code = &self.verify_code;
        self.verify_code = quote! {
            #verify_code
            #code
        };
    }

    /// Append some code to the generated `verify` function to be run
    /// before the `sent_instance` are deserialized
    pub fn verify_pre_instance_append(&mut self, code: TokenStream) {
        let verify_pre_instance_code = &self.verify_pre_instance_code;
        self.verify_pre_instance_code = quote! {
            #verify_pre_instance_code
            #code
        };
    }

    /// Append some code to both the generated `prove` and `verify`
    /// functions
    pub fn prove_verify_append(&mut self, code: TokenStream) {
        let prove_code = &self.prove_code;
        self.prove_code = quote! {
            #prove_code
            #code
        };
        let verify_code = &self.verify_code;
        self.verify_code = quote! {
            #verify_code
            #code
        };
    }

    /// Append some code to both the generated `prove` and `verify`
    /// functions, the latter to be run before the `sent_instance` are
    /// deserialized
    pub fn prove_verify_pre_instance_append(&mut self, code: TokenStream) {
        let prove_code = &self.prove_code;
        self.prove_code = quote! {
            #prove_code
            #code
        };
        let verify_pre_instance_code = &self.verify_pre_instance_code;
        self.verify_pre_instance_code = quote! {
            #verify_pre_instance_code
            #code
        };
    }

    /// Extract (as [`String`]s) the code inserted by
    /// [`prove_append`](Self::prove_append),
    /// [`verify_append`](Self::verify_append), and
    /// [`verify_pre_instance_append`](Self::verify_pre_instance_append).
    pub fn code_strings(&self) -> (String, String, String) {
        (
            self.prove_code.to_string(),
            self.verify_code.to_string(),
            self.verify_pre_instance_code.to_string(),
        )
    }

    /// Generate the code to be output by this macro.
    ///
    /// `emit_prover` and `emit_verifier` are as in
    /// [`sigma_compiler_core`](super::sigma_compiler_core).
    pub fn generate(
        &self,
        spec: &mut SigmaCompSpec,
        emit_prover: bool,
        emit_verifier: bool,
    ) -> TokenStream {
        let proto_name = &self.proto_name;
        let group_name = &self.group_name;

        let group_types = quote! {
            use super::group;
            pub type Scalar = <super::#group_name as group::Group>::Scalar;
            pub type Point = super::#group_name;
        };

        // vardict contains the variables that were defined in the macro
        // call to [`sigma_compiler`]
        let vardict = taggedvardict_to_vardict(&self.vars);
        // sigma_proofs_vardict contains the variables that we are passing
        // to sigma_proofs.  We may have removed some via substitution, and
        // we may have added some when compiling statements like range
        // assertions into underlying linear combination assertions.
        let sigma_proofs_vardict = taggedvardict_to_vardict(&spec.vars);

        // Generate the code that uses the underlying sigma_proofs API
        let mut sigma_proofs_codegen = super::sigma::codegen::CodeGen::new(
            format_ident!("sigma"),
            format_ident!("Point"),
            &sigma_proofs_vardict,
            &mut spec.statements,
        );
        let sigma_proofs_code = sigma_proofs_codegen.generate(emit_prover, emit_verifier);

        let mut pub_instance_fields = StructFieldList::default();
        pub_instance_fields.push_vars(&vardict, true);
        let mut witness_fields = StructFieldList::default();
        witness_fields.push_vars(&vardict, false);

        let mut sigma_proofs_instance_fields = StructFieldList::default();
        sigma_proofs_instance_fields.push_vars(&sigma_proofs_vardict, true);
        let mut sigma_proofs_witness_fields = StructFieldList::default();
        sigma_proofs_witness_fields.push_vars(&sigma_proofs_vardict, false);

        // Generate the public instance struct definition
        let instance_def = {
            let decls = pub_instance_fields.field_decls();
            #[cfg(feature = "dump")]
            let dump_impl = {
                let dump_chunks = pub_instance_fields.dump();
                quote! {
                    impl Instance {
                        fn dump_scalar(s: &Scalar) {
                            let bytes: &[u8] = &s.to_repr();
                            print!("{:02x?}", bytes);
                        }

                        fn dump_point(p: &Point) {
                            let bytes: &[u8] = &p.to_bytes();
                            print!("{:02x?}", bytes);
                        }

                        pub fn dump(&self) {
                            #dump_chunks
                        }
                    }
                }
            };
            #[cfg(not(feature = "dump"))]
            let dump_impl = {
                quote! {}
            };
            quote! {
                #[derive(Clone)]
                pub struct Instance {
                    #decls
                }

                #dump_impl
            }
        };

        // Generate the witness struct definition
        let witness_def = if emit_prover {
            let decls = witness_fields.field_decls();
            quote! {
                #[derive(Clone)]
                pub struct Witness {
                    #decls
                }
            }
        } else {
            quote! {}
        };

        // Generate the prove function
        let prove_func = if emit_prover {
            let instance_ids = pub_instance_fields.field_list();
            let witness_ids = witness_fields.field_list();
            let sigma_proofs_instance_ids = sigma_proofs_instance_fields.field_list();
            let sigma_proofs_witness_ids = sigma_proofs_witness_fields.field_list();
            let prove_code = &self.prove_code;
            let codegen_instance_var = format_ident!("{}sigma_instance", self.unique_prefix);
            let codegen_witness_var = format_ident!("{}sigma_witness", self.unique_prefix);
            let instance_var = format_ident!("{}instance", self.unique_prefix);
            let witness_var = format_ident!("{}witness", self.unique_prefix);
            let rng_var = format_ident!("{}rng", self.unique_prefix);
            let proof_var = format_ident!("{}proof", self.unique_prefix);
            let sid_var = format_ident!("{}session_id", self.unique_prefix);
            let sent_instance_code = {
                let chunks = self.sent_instance.fields.iter().map(|sf| match sf {
                    StructField::Point(id) => quote! {
                        #proof_var.extend(sigma_proofs::serialization::serialize_elements(
                            std::slice::from_ref(&#codegen_instance_var.#id)
                        ));
                    },
                    StructField::VecPoint(id) => quote! {
                        #proof_var.extend(sigma_proofs::serialization::serialize_elements(
                            &#codegen_instance_var.#id
                        ));
                    },
                    _ => quote! {},
                });
                quote! { #(#chunks)* }
            };

            let dumper = if cfg!(feature = "dump") {
                quote! {
                    println!("prover instance = {{");
                    #instance_var.dump();
                    println!("}}");
                }
            } else {
                quote! {}
            };

            quote! {
                pub fn prove(
                    #instance_var: &Instance,
                    #witness_var: &Witness,
                    #sid_var: &[u8],
                    #rng_var: &mut (impl CryptoRng + RngCore),
                ) -> Result<Vec<u8>, SigmaError> {
                    #dumper
                    let Instance { #instance_ids } = #instance_var.clone();
                    let Witness { #witness_ids } = #witness_var.clone();
                    #prove_code
                    let mut #proof_var = Vec::<u8>::new();
                    let #codegen_instance_var = sigma::Instance {
                        #sigma_proofs_instance_ids
                    };
                    let #codegen_witness_var = sigma::Witness {
                        #sigma_proofs_witness_ids
                    };
                    #sent_instance_code
                    #proof_var.extend(
                        sigma::prove(
                            &#codegen_instance_var,
                            &#codegen_witness_var,
                            #sid_var,
                            #rng_var,
                        )?
                    );
                    Ok(#proof_var)
                }
            }
        } else {
            quote! {}
        };

        // Generate the verify function
        let verify_func = if emit_verifier {
            let instance_ids = pub_instance_fields.field_list();
            let sigma_proofs_instance_ids = sigma_proofs_instance_fields.field_list();
            let verify_pre_instance_code = &self.verify_pre_instance_code;
            let verify_code = &self.verify_code;
            let codegen_instance_var = format_ident!("{}sigma_instance", self.unique_prefix);
            let element_len_var = format_ident!("{}element_len", self.unique_prefix);
            let offset_var = format_ident!("{}proof_offset", self.unique_prefix);
            let instance_var = format_ident!("{}instance", self.unique_prefix);
            let proof_var = format_ident!("{}proof", self.unique_prefix);
            let sid_var = format_ident!("{}session_id", self.unique_prefix);
            let sent_instance_code = {
                let element_len_code = if self.sent_instance.fields.is_empty() {
                    quote! {}
                } else {
                    quote! {
                        let #element_len_var =
                            <Point as group::GroupEncoding>::Repr::default().as_ref().len();
                    }
                };

                let chunks = self.sent_instance.fields.iter().map(|sf| match sf {
                    StructField::Point(id) => quote! {
                        let #id: Point = sigma_proofs::serialization::deserialize_elements(
                                &#proof_var[#offset_var..],
                                1,
                            ).ok_or(SigmaError::VerificationFailure)?[0];
                        #offset_var += #element_len_var;
                    },
                    StructField::VecPoint(id) => quote! {
                        #id = sigma_proofs::serialization::deserialize_elements(
                                &#proof_var[#offset_var..],
                                #id.len(),
                            ).ok_or(SigmaError::VerificationFailure)?;
                        #offset_var += #element_len_var * #id.len();
                    },
                    _ => quote! {},
                });

                quote! {
                    let mut #offset_var = 0usize;
                    #element_len_code
                    #(#chunks)*
                }
            };

            let dumper = if cfg!(feature = "dump") {
                quote! {
                    println!("verifier instance = {{");
                    #instance_var.dump();
                    println!("}}");
                }
            } else {
                quote! {}
            };

            quote! {
                pub fn verify(
                    #instance_var: &Instance,
                    #proof_var: &[u8],
                    #sid_var: &[u8],
                ) -> Result<(), SigmaError> {
                    #dumper
                    let Instance { #instance_ids } = #instance_var.clone();
                    #verify_pre_instance_code
                    #sent_instance_code
                    #verify_code
                    let #codegen_instance_var = sigma::Instance {
                        #sigma_proofs_instance_ids
                    };
                    sigma::verify(
                        &#codegen_instance_var,
                        &#proof_var[#offset_var..],
                        #sid_var,
                    )
                }
            }
        } else {
            quote! {}
        };

        // Output the generated module for this protocol
        let dump_use = if cfg!(feature = "dump") {
            quote! {
                use group::GroupEncoding;
            }
        } else {
            quote! {}
        };
        quote! {
            #[allow(non_snake_case)]
            pub mod #proto_name {
                use sigma_compiler::group::Group;
                use sigma_compiler::group::ff::{Field, PrimeField};
                use sigma_compiler::group::ff::derive::subtle::ConditionallySelectable;
                use sigma_compiler::rand::{CryptoRng, RngCore};
                use sigma_compiler::sigma_proofs;
                use sigma_compiler::sigma_proofs::errors::Error as SigmaError;
                use sigma_compiler::vecutils::*;
                use std::ops::Neg;
                #dump_use

                #group_types

                #sigma_proofs_code

                #instance_def
                #witness_def
                #prove_func
                #verify_func
            }
        }
    }
}