pragma circom 2.0.0;

include "../node_modules/circomlib/circuits/poseidon.circom";

template PasswordCheck() {
  // Public inputs
  signal input userAddress;
  signal input pwdHash;
  // Private input
  signal input pwd;

  // (Public) output
  signal output uniqueToUser;

  // Make sure password is the correct one by comparing its hash to the expected known hash
  component hasher = Poseidon(1);
  hasher.inputs[0] <== pwd;

  hasher.out === pwdHash;

  // Compute a number unique to user so that other users can't simply copy and use same proof
  // but instead have to execute this circuit to generate a proof unique to them
  component uniqueHasher = Poseidon(2);
  uniqueHasher.inputs[0] <== pwdHash;
  uniqueHasher.inputs[1] <== userAddress;

  uniqueToUser <== uniqueHasher.out;
}

component main {public [userAddress, pwdHash]} = PasswordCheck();

use super::groth16_verifier_constants::{N_PUBLIC_INPUTS, vk, ic, precomputed_lines};

#[starknet::interface]
trait IGroth16VerifierBN254<TContractState> {
    fn verify_groth16_proof_bn254(
        self: @TContractState, full_proof_with_hints: Span<felt252>,
    ) -> Option<Span<u256>>;
}

#[starknet::contract]
mod Groth16VerifierBN254 {
    use starknet::SyscallResultTrait;
    use garaga::definitions::{G1Point, G1G2Pair};
    use garaga::groth16::{multi_pairing_check_bn254_3P_2F_with_extra_miller_loop_result};
    use garaga::ec_ops::{G1PointTrait, ec_safe_add};
    use garaga::ec_ops_g2::{G2PointTrait};
    use garaga::utils::calldata::{deserialize_full_proof_with_hints_bn254};
    use super::{N_PUBLIC_INPUTS, vk, ic, precomputed_lines};

    const ECIP_OPS_CLASS_HASH: felt252 =
        0x70c1d1c709c75e3cf51d79d19cf7c84a0d4521f3a2b8bf7bff5cb45ee0dd289;
    use starknet::ContractAddress;

    #[storage]
    struct Storage {}

    #[abi(embed_v0)]
    impl IGroth16VerifierBN254 of super::IGroth16VerifierBN254<ContractState> {
        fn verify_groth16_proof_bn254(
            self: @ContractState, full_proof_with_hints: Span<felt252>,
        ) -> Option<Span<u256>> {
            // DO NOT EDIT THIS FUNCTION UNLESS YOU KNOW WHAT YOU ARE DOING.
            // This function returns an Option for the public inputs if the proof is valid.
            // If the proof is invalid, the execution will either fail or return None.
            // Read the documentation to learn how to generate the full_proof_with_hints array given
            // a proof and a verifying key.
            let fph = deserialize_full_proof_with_hints_bn254(full_proof_with_hints);
            let groth16_proof = fph.groth16_proof;
            let mpcheck_hint = fph.mpcheck_hint;
            let small_Q = fph.small_Q;
            let msm_hint = fph.msm_hint;

            groth16_proof.a.assert_on_curve(0);
            groth16_proof.b.assert_on_curve(0);
            groth16_proof.c.assert_on_curve(0);

            let ic = ic.span();

            let vk_x: G1Point = match ic.len() {
                0 => panic!("Malformed VK"),
                1 => *ic.at(0),
                _ => {
                    // Start serialization with the hint array directly to avoid copying it.
                    let mut msm_calldata: Array<felt252> = msm_hint;
                    // Add the points from VK and public inputs to the proof.
                    Serde::serialize(@ic.slice(1, N_PUBLIC_INPUTS), ref msm_calldata);
                    Serde::serialize(@groth16_proof.public_inputs, ref msm_calldata);
                    // Complete with the curve identifier (0 for BN254):
                    msm_calldata.append(0);

                    // Call the multi scalar multiplication endpoint on the Garaga ECIP ops contract
                    // to obtain vk_x.
                    let mut _vx_x_serialized = core::starknet::syscalls::library_call_syscall(
                        ECIP_OPS_CLASS_HASH.try_into().unwrap(),
                        selector!("msm_g1"),
                        msm_calldata.span(),
                    )
                        .unwrap_syscall();

                    ec_safe_add(
                        Serde::<G1Point>::deserialize(ref _vx_x_serialized).unwrap(), *ic.at(0), 0,
                    )
                },
            };
            // Perform the pairing check.
            let check = multi_pairing_check_bn254_3P_2F_with_extra_miller_loop_result(
                G1G2Pair { p: vk_x, q: vk.gamma_g2 },
                G1G2Pair { p: groth16_proof.c, q: vk.delta_g2 },
                G1G2Pair { p: groth16_proof.a.negate(0), q: groth16_proof.b },
                vk.alpha_beta_miller_loop_result,
                precomputed_lines.span(),
                mpcheck_hint,
                small_Q,
            );
            if check == true {
                return Option::Some(groth16_proof.public_inputs);
            } else {
                return Option::None;
            }
        }
    }
}

use starknet::ContractAddress;

#[starknet::interface]
trait IZkERC20Token<TContractState> {
    fn mint_with_proof(ref self: TContractState, full_proof: Span<felt252>);
    fn has_user_minted(self: @TContractState, address: ContractAddress) -> bool;
}

#[starknet::interface]
trait IGroth16VerifierBN254<TContractState> {
    fn verify_groth16_proof_bn254(
        self: @TContractState, full_proof_with_hints: Span<felt252>,
    ) -> Option<Span<u256>>;
}

mod errors {
    pub const ALREADY_MINTED: felt252 = 'User has already minted tokens';
    pub const PROOF_NOT_VERIFIED: felt252 = 'Proof is not correct';
    pub const PROOF_ALREADY_USED: felt252 = 'Generate a proof unique to you';
}

#[starknet::contract]
pub mod ZkERC20Token {
    use openzeppelin_token::erc20::{ERC20Component, ERC20HooksEmptyImpl};
    use starknet::{ContractAddress, get_caller_address};
    use super::{errors, IGroth16VerifierBN254Dispatcher, IGroth16VerifierBN254DispatcherTrait};
    use starknet::storage::{
        Map, StoragePointerReadAccess, StoragePointerWriteAccess, StoragePathEntry,
    };

    component!(path: ERC20Component, storage: erc20, event: ERC20Event);

    #[abi(embed_v0)]
    impl ERC20MixinImpl = ERC20Component::ERC20MixinImpl<ContractState>;

    impl ERC20InternalImpl = ERC20Component::InternalImpl<ContractState>;

    const MINT_WITH_PROOF_TOKEN_REWARD: u8 = 100;
    // used in the front end to generate the proof
    const PASSWORD_HASH: u256 =
        16260938803047823847354854419633652218467975114284208787981985448019235110758;

    #[storage]
    struct Storage {
        #[substorage(v0)]
        erc20: ERC20Component::Storage,
        verifier_contract: IGroth16VerifierBN254Dispatcher,
        users_who_minted: Map<ContractAddress, bool>,
    }

    #[event]
    #[derive(Drop, starknet::Event)]
    enum Event {
        #[flat]
        ERC20Event: ERC20Component::Event,
    }

    #[constructor]
    fn constructor(
        ref self: ContractState,
        initial_supply: u256,
        recipient: ContractAddress,
        name: ByteArray,
        symbol: ByteArray,
        proof_verifier_address: ContractAddress,
    ) {
        self.erc20.initializer(name, symbol);
        self.erc20.mint(recipient, initial_supply);

        self
            .verifier_contract
            .write(IGroth16VerifierBN254Dispatcher { contract_address: proof_verifier_address });
    }

    #[abi(embed_v0)]
    impl ZkERC20TokenImpl of super::IZkERC20Token<ContractState> {
        fn mint_with_proof(ref self: ContractState, full_proof: Span<felt252>) {
            let caller = get_caller_address();
            // Prevent a user from receiving tokens twice
            assert(!self.users_who_minted.entry(caller).read(), errors::ALREADY_MINTED);

            // Verify the correctness of the proof by calling the verifier contract
            // If incorrect, execution of the verifier will fail or return an Option::None
            let proof_public_inputs = self
                .verifier_contract
                .read()
                .verify_groth16_proof_bn254(full_proof);
            assert(
                proof_public_inputs.is_some() && proof_public_inputs.unwrap().len() == 3,
                errors::PROOF_NOT_VERIFIED,
            );

            // Verify the proof has been generated by the user calling this smart contract
            let user_address_dec: u256 = *proof_public_inputs.unwrap().at(1);
            let address_felt252: felt252 = caller.into();
            assert(address_felt252.into() == user_address_dec, errors::PROOF_ALREADY_USED);

            // Mint tokens only if the proof is valid and has been generated by the user
            self.erc20.mint(caller, MINT_WITH_PROOF_TOKEN_REWARD.into());

            self.users_who_minted.entry(caller).write(true);
        }

        fn has_user_minted(self: @TContractState, address: ContractAddress) -> bool {
            self.users_who_minted.entry(address).read()
        }
    }
}