contracts-imtbl/contracts/activity/NFTClaimStage2WL.sol

// SPDX-License-Identifier: MIT
pragma solidity 0.8.19;

import {ReentrancyGuard} from "@openzeppelin/contracts/security/ReentrancyGuard.sol";
import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import {AccessControl} from "@openzeppelin/contracts/access/AccessControl.sol";
import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
/**
 * Contract for the activity of NFT claim stage 2.
 */
interface IClaimAbleNFT {
  function safeMint(address to, uint256 tokenID) external;
}

contract NFTClaimStage2WL is ReentrancyGuard, AccessControl {
  using EnumerableSet for EnumerableSet.UintSet;
  using SafeERC20 for IERC20;
  /// @notice Only UPDATE_WL_ROLE can add white listing
  bytes32 public constant UPDATE_WL_ROLE = bytes32("UPDATE_WL_ROLE");
  /// @notice Only MANAGE_ROLE can change mint config
  bytes32 public constant MANAGE_ROLE = keccak256("MANAGE_ROLE");

  struct MintConfig {
    uint256 maxSupply; // max supply for phase2
    address currency; // token address which user must pay to mint
    uint256 mintPrice; // in wei
    address feeToAddress; // wallet address to receive mint fee
    uint256 airdropCount; // airdrop count
  }
  // parse: 0: not open or end, 1: phase1, 2: phase2
  uint256 public mintParse = 0;
  address public immutable nftAddress;
  uint256 public immutable nftIdStart;

  MintConfig public mintConfig;
  uint256 public totalCount;
  mapping(address user => uint256 num) private _whitelist1;
  mapping(address user => uint256 num) private _whitelist2;
  mapping(address user => EnumerableSet.UintSet tokenIdSet) private _mintedRecords;
  

  event NFTClaimed(address indexed nftAddress, address indexed to, uint256[] ids);

  event ParseUpdated(uint256 _parse);
  event MintConfigUpdated(MintConfig config);

  constructor(address _nftAddress, uint256 _nftIdStart, MintConfig memory _mintConfig) {
    _grantRole(DEFAULT_ADMIN_ROLE, _msgSender());
    _grantRole(UPDATE_WL_ROLE, _msgSender());
    _grantRole(MANAGE_ROLE, _msgSender());
    nftAddress = _nftAddress;
    mintConfig = _mintConfig;
    nftIdStart = _nftIdStart;
  }

  modifier whenNotPaused() {
    require(mintParse > 0, "NFTClaimer: not begin or ended");
    _;
  }

  function updateMintParse(uint256 _mintParse) external onlyRole(MANAGE_ROLE) {
    require(_mintParse == 0 || _mintParse == 1 || _mintParse == 2, "NFTClaimer: invalid mintParse");
    mintParse = _mintParse;
    emit ParseUpdated(_mintParse);
  }

  function updateMintConfig(MintConfig calldata config) external onlyRole(MANAGE_ROLE) {
    mintConfig = config;
    emit MintConfigUpdated(config);
  }

  function addParse1WL(address[] calldata _addressList, uint256[] calldata _nums) external onlyRole(UPDATE_WL_ROLE) {
    require(_addressList.length == _nums.length, "NFTClaimer: invalid whitelist");
    for (uint256 i = 0; i < _addressList.length; i++) {
      _whitelist1[_addressList[i]] = _nums[i];
    }
  }

  function revokeParse1WL(address[] calldata _addressList) external onlyRole(MANAGE_ROLE) {
    for (uint256 i = 0; i < _addressList.length; i++) {
      delete _whitelist1[_addressList[i]];
    }
  }

  function addParse2WL(address[] calldata _addressList) external onlyRole(UPDATE_WL_ROLE){
    for (uint256 i = 0; i < _addressList.length; i++) {
      _whitelist2[_addressList[i]] = 1;
    }
  }

  function revokeParse2WL(address[] calldata _addressList) external onlyRole(MANAGE_ROLE) {
    for (uint256 i = 0; i < _addressList.length; i++) {
      delete _whitelist2[_addressList[i]];
    }
  }


  /**
   * @dev claim NFT
   * @param nftCount nft count to claim
   */
  function claim(
    uint256 nftCount
  ) external nonReentrant whenNotPaused {
    require(nftCount > 0, "NFTClaimer: nft count must be greater than 0");
    require(nftCount <= mintConfig.maxSupply - mintConfig.airdropCount - totalCount, "NFTClaimer: exceed max supply");
    address to = _msgSender();
    uint256 _mintedCount = _mintedRecords[to].length();
    if (mintParse == 1) {
      require(_whitelist1[to] >= _mintedCount + nftCount, "NFTClaimer: not in whitelist or exceed limit");
    } else if (mintParse == 2) {
      require(_whitelist1[to] + _whitelist2[to] >= _mintedCount + nftCount, "NFTClaimer: not in whitelist or exceed limit");
    }
    uint256 _tokenAmount = mintConfig.mintPrice * nftCount;
    totalCount += nftCount;
    IERC20(mintConfig.currency).safeTransferFrom(to, mintConfig.feeToAddress, _tokenAmount);
    uint256[] memory ids = new uint256[](nftCount);
    for (uint256 i = 0; i < nftCount; ++i) {
      uint256 _nftId = nftIdStart + totalCount + i;
      ids[i] = _nftId;
      _mintedRecords[to].add(_nftId);
      IClaimAbleNFT(nftAddress).safeMint(to, _nftId);
    }
    emit NFTClaimed(nftAddress, to, ids);
  }

  function queryInfo() external view returns (uint256 num1, uint256 num2, uint256 minted) {
    num1 = _whitelist1[_msgSender()];
    num2 = _whitelist2[_msgSender()];
    minted = _mintedRecords[_msgSender()].length();
    return (num1, num2, minted);
  }


  function mintedNum() external view returns (uint256){
    return _mintedRecords[_msgSender()].length();
  }

  function mintedNft() external view returns (uint256[] memory){
    return _mintedRecords[_msgSender()].values();
  }

}