// SPDX-License-Identifier: MIT
// OpenZeppelin Tron Contracts (last updated v5.6.0) (finance/VestingWallet.sol)

pragma solidity ^0.8.20;

import {ITRC20} from "@openzeppelin/tron-contracts/token/TRC20/ITRC20.sol";
import {SafeTRC20} from "@openzeppelin/tron-contracts/token/TRC20/utils/SafeTRC20.sol";
import {Address} from "@openzeppelin/tron-contracts/utils/Address.sol";
import {ContextUpgradeable} from "../utils/ContextUpgradeable.sol";
import {OwnableUpgradeable} from "../access/OwnableUpgradeable.sol";
import {Initializable} from "@openzeppelin/tron-contracts/proxy/utils/Initializable.sol";

/**
 * @dev A vesting wallet is an ownable contract that can receive native currency and TRC-20 tokens, and release these
 * assets to the wallet owner, also referred to as "beneficiary", according to a vesting schedule.
 *
 * Any assets transferred to this contract will follow the vesting schedule as if they were locked from the beginning.
 * Consequently, if the vesting has already started, any amount of tokens sent to this contract will (at least partly)
 * be immediately releasable.
 *
 * By setting the duration to 0, one can configure this contract to behave like an asset timelock that holds tokens for
 * a beneficiary until a specified time.
 *
 * NOTE: Since the wallet is {Ownable}, and ownership can be transferred, it is possible to sell unvested tokens.
 * Preventing this in a smart contract is difficult, considering that: 1) a beneficiary address could be a
 * counterfactually deployed contract, 2) there is likely to be a migration path for EOAs to become contracts in the
 * near future.
 *
 * NOTE: When using this contract with any token whose balance is adjusted automatically (i.e. a rebase token), make
 * sure to account the supply/balance adjustment in the vesting schedule to ensure the vested amount is as intended.
 *
 * NOTE: Chains with support for native TRC20s may allow the vesting wallet to withdraw the underlying asset as both an
 * TRC20 and as native currency. For example, if chain C supports token A and the wallet gets deposited 100 A, then
 * at 50% of the vesting period, the beneficiary can withdraw 50 A as TRC20 and 25 A as native currency (totaling 75 A).
 * Consider disabling one of the withdrawal methods.
 */
contract VestingWalletUpgradeable is Initializable, ContextUpgradeable, OwnableUpgradeable {
    event TRXReleased(uint256 amount);
    event TRC20Released(address indexed token, uint256 amount);

    /// @custom:storage-location erc7201:openzeppelin.storage.VestingWallet
    struct VestingWalletStorage {
        uint256 _released;
        mapping(address token => uint256) _trc20Released;
        uint64 _start;
        uint64 _duration;
    }

    // keccak256(abi.encode(uint256(keccak256("openzeppelin.storage.VestingWallet")) - 1)) & ~bytes32(uint256(0xff))
    bytes32 private constant VestingWalletStorageLocation =
        0xa1eac494560f7591e4da38ed031587f09556afdfc4399dd2e205b935fdfa3900;

    function _getVestingWalletStorage() private pure returns (VestingWalletStorage storage $) {
        assembly {
            $.slot := VestingWalletStorageLocation
        }
    }

    function initialize(address beneficiary, uint64 startTimestamp, uint64 durationSeconds) public virtual initializer {
        __VestingWallet_init(beneficiary, startTimestamp, durationSeconds);
    }
    /**
     * @dev Sets the beneficiary (owner), the start timestamp and the vesting duration (in seconds) of the vesting
     * wallet.
     */
    function __VestingWallet_init(
        address beneficiary,
        uint64 startTimestamp,
        uint64 durationSeconds
    ) internal onlyInitializing {
        __Ownable_init_unchained(beneficiary);
        __VestingWallet_init_unchained(beneficiary, startTimestamp, durationSeconds);
    }

    function __VestingWallet_init_unchained(
        address,
        uint64 startTimestamp,
        uint64 durationSeconds
    ) internal onlyInitializing {
        VestingWalletStorage storage $ = _getVestingWalletStorage();
        $._start = startTimestamp;
        $._duration = durationSeconds;
    }

    /**
     * @dev The contract should be able to receive TRX.
     */
    receive() external payable virtual {}

    /**
     * @dev Getter for the start timestamp.
     */
    function start() public view virtual returns (uint256) {
        VestingWalletStorage storage $ = _getVestingWalletStorage();
        return $._start;
    }

    /**
     * @dev Getter for the vesting duration.
     */
    function duration() public view virtual returns (uint256) {
        VestingWalletStorage storage $ = _getVestingWalletStorage();
        return $._duration;
    }

    /**
     * @dev Getter for the end timestamp.
     */
    function end() public view virtual returns (uint256) {
        return start() + duration();
    }

    /**
     * @dev Amount of TRX already released
     */
    function released() public view virtual returns (uint256) {
        VestingWalletStorage storage $ = _getVestingWalletStorage();
        return $._released;
    }

    /**
     * @dev Amount of token already released
     */
    function released(address token) public view virtual returns (uint256) {
        VestingWalletStorage storage $ = _getVestingWalletStorage();
        return $._trc20Released[token];
    }

    /**
     * @dev Getter for the amount of releasable TRX.
     */
    function releasable() public view virtual returns (uint256) {
        return vestedAmount(uint64(block.timestamp)) - released();
    }

    /**
     * @dev Getter for the amount of releasable `token` tokens. `token` should be the address of an
     * {ITRC20} contract.
     */
    function releasable(address token) public view virtual returns (uint256) {
        return vestedAmount(token, uint64(block.timestamp)) - released(token);
    }

    /**
     * @dev Release the native tokens (TRX) that have already vested.
     *
     * Emits a {TRXReleased} event.
     */
    function release() public virtual {
        VestingWalletStorage storage $ = _getVestingWalletStorage();
        uint256 amount = releasable();
        $._released += amount;
        emit TRXReleased(amount);
        Address.sendValue(payable(owner()), amount);
    }

    /**
     * @dev Release the tokens that have already vested.
     *
     * Emits a {TRC20Released} event.
     */
    function release(address token) public virtual {
        VestingWalletStorage storage $ = _getVestingWalletStorage();
        uint256 amount = releasable(token);
        $._trc20Released[token] += amount;
        emit TRC20Released(token, amount);
        // Balance-delta verified (see {SafeTRC20-safeTransferChecked}) so a wallet funded with a false-on-success
        // token such as TRON USDT can be released rather than reverting on every payout.
        SafeTRC20.safeTransferChecked(ITRC20(token), owner(), amount);
    }

    /**
     * @dev Calculates the amount of TRX that has already vested. Default implementation is a linear vesting curve.
     */
    function vestedAmount(uint64 timestamp) public view virtual returns (uint256) {
        return _vestingSchedule(address(this).balance + released(), timestamp);
    }

    /**
     * @dev Calculates the amount of tokens that has already vested. Default implementation is a linear vesting curve.
     */
    function vestedAmount(address token, uint64 timestamp) public view virtual returns (uint256) {
        return _vestingSchedule(ITRC20(token).balanceOf(address(this)) + released(token), timestamp);
    }

    /**
     * @dev Virtual implementation of the vesting formula. This returns the amount vested, as a function of time, for
     * an asset given its total historical allocation.
     */
    function _vestingSchedule(uint256 totalAllocation, uint64 timestamp) internal view virtual returns (uint256) {
        if (timestamp < start()) {
            return 0;
        } else if (timestamp >= end()) {
            return totalAllocation;
        } else {
            return (totalAllocation * (timestamp - start())) / duration();
        }
    }
}