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

import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import "@openzeppelin/contracts/security/ReentrancyGuard.sol";

/**
 * @title AereCompliancePool — Privacy Pools + Travel Rule hybrid
 * @notice Implements the Privacy Pools construction of Buterin, Illum,
 *         Nadler, Schar & Soleimani (eprint 2023/1366) — zk-mixer with
 *         association-set proofs that let honest users prove their deposit
 *         belongs to a curated "clean" subset, while sanctioned addresses
 *         cannot exit via the pool because their commitment never enters
 *         the association set.
 *
 *         The deposit root is computed on-chain from a deterministic
 *         append-only Merkle tree of committed leaves (audit fix HIGH #28).
 *         The Foundation cannot publish a fake tree to drain the pool.
 *
 *         Hashing: keccak256 (cheap on EVM, no trusted setup, no MiMC).
 *         Depth: 20 (1,048,576 leaves per pool).
 *         History: last 64 roots remain valid simultaneously so a user with
 *         a proof built against a slightly-old root can still withdraw.
 *
 *         The Foundation still publishes ASSOCIATION roots through the 24h
 *         challenge process, since the association set is an off-chain
 *         classification.
 */
interface IPrivacyPoolVerifier {
    /// @param  publicInputs   [poolId, depositRoot, associationRoot, nullifierHash,
    ///                         recipientWord, feeRecipientWord, refundWord, feeWord]
    function verify(bytes calldata proof, uint256[] calldata publicInputs) external view returns (bool);
}

/// @dev Append-only zk-mixer Merkle tree with rolling-history root cache.
/// Construction follows the Privacy Pools paper (Buterin et al. 2023).
abstract contract MerkleTreeWithHistory {
    uint32 public constant TREE_DEPTH = 20;
    /// AUDIT FIX R5 (MED-10): raised from 64 → 256. At 0.5s blocks, 64 deposits
    /// fit in ~32s; honest provers building proofs against a 30s-old root
    /// were getting DoS'd by deposit bursts. 256 gives ~2min headroom which
    /// matches expected prover latency on AERE.
    uint32 public constant ROOT_HISTORY_SIZE = 256;

    bytes32 public constant ZERO_VALUE = keccak256("aere.compliance-pool.v1.empty-leaf");

    bytes32[TREE_DEPTH] public filledSubtrees;
    bytes32[TREE_DEPTH] public zeros;
    bytes32[ROOT_HISTORY_SIZE] public roots;
    uint32 public currentRootIndex;
    uint32 public nextLeafIndex;

    error TreeFull();

    event TreeNearCapacity(uint32 indexed leaves, uint32 capacity); // AUDIT FIX R4 (MED)

    constructor() {
        // AUDIT FIX R4 (LOW): do NOT seed the empty-tree root into the rolling
        // history — that would let withdrawals against an empty-tree proof
        // succeed for the first ROOT_HISTORY_SIZE deposits. Instead, leave all
        // history slots at bytes32(0); isKnownRoot already returns false for
        // zero. The first real _insert writes the first valid root.
        bytes32 z = ZERO_VALUE;
        for (uint32 i = 0; i < TREE_DEPTH; i++) {
            zeros[i] = z;
            filledSubtrees[i] = z;
            z = keccak256(abi.encodePacked(z, z));
        }
        // history intentionally left as zeros.
    }

    function _insert(bytes32 leaf) internal returns (uint32 idx) {
        idx = nextLeafIndex;
        if (idx >= uint32(1 << TREE_DEPTH)) revert TreeFull();
        nextLeafIndex = idx + 1;

        uint32 cur = idx;
        bytes32 left;
        bytes32 right;
        bytes32 node = leaf;
        for (uint32 i = 0; i < TREE_DEPTH; i++) {
            if (cur & 1 == 0) {
                left = node;
                right = zeros[i];
                filledSubtrees[i] = node;
            } else {
                left = filledSubtrees[i];
                right = node;
            }
            node = keccak256(abi.encodePacked(left, right));
            cur >>= 1;
        }

        uint32 newRootIdx = (currentRootIndex + 1) % ROOT_HISTORY_SIZE;
        currentRootIndex = newRootIdx;
        roots[newRootIdx] = node;

        // AUDIT FIX R4 (MED): warn at 90% and 99% capacity so operators can
        // schedule a sibling-pool deployment before depositors hit TreeFull.
        uint32 cap = uint32(1 << TREE_DEPTH);
        uint32 leaves = nextLeafIndex;
        if (leaves == (cap * 90) / 100 || leaves == (cap * 99) / 100) {
            emit TreeNearCapacity(leaves, cap);
        }
    }

    /// @notice Returns true if `root` is in the recent ROOT_HISTORY_SIZE roots.
    function isKnownRoot(bytes32 root) public view returns (bool) {
        if (root == bytes32(0)) return false;
        uint32 i = currentRootIndex;
        do {
            if (roots[i] == root) return true;
            if (i == 0) i = ROOT_HISTORY_SIZE;
            i--;
        } while (i != currentRootIndex);
        return false;
    }

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

contract AereCompliancePool is MerkleTreeWithHistory, ReentrancyGuard {

    /* ================================ config ================================ */

    IERC20 public immutable TOKEN;
    uint256 public immutable DENOMINATION;
    address public immutable FOUNDATION;
    IPrivacyPoolVerifier public immutable VERIFIER;
    bytes32 public immutable POOL_ID;

    /* ================================= state ================================= */

    /// AUDIT FIX (MED #28): association roots still go through 24h challenge.
    mapping(bytes32 => uint64) public associationRootPublishedAt;
    mapping(bytes32 => uint64) public associationRootProposedAt;
    mapping(bytes32 => bool)   public associationRootChallenged;
    /// AUDIT FIX R5 (HIGH-2): cap dismissals + slashable challenge bond.
    mapping(bytes32 => uint8) public associationRootDismissCount;
    mapping(bytes32 => address) public associationRootChallenger;
    uint8  public constant MAX_DISMISSALS_PER_ROOT = 2;
    uint256 public constant CHALLENGE_BOND = 100 ether; // 100 AERE — burnt on dismissal
    address public constant CHALLENGE_BOND_BURN = 0x000000000000000000000000000000000000dEaD;
    IERC20 public immutable BOND_TOKEN; // bond paid in same TOKEN as DENOMINATION
    uint64 public constant ASSOCIATION_ROOT_CHALLENGE_WINDOW = 24 hours;

    mapping(bytes32 => bool) public spentNullifiers;
    mapping(bytes32 => bool) public commitmentSeen; // prevents duplicate commitment deposits
    uint256 public depositCount;

    mapping(address => bool) public isComplianceProvider;
    address[] public complianceProviderList;

    /* ================================= events ================================= */

    event Deposited(
        uint256 indexed depositId,
        address indexed sender,
        bytes32 commitment,
        bytes32 travelRuleHash,
        address indexed complianceProvider,
        string  travelRuleUri,
        uint32  leafIndex,
        bytes32 newRoot
    );
    event AssociationRootProposed(bytes32 indexed root, uint64 challengeEndsAt);
    event AssociationRootChallenged(bytes32 indexed root, address indexed challenger, string reason);
    event AssociationRootChallengeDismissed(bytes32 indexed root, address indexed challenger, uint8 dismissCount); // R5 HIGH-2
    event AssociationRootPublished(bytes32 indexed root);
    event Withdrawn(
        bytes32 indexed nullifierHash,
        address indexed recipient,
        address indexed feeRecipient,
        uint256 fee,
        uint256 refund,
        bytes32 depositRoot,
        bytes32 associationRoot
    );
    event ComplianceProviderSet(address indexed provider, bool allowed);

    /* ================================= errors ================================= */

    error NotFoundation();
    error ZeroAddr();
    error NotAComplianceProvider();
    error AlreadySpent();
    error UnknownAssociationRoot();
    error UnknownDepositRoot();
    error InvalidProof();
    error WrongDenomination();
    error TransferFailed();
    error FeeTooLarge();
    error CommitmentReused();
    error ChallengeBondTransferFailed(); // R5 HIGH-2
    error MaxDismissalsReached();        // R5 HIGH-2

    modifier onlyFoundation() {
        if (msg.sender != FOUNDATION) revert NotFoundation();
        _;
    }

    constructor(
        IERC20 token,
        uint256 denomination,
        address foundation,
        IPrivacyPoolVerifier verifier
    ) MerkleTreeWithHistory() {
        if (address(token) == address(0) || foundation == address(0) || address(verifier) == address(0)) {
            revert ZeroAddr();
        }
        require(denomination > 0, "zero-denom");
        TOKEN = token;
        BOND_TOKEN = token; // R5 HIGH-2: same asset as DENOMINATION
        DENOMINATION = denomination;
        FOUNDATION = foundation;
        VERIFIER = verifier;
        POOL_ID = keccak256(abi.encodePacked(address(this), block.chainid));
    }

    /* ============================ compliance provider ========================= */

    function setComplianceProvider(address provider, bool allowed) external onlyFoundation {
        if (provider == address(0)) revert ZeroAddr();
        bool prev = isComplianceProvider[provider];
        isComplianceProvider[provider] = allowed;
        if (!prev && allowed) complianceProviderList.push(provider);
        emit ComplianceProviderSet(provider, allowed);
    }

    /* ============================== deposit path ============================= */

    function deposit(
        bytes32 commitment,
        bytes32 travelRuleHash,
        address complianceProvider,
        string calldata travelRuleUri
    ) external nonReentrant {
        if (!isComplianceProvider[complianceProvider]) revert NotAComplianceProvider();
        if (commitment == bytes32(0) || commitment == ZERO_VALUE) revert CommitmentReused();
        // AUDIT FIX R5 (MED-11): bind commitment to msg.sender. Mempool
        // front-run by another address now produces a DIFFERENT actualLeaf so
        // the victim's deposit still succeeds. The ZK circuit will need to
        // constrain that the witnessed leaf = keccak256(commitment, depositor).
        bytes32 actualLeaf = keccak256(abi.encodePacked(commitment, msg.sender));
        if (commitmentSeen[actualLeaf]) revert CommitmentReused();
        commitmentSeen[actualLeaf] = true;

        if (!TOKEN.transferFrom(msg.sender, address(this), DENOMINATION)) revert TransferFailed();

        uint32 leafIndex = _insert(actualLeaf);
        bytes32 newRoot = roots[currentRootIndex];

        uint256 id = depositCount++;
        emit Deposited(id, msg.sender, commitment, travelRuleHash, complianceProvider, travelRuleUri, leafIndex, newRoot);
    }

    /* ========================== association set publish ====================== */

    function proposeAssociationRoot(bytes32 root) external onlyFoundation {
        if (associationRootPublishedAt[root] != 0) return;
        if (associationRootProposedAt[root] == 0) {
            associationRootProposedAt[root] = uint64(block.timestamp);
            associationRootChallenged[root] = false;
            emit AssociationRootProposed(root, uint64(block.timestamp) + ASSOCIATION_ROOT_CHALLENGE_WINDOW);
        }
    }

    function challengeAssociationRoot(bytes32 root, string calldata reason) external {
        if (associationRootProposedAt[root] == 0) revert UnknownAssociationRoot();
        if (associationRootPublishedAt[root] != 0) revert UnknownAssociationRoot();
        if (associationRootChallenged[root]) revert UnknownAssociationRoot(); // already challenged — wait for dismissal
        if (associationRootDismissCount[root] >= MAX_DISMISSALS_PER_ROOT) revert MaxDismissalsReached();
        // AUDIT FIX R5 (HIGH-2): require slashable CHALLENGE_BOND so a griefer
        // can no longer freeze withdrawals at ~30k gas per 24h cycle.
        if (!BOND_TOKEN.transferFrom(msg.sender, address(this), CHALLENGE_BOND)) revert ChallengeBondTransferFailed();
        associationRootChallenged[root] = true;
        associationRootChallenger[root] = msg.sender;
        emit AssociationRootChallenged(root, msg.sender, reason);
    }

    function dismissAssociationRootChallenge(bytes32 root) external onlyFoundation {
        if (!associationRootChallenged[root]) revert UnknownAssociationRoot();
        associationRootChallenged[root] = false;
        associationRootDismissCount[root] += 1;
        // AUDIT FIX R6 (CRITICAL): reset window on dismiss. Combined with
        // MAX_DISMISSALS_PER_ROOT=2 this caps the worst-case Foundation+
        // sockpuppet bypass at 48h (vs R5's 0h instant publish). Each
        // dismissal genuinely costs Foundation 24h delay AND burns the
        // challenger's bond.
        associationRootProposedAt[root] = uint64(block.timestamp);
        address challenger = associationRootChallenger[root];
        delete associationRootChallenger[root];
        if (challenger != address(0)) {
            // best-effort burn
            BOND_TOKEN.transfer(CHALLENGE_BOND_BURN, CHALLENGE_BOND);
        }
        emit AssociationRootChallengeDismissed(root, challenger, associationRootDismissCount[root]);
        emit AssociationRootProposed(root, uint64(block.timestamp) + ASSOCIATION_ROOT_CHALLENGE_WINDOW);
    }

    function publishAssociationRoot(bytes32 root) external {
        if (associationRootPublishedAt[root] != 0) return;
        uint64 proposed = associationRootProposedAt[root];
        if (proposed == 0) revert UnknownAssociationRoot();
        if (associationRootChallenged[root]) revert UnknownAssociationRoot();
        if (block.timestamp < uint256(proposed) + ASSOCIATION_ROOT_CHALLENGE_WINDOW) {
            revert UnknownAssociationRoot();
        }
        associationRootPublishedAt[root] = uint64(block.timestamp);
        emit AssociationRootPublished(root);
    }

    /* ============================== withdraw path ============================ */

    function withdraw(
        bytes calldata proof,
        bytes32 depositRoot,
        bytes32 associationRoot,
        bytes32 nullifierHash,
        address recipient,
        address feeRecipient,
        uint256 fee,
        uint256 refund
    ) external nonReentrant {
        if (spentNullifiers[nullifierHash]) revert AlreadySpent();
        if (associationRootPublishedAt[associationRoot] == 0) revert UnknownAssociationRoot();
        // AUDIT FIX (HIGH #28): the depositRoot must be one of the last 64
        // on-chain-computed roots — there is no Foundation override.
        if (!isKnownRoot(depositRoot)) revert UnknownDepositRoot();
        if (fee + refund > DENOMINATION) revert FeeTooLarge();

        uint256[] memory pubs = new uint256[](8);
        pubs[0] = uint256(POOL_ID);
        pubs[1] = uint256(depositRoot);
        pubs[2] = uint256(associationRoot);
        pubs[3] = uint256(nullifierHash);
        pubs[4] = uint256(uint160(recipient));
        pubs[5] = uint256(uint160(feeRecipient));
        pubs[6] = refund;
        pubs[7] = fee;

        if (!VERIFIER.verify(proof, pubs)) revert InvalidProof();

        spentNullifiers[nullifierHash] = true;

        // AUDIT FIX R6 (HIGH): refund is paid out to msg.sender (the relayer
        // or sponsor) as a gas-reimbursement kickback — the standard
        // zk-mixer relayer pattern. Previously refund was subtracted from
        // recipient amount but NEVER paid out — pool accounting was unbalanced.
        uint256 toRecipient = DENOMINATION - fee - refund;
        if (toRecipient > 0 && !TOKEN.transfer(recipient, toRecipient)) revert TransferFailed();
        if (fee > 0 && !TOKEN.transfer(feeRecipient, fee)) revert TransferFailed();
        if (refund > 0 && !TOKEN.transfer(msg.sender, refund)) revert TransferFailed();

        emit Withdrawn(nullifierHash, recipient, feeRecipient, fee, refund, depositRoot, associationRoot);
    }

    /* ================================= views ================================= */

    function complianceProviders() external view returns (address[] memory) {
        return complianceProviderList;
    }

    function isAssociationRootValid(bytes32 root) external view returns (bool) {
        return associationRootPublishedAt[root] != 0;
    }

    /// @notice Compatibility alias for older indexers/clients that expect a
    /// `latestDepositRoot()` getter. Returns the current on-chain Merkle root.
    function latestDepositRoot() external view returns (bytes32) {
        return roots[currentRootIndex];
    }
}