// SPDX-License-Identifier: MIT
pragma solidity ^0.8.26;

interface IERC20 {
    function balanceOf(address) external view returns (uint256);
    function transfer(address, uint256) external returns (bool);
    function transferFrom(address, address, uint256) external returns (bool);
}

interface IHasher {
    function hash(uint256[2] memory) external pure returns (uint256);
}

interface IVerifier {
    function verifyProof(uint256[2] calldata, uint256[2][2] calldata, uint256[2] calldata, uint256[7] calldata)
        external
        view
        returns (bool);
}

/// @title Lethe Pool
/// @notice A shielded pool for one asset (USDC, or native HYPE when `token` is zero). A shield adds
/// a private note to a Merkle tree; a withdrawal proves with a Groth16 proof that it spends some
/// note in the tree, without saying which. Notes are amounts of the asset (Tornado Nova design,
/// join-split 2 in / 2 out).
///
/// A withdrawal can be sent by a relayer, who is paid `fee` in the asset and may forward `refund`
/// HYPE to the recipient so a new address can pay its own gas. Both are bound into the proof, so a
/// relayer cannot change them.
///
/// Everything that leaves the pool pays `exitFeeBps`, in the pool asset. It is accrued here and
/// sent to `feeRecipient` by `collectFees()`, so a withdrawal never depends on that address.
///
/// The owner can pause new shields, change the two deposit limits, set the exit fee under the hard
/// cap and choose where the fee goes. The owner has no path to the notes: withdrawals and private
/// transfers cannot be paused, and there is no upgrade path.
contract LethePool {
    uint256 public constant FIELD_SIZE =
        21888242871839275222246405745257275088548364400416034343698204186575808495617;
    uint256 public constant MAX_EXT_AMOUNT = 2 ** 248;
    uint32 public constant LEVELS = 20;
    uint32 public constant ROOT_HISTORY_SIZE = 100;
    uint256 public constant MAX_ENCRYPTED_OUTPUT = 512;
    uint256 public constant ZERO_VALUE = uint256(keccak256("lethe pool")) % FIELD_SIZE;
    /// @notice Hard cap on the exit fee: 1%. No function can set it higher.
    uint256 public constant MAX_EXIT_FEE_BPS = 100;
    uint256 private constant BPS = 10_000;

    /// @notice The pool's asset. Zero means native HYPE.
    address public immutable token;
    IVerifier public immutable verifier;
    IHasher public immutable hasher;
    /// @notice Largest HYPE refund a relayer may forward with one withdrawal (USDC pool only).
    uint256 public immutable maxRefund;

    address public owner;
    bool public depositsPaused;
    /// @notice Largest single shield, in the asset's base units.
    uint256 public maxDeposit;
    /// @notice Most the notes in this pool may add up to. Only new shields are checked against it.
    uint256 public tvlCap;
    /// @notice Share of every exit taken as a fee, in basis points.
    uint256 public exitFeeBps;
    /// @notice Where `collectFees()` sends the accrued fee.
    address public feeRecipient;
    /// @notice Sum of all unspent notes, in base units.
    uint256 public totalShielded;
    /// @notice Exit fees taken and not collected yet.
    uint256 public accruedFees;

    uint256[LEVELS + 1] public zeros;
    uint256[LEVELS] public filledSubtrees;
    uint256[ROOT_HISTORY_SIZE] public roots;
    uint32 public currentRootIndex;
    uint32 public nextIndex;
    mapping(uint256 => bool) public nullifierHashes;

    uint256 private _lock = 1;

    struct ExtData {
        address recipient;
        int256 extAmount;
        address relayer;
        uint256 fee;
        uint256 refund;
        bytes encryptedOutput1;
        bytes encryptedOutput2;
    }

    struct Proof {
        uint256[2] a;
        uint256[2][2] b;
        uint256[2] c;
        uint256 root;
        uint256 publicAmount;
        uint256 extDataHash;
        uint256[2] inputNullifiers;
        uint256[2] outputCommitments;
    }

    event NewCommitment(uint256 commitment, uint256 index, bytes encryptedOutput);
    event NewNullifier(uint256 nullifier);
    event Shield(address indexed sender, uint256 amount);
    /// @param amount What left the pool toward `recipient`, before the exit fee.
    /// @param exitFee The fee taken from `amount` and from the relayer's `fee`.
    event Unshield(
        address indexed recipient, uint256 amount, address indexed relayer, uint256 fee, uint256 refund, uint256 exitFee
    );
    event FeesCollected(address indexed recipient, uint256 amount);
    event DepositsPaused(bool paused);
    event MaxDepositSet(uint256 maxDeposit);
    event TvlCapSet(uint256 tvlCap);
    event ExitFeeSet(uint256 exitFeeBps);
    event FeeRecipientSet(address feeRecipient);
    event OwnershipTransferred(address indexed previousOwner, address indexed newOwner);

    error NotOwner();
    error Reentrancy();
    error UnknownRoot();
    error NullifierSpent();
    error BadExtData();
    error BadFee();
    error BadRefund();
    error BadValue();
    error BadPublicAmount();
    error BadExtAmount();
    error InvalidProof();
    error Paused();
    error DepositTooLarge();
    error TvlCapReached();
    error ExitFeeTooHigh();
    error ZeroAddress();
    error NothingToCollect();
    error ZeroRecipient();
    error TransferFailed();
    error TreeFull();
    error OutputTooLong();

    modifier onlyOwner() {
        if (msg.sender != owner) revert NotOwner();
        _;
    }

    modifier nonReentrant() {
        if (_lock != 1) revert Reentrancy();
        _lock = 2;
        _;
        _lock = 1;
    }

    constructor(
        address _token,
        IVerifier _verifier,
        IHasher _hasher,
        address _owner,
        uint256 _maxDeposit,
        uint256 _tvlCap,
        uint256 _maxRefund,
        uint256 _exitFeeBps,
        address _feeRecipient
    ) {
        if (_exitFeeBps > MAX_EXIT_FEE_BPS) revert ExitFeeTooHigh();
        if (_feeRecipient == address(0) || _owner == address(0)) revert ZeroAddress();
        token = _token;
        verifier = _verifier;
        hasher = _hasher;
        maxDeposit = _maxDeposit;
        tvlCap = _tvlCap;
        maxRefund = _token == address(0) ? 0 : _maxRefund;
        exitFeeBps = _exitFeeBps;
        feeRecipient = _feeRecipient;
        owner = _owner;
        emit OwnershipTransferred(address(0), _owner);

        uint256 z = ZERO_VALUE;
        zeros[0] = z;
        for (uint32 i = 0; i < LEVELS; i++) {
            filledSubtrees[i] = z;
            z = _hash(z, z);
            zeros[i + 1] = z;
        }
        roots[0] = z;
    }

    // ------------------------------------------------------------------ transact

    /// @notice Shield (extAmount > 0), withdraw (< 0) or move privately (0). Every call proves.
    function transact(Proof calldata p, ExtData calldata e) external payable nonReentrant {
        if (!isKnownRoot(p.root)) revert UnknownRoot();
        if (nullifierHashes[p.inputNullifiers[0]] || nullifierHashes[p.inputNullifiers[1]]) revert NullifierSpent();
        if (uint256(keccak256(abi.encode(e))) % FIELD_SIZE != p.extDataHash) revert BadExtData();
        if (e.fee >= MAX_EXT_AMOUNT || (e.fee > 0 && e.relayer == address(0))) revert BadFee();
        if (e.encryptedOutput1.length > MAX_ENCRYPTED_OUTPUT || e.encryptedOutput2.length > MAX_ENCRYPTED_OUTPUT) {
            revert OutputTooLong();
        }
        if (p.publicAmount != publicAmountOf(e.extAmount, e.fee)) revert BadPublicAmount();
        if (!_verify(p)) revert InvalidProof();

        nullifierHashes[p.inputNullifiers[0]] = true;
        nullifierHashes[p.inputNullifiers[1]] = true;
        emit NewNullifier(p.inputNullifiers[0]);
        emit NewNullifier(p.inputNullifiers[1]);
        uint32 index = _insert(p.outputCommitments[0], p.outputCommitments[1]);
        emit NewCommitment(p.outputCommitments[0], index, e.encryptedOutput1);
        emit NewCommitment(p.outputCommitments[1], index + 1, e.encryptedOutput2);

        _settle(e);
    }

    /// @notice Sends the accrued exit fees to `feeRecipient`. Anyone can call it.
    function collectFees() external nonReentrant {
        uint256 amount = accruedFees;
        if (amount == 0) revert NothingToCollect();
        accruedFees = 0;
        _push(feeRecipient, amount);
        emit FeesCollected(feeRecipient, amount);
    }

    // ------------------------------------------------------------------ views

    function publicAmountOf(int256 extAmount, uint256 fee) public pure returns (uint256) {
        if (extAmount <= -int256(MAX_EXT_AMOUNT) || extAmount >= int256(MAX_EXT_AMOUNT)) revert BadExtAmount();
        int256 amount = extAmount - int256(fee);
        return amount >= 0 ? uint256(amount) : FIELD_SIZE - uint256(-amount);
    }

    function isKnownRoot(uint256 root) public view returns (bool) {
        if (root == 0) return false;
        uint32 i = currentRootIndex;
        do {
            if (roots[i] == root) return true;
            i = i == 0 ? ROOT_HISTORY_SIZE - 1 : i - 1;
        } while (i != currentRootIndex);
        return false;
    }

    function getLastRoot() external view returns (uint256) {
        return roots[currentRootIndex];
    }

    // ------------------------------------------------------------------ owner

    /// @notice Stops or restarts new shields. Withdrawals and private transfers are not affected.
    function setDepositsPaused(bool paused) external onlyOwner {
        depositsPaused = paused;
        emit DepositsPaused(paused);
    }

    /// @notice Sets the largest single shield. Only new shields are checked against it.
    function setMaxDeposit(uint256 newMaxDeposit) external onlyOwner {
        maxDeposit = newMaxDeposit;
        emit MaxDepositSet(newMaxDeposit);
    }

    /// @notice Sets the most the pool may hold. Only new shields are checked against it.
    function setTvlCap(uint256 newTvlCap) external onlyOwner {
        tvlCap = newTvlCap;
        emit TvlCapSet(newTvlCap);
    }

    /// @notice Sets the exit fee, never above `MAX_EXIT_FEE_BPS`.
    function setExitFeeBps(uint256 newExitFeeBps) external onlyOwner {
        if (newExitFeeBps > MAX_EXIT_FEE_BPS) revert ExitFeeTooHigh();
        exitFeeBps = newExitFeeBps;
        emit ExitFeeSet(newExitFeeBps);
    }

    /// @notice Sets where `collectFees()` sends the accrued fee.
    function setFeeRecipient(address newFeeRecipient) external onlyOwner {
        if (newFeeRecipient == address(0)) revert ZeroAddress();
        feeRecipient = newFeeRecipient;
        emit FeeRecipientSet(newFeeRecipient);
    }

    function transferOwnership(address newOwner) external onlyOwner {
        if (newOwner == address(0)) revert ZeroAddress();
        emit OwnershipTransferred(owner, newOwner);
        owner = newOwner;
    }

    // ------------------------------------------------------------------ internals

    /// @dev Moves the asset for one transaction, after the tree is updated.
    function _settle(ExtData calldata e) internal {
        if (e.extAmount > 0) {
            if (depositsPaused) revert Paused();
            uint256 deposit = uint256(e.extAmount);
            if (deposit > maxDeposit) revert DepositTooLarge();
            if (totalShielded + deposit > tvlCap) revert TvlCapReached();
            // A shield is sent by its owner: there is no relayer to pay.
            if (e.fee != 0 || e.refund != 0) revert BadFee();
            totalShielded += deposit;
            if (token == address(0)) {
                if (msg.value != deposit) revert BadValue();
            } else {
                if (msg.value != 0) revert BadValue();
                _pull(msg.sender, deposit);
            }
            emit Shield(msg.sender, deposit);
            return;
        }

        if (e.refund > maxRefund) revert BadRefund();
        if (msg.value != e.refund) revert BadValue();
        uint256 amount;
        if (e.extAmount < 0) {
            if (e.recipient == address(0)) revert ZeroRecipient();
            amount = uint256(-e.extAmount);
        } else if (e.refund > 0) {
            revert BadRefund();
        }

        // Value leaves the pool two ways: to the recipient and to the relayer. Both pay the exit
        // fee, otherwise a user could name their own wallet as the relayer and exit for free.
        uint256 feeOnAmount = (amount * exitFeeBps) / BPS;
        uint256 feeOnRelayer = (e.fee * exitFeeBps) / BPS;
        totalShielded -= amount + e.fee;
        accruedFees += feeOnAmount + feeOnRelayer;

        if (amount > 0) {
            _push(e.recipient, amount - feeOnAmount);
            if (e.refund > 0) _pushHype(e.recipient, e.refund);
        }
        if (e.fee > 0) _push(e.relayer, e.fee - feeOnRelayer);
        if (e.extAmount < 0) emit Unshield(e.recipient, amount, e.relayer, e.fee, e.refund, feeOnAmount + feeOnRelayer);
    }

    function _verify(Proof calldata p) internal view returns (bool) {
        uint256[7] memory pub = [
            p.root,
            p.publicAmount,
            p.extDataHash,
            p.inputNullifiers[0],
            p.inputNullifiers[1],
            p.outputCommitments[0],
            p.outputCommitments[1]
        ];
        return verifier.verifyProof(p.a, p.b, p.c, pub);
    }

    function _insert(uint256 leaf1, uint256 leaf2) internal returns (uint32 index) {
        uint32 next = nextIndex;
        if (next == uint32(2) ** LEVELS) revert TreeFull();
        uint32 current = next / 2;
        uint256 node = _hash(leaf1, leaf2);
        for (uint32 i = 1; i < LEVELS; i++) {
            if (current % 2 == 0) {
                filledSubtrees[i] = node;
                node = _hash(node, zeros[i]);
            } else {
                node = _hash(filledSubtrees[i], node);
            }
            current /= 2;
        }
        uint32 r = (currentRootIndex + 1) % ROOT_HISTORY_SIZE;
        currentRootIndex = r;
        roots[r] = node;
        nextIndex = next + 2;
        return next;
    }

    function _hash(uint256 left, uint256 right) internal view returns (uint256) {
        return hasher.hash([left, right]);
    }

    function _pull(address from, uint256 amount) internal {
        uint256 before = IERC20(token).balanceOf(address(this));
        (bool ok, bytes memory ret) =
            token.call(abi.encodeWithSelector(IERC20.transferFrom.selector, from, address(this), amount));
        if (!ok || (ret.length != 0 && !abi.decode(ret, (bool)))) revert TransferFailed();
        // A token that takes a cut on the way in would let a shield create a note it did not pay for.
        if (IERC20(token).balanceOf(address(this)) - before < amount) revert TransferFailed();
    }

    function _push(address to, uint256 amount) internal {
        if (token == address(0)) return _pushHype(to, amount);
        (bool ok, bytes memory ret) = token.call(abi.encodeWithSelector(IERC20.transfer.selector, to, amount));
        if (!ok || (ret.length != 0 && !abi.decode(ret, (bool)))) revert TransferFailed();
    }

    function _pushHype(address to, uint256 amount) internal {
        (bool ok,) = to.call{value: amount}("");
        if (!ok) revert TransferFailed();
    }
}