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

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

contract SimpleStake is Ownable, ReentrancyGuard {
  using SafeERC20 for IERC20;
  IERC20 public stakingToken;
  uint256 public startTime;
  uint256 public endTime;
  uint256 public maxStakeAmount;
  uint256 public totalStaked;

  struct Stake {
    uint256 amount;
    uint256 stakeTime;
    bool unlocked;
  }

  mapping(address account => Stake[] stakes) public stakes;

  event StakeCreated(address indexed user, uint256 amount, uint256 stakeId);
  event Unlocked(address indexed user, uint256 amount);
  event PeriodUpdated(uint256 startTime, uint256 endTime);
  event MaxStakeAmountUpdated(uint256 maxStakeAmount);

  constructor(IERC20 _stakingToken, uint256 _startTime, uint256 _endTime, uint256 _maxStakeAmount) {
    require(_endTime > _startTime, "End time must be after start time");
    stakingToken = _stakingToken;
    startTime = _startTime;
    endTime = _endTime;
    maxStakeAmount = _maxStakeAmount;
  }

  function stake(uint256 _amount) public nonReentrant {
    require(block.timestamp >= startTime, "Staking has not started yet");
    require(block.timestamp <= endTime, "Staking has ended");
    require(_amount > 0, "Amount must be greater than 0");
    require(_amount <= maxStakeAmount, "Amount exceeds max stake amount");
    require(totalStaked + _amount <= maxStakeAmount, "Exceeds max stake amount");
    address user = _msgSender();
    totalStaked += _amount;
    stakingToken.safeTransferFrom(user, address(this), _amount);

    stakes[user].push(Stake(_amount, block.timestamp, false));
    emit StakeCreated(user, _amount, stakes[user].length - 1);
  }

  /**
   * @dev Unstake all stakes that have not been unlocked yet
   */
  function unstake() public nonReentrant {
    require(
      block.timestamp > endTime,
      "Unlock time not reached"
    );
    address user = _msgSender();
    require(stakes[user].length > 0, "No stakes to unstake");
   
    uint256 totalAmount = 0;
    for (uint256 i = 0; i < stakes[user].length; i++) {
      Stake storage stakeInfo = stakes[user][i];
      if (!stakeInfo.unlocked) {
        totalAmount += stakeInfo.amount;
        stakeInfo.unlocked = true;
      }
    }
    require(totalAmount > 0, "No stakes to unstake");
    totalStaked -= totalAmount;
    emit Unlocked(user, totalAmount);
    stakingToken.safeTransfer(user, totalAmount);
  }

  function setTime(uint256 _startTime, uint256 _endTime) public onlyOwner {
    require(_endTime > _startTime, "End time must be after start time");
    startTime = _startTime;
    endTime = _endTime;
    emit PeriodUpdated(_startTime, _endTime);
  }

  function setMaxStakeAmount(uint256 _maxStakeAmount) public onlyOwner {
    maxStakeAmount = _maxStakeAmount;
    emit MaxStakeAmountUpdated(_maxStakeAmount);
  }

  function userStakes(address _user) public view returns (Stake[] memory) {
    return stakes[_user];
  }
}