// SPDX-License-Identifier: MIT
// OpenZeppelin Contracts v4.4.1 (interfaces/IERC20.sol)

pragma solidity ^0.8.0;

import "../token/ERC20/IERC20.sol";

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

import {IERC20} from "@openzeppelin/interfaces/IERC20.sol";

error NotStaked();
error NoStakers();
error CannotStakeZeroTokens();

contract Staking {
    event UserStaked(address indexed who, uint256 amount);
    event UserUnstaked(address indexed who, uint256 amount);

    event RewardsPayout(address indexed to, uint256 amount);
    event CapturedRewards(address indexed from, uint256 amount);

    event RewardsToTreasury(uint256 amount);

    struct StakingRecord {
        uint256 stakedAmount;
        uint256 rewardsPerStakedTokenSnapshot; // Scaled by a WAD
    }

    uint256 private constant WAD = 1e18;

    /// @notice Address of the treasury where 50% of the profits
    /// will be sent.
    address public immutable treasury;

    IERC20 public immutable weth;
    IERC20 public immutable prtc;

    /// @notice Amount of `WETH` rewards per staked token.
    /// @dev Scaled by a WAD.
    uint256 public rewardsPerStakedToken;
    /// @notice Total amount of `prtc` staked.
    uint256 public totalStaked;

    mapping(address => StakingRecord) public stakingRecords;

    constructor(IERC20 _prtc, IERC20 _weth, address _treasury) {
        prtc = _prtc;
        weth = _weth;
        treasury = _treasury;
    }

    /// @notice Method to claim `WETH` rewards accumulated in the
    /// contract.
    function claimRewards() public {
        StakingRecord storage userRecord = stakingRecords[msg.sender];

        if (userRecord.stakedAmount == 0) revert NotStaked();

        uint256 userReward = userRecord.stakedAmount
            * (rewardsPerStakedToken - userRecord.rewardsPerStakedTokenSnapshot) / WAD;

        userRecord.rewardsPerStakedTokenSnapshot = rewardsPerStakedToken;

        if (userReward != 0) {
            weth.transfer(msg.sender, userReward);
            emit RewardsPayout(msg.sender, userReward);
        }
    }

    /// @notice Retrieve any users' `WETH` rewards.
    function previewRewards(address user) public view returns (uint256) {
        StakingRecord storage userRecord = stakingRecords[user];

        uint256 totalStake = userRecord.stakedAmount;

        if (totalStake == 0) return 0;

        uint256 userReward =
            totalStake * (rewardsPerStakedToken - userRecord.rewardsPerStakedTokenSnapshot) / WAD;

        return userReward;
    }

    /// @notice Retrieve users' staked amount.
    function balanceOf(address user) public view returns (uint256) {
        return stakingRecords[user].stakedAmount;
    }

    /// @notice Staking method that accepts `prtc`. Profits from the
    /// vault will be distributed to this contract and staking users
    /// will be rewarded with `WETH`.
    /// @param amount Amount of `prtc` to be staked.
    function stake(uint256 amount) external {
        if (amount == 0) revert CannotStakeZeroTokens();

        StakingRecord storage userPosition = stakingRecords[msg.sender];

        if (userPosition.stakedAmount != 0) claimRewards();

        prtc.transferFrom(msg.sender, address(this), amount);

        totalStaked += amount;

        userPosition.stakedAmount += amount;
        userPosition.rewardsPerStakedTokenSnapshot = rewardsPerStakedToken;

        emit UserStaked(msg.sender, amount);
    }

    /// @notice Method to unstake all `prtc`.
    function unstake() external {
        unstake(type(uint256).max);
    }

    /// @notice Method to unstake `prtc`. It will calculate the portion
    /// of `WETH` rewards for the user and send them.
    function unstake(uint256 _amount) public {
        if (_amount == 0) _amount = type(uint256).max;

        StakingRecord storage userRecord = stakingRecords[msg.sender];

        uint256 userStakedAmount = userRecord.stakedAmount;

        if (userStakedAmount == 0) revert NotStaked();

        uint256 userReward = previewRewards(msg.sender);

        if (_amount >= userStakedAmount) {
            totalStaked -= userStakedAmount;
            delete stakingRecords[msg.sender];

            prtc.transfer(msg.sender, userStakedAmount);
            emit UserUnstaked(msg.sender, userStakedAmount);
        } else {
            totalStaked -= _amount;
            userRecord.stakedAmount -= _amount;
            userRecord.rewardsPerStakedTokenSnapshot = rewardsPerStakedToken;

            prtc.transfer(msg.sender, _amount);
            emit UserUnstaked(msg.sender, _amount);
        }

        if (userReward != 0) {
            weth.transfer(msg.sender, userReward);
            emit RewardsPayout(msg.sender, userReward);
        }
    }

    /// @notice Method to distribute `WETH` rewards.
    /// @dev Anyone that wants to distribute `WETH` can call this.
    /// @param amount The amount of `WETH` to be distributed to the contract.
    function distribute(uint256 amount) external {
        if (totalStaked == 0) revert NoStakers();

        uint256 amountForTreasury = amount / 2;
        uint256 amountForStakers = amount - amountForTreasury;

        weth.transferFrom(msg.sender, treasury, amountForTreasury);
        weth.transferFrom(msg.sender, address(this), amountForStakers);

        // Distribute the other 50% to the stakers.
        rewardsPerStakedToken += amountForStakers * WAD / totalStaked;

        emit RewardsToTreasury(amountForTreasury);
        emit CapturedRewards(msg.sender, amountForStakers);
    }
}

{
  "compilationTarget": {
    "src/Staking.sol": "Staking"
  },
  "evmVersion": "paris",
  "libraries": {},
  "metadata": {
    "bytecodeHash": "ipfs"
  },
  "optimizer": {
    "enabled": true,
    "runs": 200
  },
  "remappings": [
    ":@openzeppelin/=lib/openzeppelin-contracts/contracts/",
    ":ds-test/=lib/forge-std/lib/ds-test/src/",
    ":erc4626-tests/=lib/openzeppelin-contracts/lib/erc4626-tests/",
    ":forge-std/=lib/forge-std/src/",
    ":openzeppelin-contracts/=lib/openzeppelin-contracts/",
    ":openzeppelin/=lib/openzeppelin-contracts/contracts/"
  ]
}