pragma :rubidity, "1.0.0"

import 'Upgradeable'
import 'UniswapV2Router'

contract :UniswapV2RouterWithRewards, is: [:UniswapV2Router, :Upgradeable], upgradeable: true do
  address :public, :feeAdmin
  uint256 :public, :feeBPS
  
  uint256 :public, :swapperFeePct
  uint256 :public, :stakerFeePct
  uint256 :public, :protocolFeePct
  
  mapping ({ address: :uint256 }), :public, :stakerRewardsPool
  uint256 :public, :swapperRewardsPool
  uint256 :public, :protocolFeePool
  
  mapping ({ address: mapping(address: :uint256) }), :public, :stakedLP
  mapping ({ address: mapping(address: :uint256) }), :public, :rewardDebt
  mapping ({ address: :uint256 }), :public, :totalStakedLP
  mapping ({ address: :uint256 }), :public, :accRewardsPerShare
  
  array :address, :public, :topFiveSwappers, initial_length: 5
  mapping ({ address: :uint256 }), :public, :feesGeneratedBySwapper
  mapping ({ address: :uint256 }), :public, :swapperRewards
  
  constructor(
    _factory: :address,
    _WETH: :address,
    feeBPS: :uint256,
    swapperFeePct: :uint256,
    stakerFeePct: :uint256,
    protocolFeePct: :uint256,
    feeAdmin: :address
  ) {
    updateFees(
      feeBPS: feeBPS,
      swapperFeePct: swapperFeePct,
      stakerFeePct: stakerFeePct,
      protocolFeePct: protocolFeePct,
      feeAdmin: feeAdmin
    )
    
    UniswapV2Router.constructor(_factory: _factory, _WETH: _WETH)
    Upgradeable.constructor(upgradeAdmin: msg.sender)
  }
  
  function :updateFees, {
    feeBPS: :uint256,
    swapperFeePct: :uint256,
    stakerFeePct: :uint256,
    protocolFeePct: :uint256,
    feeAdmin: :address
  }, :public do
    require(msg.sender == feeAdmin || s.feeAdmin == address(0), 'Only fee admin can update fees')
    require(feeBPS <= 10000, 'Fee cannot be greater than 100%')
    require(
      swapperFeePct + stakerFeePct + protocolFeePct == 100,
      'Fees must add up to 100%'
    )
    require(feeAdmin != address(0), 'Fee admin cannot be address(0)')
    
    s.feeBPS = feeBPS
    s.stakerFeePct = stakerFeePct
    s.swapperFeePct = swapperFeePct
    s.protocolFeePct = protocolFeePct
    
    s.feeAdmin = feeAdmin
    nil
  end
  
  function :stakeLP, { lpToken: :address, amount: :uint256 }, :public, returns: :bool do
    lpPair = UniswapV2Pair(lpToken)
    token0 = lpPair.token0()
    token1 = lpPair.token1()
    
    require(token0 == s.WETH || token1 == s.WETH, 'One of the tokens must be WETH')
    require(amount > 0, 'Amount must be greater than 0')
    
    updateStakeAmount(lpToken: lpToken, amount: amount, isStaking: true, user: msg.sender)
  
    ERC20(lpToken).transferFrom(msg.sender, address(this), amount)
  end
  
  function :unstakeLP, { lpToken: :address, amount: :uint256 }, :public, returns: :bool do
    require(s.stakedLP[msg.sender][lpToken] >= amount, 'Insufficient staked amount')
    require(amount > 0, 'Amount must be greater than 0')
    
    updateStakeAmount(lpToken: lpToken, amount: amount, isStaking: false, user: msg.sender)
    
    ERC20(lpToken).transfer(msg.sender, amount)
  end
  
  function :withdrawStakingRewards, { lpToken: :address }, :public do
    updateStakeAmount(lpToken: lpToken, amount: 0, isStaking: true, user: msg.sender)
  end
  
  function :updateStakeAmount, {
    lpToken: :address,
    amount: :uint256,
    isStaking: :bool,
    user: :address
   }, :internal do
    updateStakingRewards(lpToken: lpToken)

    pending = pendingStakingRewards(user: user, lpToken: lpToken)
    
    require(pending > 0 || amount > 0, 'Nothing to do')
    
    if pending > 0
      ERC20(s.WETH).transfer(user, pending)
    end
    
    if isStaking
      s.stakedLP[user][lpToken] += amount
      s.totalStakedLP[lpToken] += amount
    else
      s.stakedLP[user][lpToken] -= amount
      s.totalStakedLP[lpToken] -= amount
    end
    
    unScaledDebt = s.stakedLP[user][lpToken] * s.accRewardsPerShare[lpToken]
    
    s.rewardDebt[user][lpToken] = unScaledDebt.div(1.ether)
    nil
  end
  
  function :swapExactTokensForTokens, { 
    amountIn: :uint256,
    amountOutMin: :uint256,
    path: [:address],
    to: :address,
    deadline: :uint256
  }, :public, :virtual, :override, returns: [:uint256] do
    amounts = UniswapV2Router.swapExactTokensForTokens(
      amountIn: amountIn - calculateFeeAmount(amountIn),
      amountOutMin: amountOutMin,
      path: path,
      to: to,
      deadline: deadline
    )
    
    feeInWETH = chargeFeeInWETH(amountIn, path[0])
    
    lpToken = pairFor(factory, path[0], path[1])
    
    s.stakerRewardsPool[lpToken] += (feeInWETH * s.stakerFeePct).div(100)
    s.protocolFeePool += (feeInWETH * s.protocolFeePct).div(100)
    
    updateTopFiveSwappers(
      currentSwapper: msg.sender,
      currentFee: feeInWETH
    )
    
    amounts
  end
  
  function :swapTokensForExactTokens, {
    amountOut: :uint256,
    amountInMax: :uint256,
    path: [:address],
    to: :address,
    deadline: :uint256
  }, :public, :virtual, :override, returns: [:uint256] do
    amounts = UniswapV2Router.swapTokensForExactTokens(
      amountOut: amountOut + calculateFeeAmount(amountOut),
      amountInMax: amountInMax,
      path: path,
      to: to,
      deadline: deadline
    )
    
    feeInWETH = chargeFeeInWETH(amountOut, path[1])
    
    lpToken = pairFor(factory, path[0], path[1])
    
    s.stakerRewardsPool[lpToken] += (feeInWETH * s.stakerFeePct).div(100)
    s.protocolFeePool += (feeInWETH * s.protocolFeePct).div(100)
    
    updateTopFiveSwappers(
      currentSwapper: msg.sender,
      currentFee: feeInWETH
    )
    
    amounts
  end
  
  function :calculateFeeAmount, { amount: :uint256 }, :public, :virtual, returns: :uint256 do
    return (amount * s.feeBPS).div(10000)
  end
  
  function :calculateAccRewardsPerShare, { lpToken: :address }, :internal, :view, returns: :uint256 do
    accRewardsPerShare = s.accRewardsPerShare[lpToken]
    
    if s.totalStakedLP[lpToken] > 0 && s.stakerRewardsPool[lpToken] > 0
      accRewardPerShareIncrement = (s.stakerRewardsPool[lpToken] * 1.ether).div(s.totalStakedLP[lpToken])
      accRewardsPerShare += accRewardPerShareIncrement
    end
    
    return accRewardsPerShare
  end
  
  function :updateStakingRewards, { lpToken: :address }, :internal do
    s.accRewardsPerShare[lpToken] = calculateAccRewardsPerShare(lpToken: lpToken)
    s.stakerRewardsPool[lpToken] = 0
    nil
  end
  
  function :pendingStakingRewards, { user: :address, lpToken: :address }, :public, :view, returns: :uint256 do
    accRewardsPerShare = calculateAccRewardsPerShare(lpToken: lpToken)
    
    topLine = (s.stakedLP[user][lpToken] * accRewardsPerShare).div(1.ether)
    
    return topLine - s.rewardDebt[user][lpToken]
  end
  
  function :chargeFeeInWETH, {
    amount: :uint256,
    token: :address
  }, :internal, :virtual, returns: :uint256 do
    feeAmount = calculateFeeAmount(amount)
    
    if token == s.WETH
      ERC20(token).transferFrom(
        from: msg.sender,
        to: address(this),
        amount: feeAmount
      )
      
      return feeAmount
    end
    
    path = array(:address, 2)
    path[0] = token
    path[1] = s.WETH
    
    feeInWETH = UniswapV2Router.swapExactTokensForTokens(
      amountIn: feeAmount,
      amountOutMin: 0,
      path: path,
      to: address(this),
      deadline: block.timestamp + 1
    )[1]
    
    return feeInWETH
  end
  
  function :updateSwapperRewards, :internal do
    nonNullSwapperCount = 0

    for_loop(
      condition: ->(i) { i < s.topFiveSwappers.length },
      max_iterations: 5
    ) do |i|
      if s.topFiveSwappers[i] != address(0)
        nonNullSwapperCount += 1
      end
    end
    
    return if nonNullSwapperCount == 0
    
    individualSwapperReward = s.swapperRewardsPool.div(nonNullSwapperCount)
    
    for_loop(
      condition: ->(i) { i < s.topFiveSwappers.length },
      max_iterations: 5
    ) do |i|
      swapper = s.topFiveSwappers[i]
      if swapper != address(0)
        s.swapperRewards[swapper] += individualSwapperReward
      end
    end
    
    s.swapperRewardsPool = 0
    nil
  end
  
  function :withdrawProtocolRewards, { to: :address }, :public, returns: :bool do
    require(msg.sender == s.feeAdmin, "Only fee admin can withdraw protocol rewards")
    
    amount = s.protocolFeePool
    require(amount > 0, "No rewards to withdraw")
    
    s.protocolFeePool = 0
    
    ERC20(s.WETH).transfer(to, amount)
  end
  
  function :withdrawSwapperRewards, :public, returns: :bool do
    updateSwapperRewards()
    
    amount = s.swapperRewards[msg.sender]
    require(amount > 0, "No rewards to withdraw")
    
    s.swapperRewards[msg.sender] = 0
    
    ERC20(s.WETH).transfer(msg.sender, amount)
  end
  
  function :updateTopFiveSwappers, { currentSwapper: :address, currentFee: :uint256 }, :internal do
    updateSwapperRewards()
    
    s.feesGeneratedBySwapper[currentSwapper] += currentFee
    s.swapperRewardsPool += (currentFee * s.swapperFeePct).div(100)

    newTotal = s.feesGeneratedBySwapper[currentSwapper]
  
    for_loop(
      condition: ->(i) { i < s.topFiveSwappers.length },
      max_iterations: 5
    ) do |i|
      return if s.topFiveSwappers[i] == currentSwapper
    end
    
    minFee = 2 ** 256 - 1
    minIndex = 0
    
    for_loop(
      condition: -> i { i < s.topFiveSwappers.length },
      max_iterations: 5
    ) do |i|
      swapper = s.topFiveSwappers[i]
      if swapper == address(0) || s.feesGeneratedBySwapper[swapper] < minFee
        minFee = s.feesGeneratedBySwapper[swapper]
        minIndex = i
      end
    end
    
    if newTotal > minFee
      s.topFiveSwappers[minIndex] = currentSwapper
    end
    
    nil
  end
end