// SPDX-License-Identifier: UNLICENSED

pragma solidity =0.8.18;

import "@openzeppelin/contracts-upgradeable/token/ERC20/ERC20Upgradeable.sol";
import "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";

import "./IERC20Partition.sol";

abstract contract ERC20Partition is ERC20Upgradeable, IERC20Partition {
    using EnumerableSet for EnumerableSet.Bytes32Set;

    bytes32 public constant DEFAULT_PARTITION = 0x00;

    mapping(address => EnumerableSet.Bytes32Set) private _partitions;
    mapping(bytes32 => uint256) private _totalSupplies;
    mapping(bytes32 => bool) private _deprecated;
    mapping(address => mapping(bytes32 => uint256)) private _balances;
    mapping(address => mapping(address => mapping(bytes32 => uint256))) private _allowances;

    uint256[50] private __gap;

    /**
     * @dev Returns the amount of tokens in existence.
     */
    function totalSupply(bytes32 id) public view returns (uint256) {
        return _totalSupplies[id];
    }

    /**
     * @dev Returns a boolean indicating if a partition is deprecated.
     */
    function deprecated(bytes32 id) public view returns (bool) {
        return _deprecated[id];
    }

    /**
     * @dev Returns the amount of tokens owned by `account`.
     */
    function balanceOf(address account, bytes32 id) public view returns (uint256) {
        return _balances[account][id];
    }

    /**
     * @dev Moves `amount` tokens from the caller's account to `to`.
     *
     * Returns a boolean value indicating whether the operation succeeded.
     *
     * Emits a {Transfer} event.
     */
    function transfer(address to, bytes32 id, uint256 amount, bytes memory data) public returns (bool) {
        address owner = _msgSender();
        _transfer(owner, to, id, id, amount, data);
        return true;
    }

    /**
     * @dev Returns the remaining number of tokens that `spender` will be
     * allowed to spend on behalf of `owner` through {transferFrom}. This is
     * zero by default.
     *
     * This value changes when {approve} or {transferFrom} are called.
     */
    function allowance(address owner, address spender, bytes32 id) public view returns (uint256) {
        return _allowances[owner][spender][id];
    }

    /**
     * @dev Sets `amount` as the allowance of `spender` over the caller's tokens.
     *
     * Returns a boolean value indicating whether the operation succeeded.
     *
     * IMPORTANT: Beware that changing an allowance with this method brings the risk
     * that someone may use both the old and the new allowance by unfortunate
     * transaction ordering. One possible solution to mitigate this race
     * condition is to first reduce the spender's allowance to 0 and set the
     * desired value afterwards:
     * https://github.com/ethereum/EIPs/issues/20#issuecomment-263524729
     *
     * Emits an {Approval} event.
     */
    function approve(address spender, bytes32 id, uint256 amount, bytes memory data) public returns (bool) {
        address owner = _msgSender();
        _approve(owner, spender, id, amount, data);
        return true;
    }

    /**
     * @dev Moves `amount` tokens from `from` to `to` using the
     * allowance mechanism. `amount` is then deducted from the caller's
     * allowance.
     *
     * Returns a boolean value indicating whether the operation succeeded.
     *
     * Emits a {Transfer} event.
     */
    function transferFrom(address from, address to, bytes32 id, uint256 amount, bytes memory data)
        public
        returns (bool)
    {
        address spender = _msgSender();
        _spendAllowance(from, spender, id, amount, data);
        _transfer(from, to, id, id, amount, data);
        return true;
    }

    function _transfer(address from, address to, uint256 amount) internal virtual override {
        require(from != address(0), "ERC20Partition: transfer from the zero address");
        require(to != address(0), "ERC20Partition: transfer to the zero address");

        EnumerableSet.Bytes32Set storage _userPartitions = _partitions[from];

        uint256 _remainingBalance = amount;
        uint256 _fromBalance = _balances[from][DEFAULT_PARTITION];
        uint256 _useAmount;

        if (_fromBalance > 0) {
            _useAmount = _calculateUseAmount(_fromBalance, _remainingBalance);
            _transferPartition(from, to, DEFAULT_PARTITION, DEFAULT_PARTITION, _fromBalance, _useAmount, "");

            unchecked {
                _remainingBalance -= _useAmount;
            }
        }

        if (_remainingBalance > 0) {
            uint256 _totalPartitions = _userPartitions.length();

            uint256 i;
            while (i < _totalPartitions) {
                bytes32 _id = _userPartitions.at(i);
                _fromBalance = _balances[from][_id];
                if (_fromBalance == 0) continue;

                _useAmount = _calculateUseAmount(_fromBalance, _remainingBalance);

                unchecked {
                    _remainingBalance -= _useAmount;
                }

                _transferPartition(from, to, _id, _id, _fromBalance, _useAmount, "");

                if (_fromBalance == _useAmount) {
                    unchecked {
                        _totalPartitions--;
                    }
                } else {
                    unchecked {
                        i++;
                    }
                }
            }

            require(_remainingBalance > 0, "ERC20Partition: transfer amount exceeds balance");
        }

        super._transfer(from, to, amount);
    }

    function _transferPartition(
        address from,
        address to,
        bytes32 fromPartitionId,
        bytes32 toPartitionId,
        uint256 fromBalance,
        uint256 amount,
        bytes memory data
    ) internal virtual {
        if (toPartitionId != DEFAULT_PARTITION && _deprecated[toPartitionId]) {
            toPartitionId = DEFAULT_PARTITION;
        }

        require(fromBalance >= amount, "ERC20Partition: insufficient partition balance");

        _beforeTokenTransfer(from, to, fromPartitionId, toPartitionId, amount, data);

        unchecked {
            _balances[from][fromPartitionId] = fromBalance - amount;
            _balances[to][toPartitionId] += amount;
        }

        if (fromBalance == amount && fromPartitionId != DEFAULT_PARTITION) {
            _partitions[from].remove(fromPartitionId);
        }

        if (toPartitionId != DEFAULT_PARTITION) {
            _partitions[to].add(toPartitionId);
        }

        if (fromPartitionId != toPartitionId) {
            uint256 _fromTotalSupply = _totalSupplies[fromPartitionId];

            unchecked {
                _totalSupplies[fromPartitionId] = _fromTotalSupply - amount;
                _totalSupplies[toPartitionId] += amount;
            }

            if (_fromTotalSupply == amount) {
                delete _deprecated[fromPartitionId];
            }
        }

        emit TransferPartition(from, to, fromPartitionId, toPartitionId, amount, data);

        _afterTokenTransfer(from, to, fromPartitionId, toPartitionId, amount, data);
    }

    /**
     * @dev Moves `amount` of tokens from `from` to `to`.
     *
     * This internal function is equivalent to {transfer}, and can be used to
     * e.g. implement automatic token fees, slashing mechanisms, etc.
     *
     * Emits a {Transfer} event.
     *
     * Requirements:
     *
     * - `from` cannot be the zero address.
     * - `to` cannot be the zero address.
     * - `from` must have a balance of at least `amount`.
     */
    function _transfer(address from, address to, bytes32 fromId, bytes32 toId, uint256 amount, bytes memory data)
        internal
        virtual
    {
        require(from != address(0), "ERC20Partition: transfer from the zero address");
        require(to != address(0), "ERC20Partition: transfer to the zero address");

        _transferPartition(from, to, fromId, toId, _balances[from][fromId], amount, data);

        super._transfer(from, to, amount);
    }

    /**
     * @dev Creates `amount` tokens and assigns them to `account`, increasing
     * the total supply.
     *
     * Emits a {Transfer} event with `from` set to the zero address.
     *
     * Requirements:
     *
     * - `account` cannot be the zero address.
     */
    function _mint(address account, bytes32 id, uint256 amount, bytes memory data) internal virtual {
        require(account != address(0), "ERC20Partition: mint to the zero address");

        _beforeTokenTransfer(address(0), account, id, id, amount, data);

        _totalSupplies[id] += amount;
        unchecked {
            _balances[account][id] += amount;
        }

        if (id != DEFAULT_PARTITION) {
            _partitions[account].add(id);
        }

        emit TransferPartition(address(0), account, id, id, amount, data);

        super._mint(account, amount);

        _afterTokenTransfer(address(0), account, id, id, amount, data);
    }

    /**
     * @dev Destroys `amount` tokens from `account`, reducing the
     * total supply.
     *
     * Emits a {Transfer} event with `to` set to the zero address.
     *
     * Requirements:
     *
     * - `account` cannot be the zero address.
     * - `account` must have at least `amount` tokens.
     */
    function _burn(address account, bytes32 id, uint256 amount, bytes memory data) internal virtual {
        require(account != address(0), "ERC20Partition: burn from the zero address");

        _beforeTokenTransfer(account, address(0), id, id, amount, data);

        uint256 accountBalance = _balances[account][id];
        require(accountBalance >= amount, "ERC20Partition: burn amount exceeds balance");
        unchecked {
            _balances[account][id] = accountBalance - amount;
        }
        _totalSupplies[id] -= amount;

        if (id != DEFAULT_PARTITION && amount == accountBalance) {
            _partitions[account].remove(id);
        }

        emit TransferPartition(account, address(0), id, id, amount, data);

        super._burn(account, amount);

        _afterTokenTransfer(account, address(0), id, id, amount, data);
    }

    /**
     * @dev Sets `amount` as the allowance of `spender` over the `owner` s tokens.
     *
     * This internal function is equivalent to `approve`, and can be used to
     * e.g. set automatic allowances for certain subsystems, etc.
     *
     * Emits an {Approval} event.
     *
     * Requirements:
     *
     * - `owner` cannot be the zero address.
     * - `spender` cannot be the zero address.
     */
    function _approve(address owner, address spender, bytes32 id, uint256 amount, bytes memory data) internal virtual {
        require(owner != address(0), "ERC20Partition: approve from the zero address");
        require(spender != address(0), "ERC20Partition: approve to the zero address");

        _allowances[owner][spender][id] = amount;
        emit ApprovalPartition(owner, spender, id, amount, data);
    }

    /**
     * @dev Deprecates a partition.
     *
     * Emits an {DeprecatePartition} event.
     *
     * Requirements:
     *
     * - `id` cannot be an already deprecated partition.
     * - `id` cannot be the DEFAULT_PARTITION.
     */
    function _deprecatePartition(bytes32 id, bytes memory data) internal virtual {
        require(!_deprecated[id], "ERC20Partition: already deprecated");
        require(id != DEFAULT_PARTITION, "ERC20Partition: cannot deprecate the default partition");

        _deprecated[id] = true;
        emit DeprecatePartition(id, data);
    }

    /**
     *
     * Does not update the allowance amount in case of infinite allowance.
     * Revert if not enough allowance is available.
     *
     * Might emit an {Approval} event.
     */
    function _spendAllowance(address owner, address spender, bytes32 id, uint256 amount, bytes memory data)
        internal
        virtual
    {
        uint256 currentAllowance = allowance(owner, spender, id);
        if (currentAllowance == type(uint256).max) return;
        if (currentAllowance >= amount) {
            unchecked {
                _approve(owner, spender, id, currentAllowance - amount, data);
            }
        } else {
            if (currentAllowance > 0) {
                _approve(owner, spender, id, 0, data);
            }
            unchecked {
                _spendAllowance(owner, spender, amount - currentAllowance);
            }
        }
    }

    /**
     * @dev Hook that is called before any transfer of tokens. This includes
     * minting and burning.
     *
     * Calling conditions:
     *
     * - when `from` and `to` are both non-zero, `amount` of ``from``'s tokens
     * will be transferred to `to`.
     * - when `from` is zero, `amount` tokens will be minted for `to`.
     * - when `to` is zero, `amount` of ``from``'s tokens will be burned.
     * - `from` and `to` are never both zero.
     *
     * To learn more about hooks, head to xref:ROOT:extending-contracts.adoc#using-hooks[Using Hooks].
     */
    function _beforeTokenTransfer(
        address from,
        address to,
        bytes32 fromId,
        bytes32 toId,
        uint256 amount,
        bytes memory data
    ) internal virtual {}

    /**
     * @dev Hook that is called after any transfer of tokens. This includes
     * minting and burning.
     *
     * Calling conditions:
     *
     * - when `from` and `to` are both non-zero, `amount` of ``from``'s tokens
     * has been transferred to `to`.
     * - when `from` is zero, `amount` tokens have been minted for `to`.
     * - when `to` is zero, `amount` of ``from``'s tokens have been burned.
     * - `from` and `to` are never both zero.
     *
     * To learn more about hooks, head to xref:ROOT:extending-contracts.adoc#using-hooks[Using Hooks].
     */
    function _afterTokenTransfer(
        address from,
        address to,
        bytes32 fromId,
        bytes32 toId,
        uint256 amount,
        bytes memory data
    ) internal virtual {}

    function _calculateUseAmount(uint256 balance, uint256 amount) private pure returns (uint256 useAmount) {
        unchecked {
            useAmount = balance > amount ? amount : balance;
        }
    }
}