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

import {Ownable, Ownable2Step} from "@openzeppelin/contracts/access/Ownable2Step.sol";
import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
import {ReentrancyGuard} from "@openzeppelin/contracts/utils/ReentrancyGuard.sol";

/// @title TokenVesting
/// @author sigmacode.io (Sanicura d.o.o.)
/// @notice Holds ERC-20 tokens for many beneficiaries and releases them on a cliff + linear schedule.
/// @dev The owner funds each schedule at creation (tokens are pulled with `transferFrom`), so the
///      contract is always fully collateralised: `balance >= totalAllocated - totalReleased`.
///
///      Vested amount at time `t` for a schedule `(start, cliff, duration, total)`:
///        - `t < start + cliff`          -> 0
///        - `t >= start + duration`      -> total
///        - otherwise                    -> total * (t - start) / duration
///
///      Revocable schedules can be revoked by the owner: the part vested so far stays claimable by the
///      beneficiary, the unvested remainder returns to the owner.
contract TokenVesting is Ownable2Step, ReentrancyGuard {
    using SafeERC20 for IERC20;

    /// @notice A single vesting schedule.
    struct Schedule {
        address beneficiary;
        uint64 start;
        uint64 cliffDuration;
        uint64 duration;
        bool revocable;
        bool revoked;
        uint256 totalAmount; // reduced to the vested amount on revocation
        uint256 released;
    }

    /// @notice Token being vested.
    IERC20 public immutable token;

    /// @notice Number of schedules created. Schedule ids are `0..scheduleCount-1`.
    uint256 public scheduleCount;

    /// @notice Sum of `totalAmount` over all schedules (after revocations).
    uint256 public totalAllocated;

    /// @notice Sum of tokens already released to beneficiaries.
    uint256 public totalReleased;

    mapping(uint256 scheduleId => Schedule) private _schedules;
    mapping(address beneficiary => uint256[] scheduleIds) private _beneficiarySchedules;

    /// @notice Emitted when a schedule is created and funded.
    event ScheduleCreated(
        uint256 indexed scheduleId,
        address indexed beneficiary,
        uint256 amount,
        uint64 start,
        uint64 cliffDuration,
        uint64 duration,
        bool revocable
    );
    /// @notice Emitted when vested tokens are released to a beneficiary.
    event TokensReleased(uint256 indexed scheduleId, address indexed beneficiary, uint256 amount);
    /// @notice Emitted when a revocable schedule is revoked.
    event ScheduleRevoked(uint256 indexed scheduleId, uint256 vestedAmount, uint256 refundedAmount);

    error ZeroAddress();
    error ZeroAmount();
    error ZeroDuration();
    error CliffLongerThanDuration(uint64 cliffDuration, uint64 duration);
    error ScheduleNotFound(uint256 scheduleId);
    error NotBeneficiary();
    error NothingToRelease();
    error NotRevocable(uint256 scheduleId);
    error AlreadyRevoked(uint256 scheduleId);

    /// @param token_ ERC-20 token to vest.
    /// @param initialOwner Owner who creates and revokes schedules.
    constructor(IERC20 token_, address initialOwner) Ownable(initialOwner) {
        if (address(token_) == address(0)) revert ZeroAddress();
        token = token_;
    }

    /// @notice Creates and funds a vesting schedule. The owner must have approved `amount` tokens.
    /// @param beneficiary Receiver of the vested tokens.
    /// @param start Unix timestamp when vesting starts (may be in the past or future).
    /// @param cliffDuration Seconds after `start` before anything vests.
    /// @param duration Total vesting duration in seconds (>= `cliffDuration`).
    /// @param amount Total tokens to vest.
    /// @param revocable Whether the owner may revoke the unvested remainder.
    /// @return scheduleId Identifier of the new schedule.
    function createSchedule(
        address beneficiary,
        uint64 start,
        uint64 cliffDuration,
        uint64 duration,
        uint256 amount,
        bool revocable
    ) external onlyOwner nonReentrant returns (uint256 scheduleId) {
        if (beneficiary == address(0)) revert ZeroAddress();
        if (amount == 0) revert ZeroAmount();
        if (duration == 0) revert ZeroDuration();
        if (cliffDuration > duration) revert CliffLongerThanDuration(cliffDuration, duration);

        scheduleId = scheduleCount++;
        _schedules[scheduleId] = Schedule({
            beneficiary: beneficiary,
            start: start,
            cliffDuration: cliffDuration,
            duration: duration,
            revocable: revocable,
            revoked: false,
            totalAmount: amount,
            released: 0
        });
        _beneficiarySchedules[beneficiary].push(scheduleId);
        totalAllocated += amount;

        emit ScheduleCreated(scheduleId, beneficiary, amount, start, cliffDuration, duration, revocable);
        token.safeTransferFrom(msg.sender, address(this), amount);
    }

    /// @notice Releases all currently releasable tokens of a schedule to its beneficiary.
    /// @dev Callable by the beneficiary only.
    /// @param scheduleId Schedule identifier.
    /// @return amount Amount released.
    function release(uint256 scheduleId) external nonReentrant returns (uint256 amount) {
        Schedule storage s = _existing(scheduleId);
        if (msg.sender != s.beneficiary) revert NotBeneficiary();
        amount = _vested(s, block.timestamp) - s.released;
        if (amount == 0) revert NothingToRelease();

        s.released += amount;
        totalReleased += amount;
        emit TokensReleased(scheduleId, s.beneficiary, amount);
        token.safeTransfer(s.beneficiary, amount);
    }

    /// @notice Revokes a revocable schedule. Vested tokens stay claimable; the rest returns to the owner.
    /// @param scheduleId Schedule identifier.
    function revoke(uint256 scheduleId) external onlyOwner nonReentrant {
        Schedule storage s = _existing(scheduleId);
        if (!s.revocable) revert NotRevocable(scheduleId);
        if (s.revoked) revert AlreadyRevoked(scheduleId);

        uint256 vested = _vested(s, block.timestamp);
        uint256 refund = s.totalAmount - vested;
        s.revoked = true;
        s.totalAmount = vested;
        totalAllocated -= refund;

        emit ScheduleRevoked(scheduleId, vested, refund);
        if (refund > 0) token.safeTransfer(owner(), refund);
    }

    /// @notice Returns a schedule.
    function getSchedule(uint256 scheduleId) external view returns (Schedule memory) {
        return _schedules[scheduleId];
    }

    /// @notice Returns all schedule ids of a beneficiary.
    function schedulesOf(address beneficiary) external view returns (uint256[] memory) {
        return _beneficiarySchedules[beneficiary];
    }

    /// @notice Amount vested at `timestamp` (including already released tokens).
    function vestedAmount(uint256 scheduleId, uint64 timestamp) external view returns (uint256) {
        return _vested(_existing(scheduleId), timestamp);
    }

    /// @notice Amount the beneficiary can release right now.
    function releasableAmount(uint256 scheduleId) external view returns (uint256) {
        Schedule storage s = _existing(scheduleId);
        return _vested(s, block.timestamp) - s.released;
    }

    function _existing(uint256 scheduleId) private view returns (Schedule storage s) {
        s = _schedules[scheduleId];
        if (s.beneficiary == address(0)) revert ScheduleNotFound(scheduleId);
    }

    function _vested(Schedule storage s, uint256 timestamp) private view returns (uint256) {
        if (s.revoked) return s.totalAmount; // frozen at revocation
        if (timestamp < uint256(s.start) + s.cliffDuration) return 0;
        if (timestamp >= uint256(s.start) + s.duration) return s.totalAmount;
        return (s.totalAmount * (timestamp - s.start)) / s.duration;
    }
}