From 61d5a855bedefcd4e1a9c8c9dd1d56e9f6e4efcb Mon Sep 17 00:00:00 2001
From: naman mohnani <mohnaninaman@gmail.com>
Date: Wed, 27 Nov 2024 12:00:21 +0530
Subject: [PATCH] created governance token

---
 src/GovToken.sol             | 14 ++++++++++++++
 src/JITHook.sol              | 24 ++++++++++++++++++++----
 src/interfaces/IStrategy.sol |  2 +-
 3 files changed, 35 insertions(+), 5 deletions(-)
 create mode 100644 src/GovToken.sol

diff --git a/src/GovToken.sol b/src/GovToken.sol
new file mode 100644
index 0000000..a6b41b9
--- /dev/null
+++ b/src/GovToken.sol
@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.26;
+
+import {ERC20} from "solmate/src/tokens/ERC20.sol";
+import {Owned} from "solmate/src/auth/Owned.sol";
+
+contract GovToken is ERC20, Owned {
+
+    constructor(string memory name, string memory symbol) ERC20(name, symbol, 18) Owned(msg.sender){}
+
+    function mint(address to, uint256 amount) external override onlyOwner() {
+        _mint(to, amount);
+    }
+}
\ No newline at end of file
diff --git a/src/JITHook.sol b/src/JITHook.sol
index eb589f8..8c0e751 100644
--- a/src/JITHook.sol
+++ b/src/JITHook.sol
@@ -17,6 +17,7 @@ import {IPositionManager} from "v4-periphery/src/interfaces/IPositionManager.sol
 import {Actions} from "v4-periphery/src/libraries/Actions.sol";
 import {LiquidityAmounts} from "v4-core/test/utils/LiquidityAmounts.sol";
 import {TickMath} from "v4-core/src/libraries/TickMath.sol";
+import {GovToken} from "./GovToken.sol";
 
 contract JITHook is BaseHook {
     using PoolIdLibrary for PoolId;
@@ -32,6 +33,9 @@ contract JITHook is BaseHook {
     uint256 public currentPositionId;
     uint256 public currentActiveStrategyId;
 
+    mapping(PoolKey => GovToken) public govTokens;
+    // mapping(PoolId => GovToken) public govTokens; 
+
     constructor(
         IPoolManager _manager,
         address _strategiesController,
@@ -56,7 +60,7 @@ contract JITHook is BaseHook {
         return
             Hooks.Permissions({
                 beforeInitialize: false,
-                afterInitialize: false,
+                afterInitialize: true,
                 beforeAddLiquidity: false,
                 afterAddLiquidity: false,
                 beforeRemoveLiquidity: false,
@@ -72,6 +76,13 @@ contract JITHook is BaseHook {
             });
     }
 
+    function afterInitialize(address, PoolKey calldata key, uint160, int24) external virtual returns (bytes4) {
+        string memory name = string.concat("Governance token: ", key.currency0.symbol, "-", key.currency1.symbol);
+        string memory symbol = string.concat(key.currency0.symbol, "-", key.currency1.symbol);
+        govTokens[key] = new GovToken(name, symbol);
+        return this.afterInitialize.selector;
+    }
+
     // if swap is big => remove lquidity from external protocol and add liquidity
     function beforeSwap(
         address,
@@ -150,9 +161,11 @@ contract JITHook is BaseHook {
 
         if(amount0 > 0) {
             _depositToStrategy(currentActiveStrategyId, key.currency0, amount0);
+            govTokens[key].mint(msg.sender, amount0);
         } 
         if(amount1 > 0) {
             _depositToStrategy(currentActiveStrategyId, key.currency1, amount1);
+            govTokens[key].mint(msg.sender, amount1);
         }
 
         // todo emit event
@@ -160,7 +173,10 @@ contract JITHook is BaseHook {
 
     // funds transferred to small LPs from external protocol
     function withdraw(uint256 amount0, uint256 amount1, PoolKey calldata key) external {
-        uint256 withdrawAmount = _withdrawFromStrategy(currentActiveStrategyId);
+        address token0 = Currency.unwrap(key.currency0);
+        address token1 = Currency.unwrap(key.currency1);
+        uint256 withdrawAmount = _withdrawFromStrategy(currentActiveStrategyId, token0, amount0);
+        uint256 withdrawAmountToken1 = _withdrawFromStrategy(currentActiveStrategyId, token1, amount1); 
 
         // todo calculate token amount for the user, and transfer to user 
         // todo multiple user wil store funds here, when one user call this only his liquidty should be removed and transferred to user NOT ALL
@@ -177,8 +193,8 @@ contract JITHook is BaseHook {
         IStrategy(controller.getStrategyAddress(_id)).deposit(_token, _amount);
     }
 
-    function _withdrawFromStrategy(uint256 _id) internal returns (uint256) {
-        return IStrategy(controller.getStrategyAddress(_id)).withdraw();
+    function _withdrawFromStrategy(uint256 _id, address token, uint256 amount) internal returns (uint256) {
+        return IStrategy(controller.getStrategyAddress(_id)).withdraw(token, amount);
     }
 
     function _addLiquidityToPool(
diff --git a/src/interfaces/IStrategy.sol b/src/interfaces/IStrategy.sol
index f9116d8..2d90a02 100644
--- a/src/interfaces/IStrategy.sol
+++ b/src/interfaces/IStrategy.sol
@@ -4,6 +4,6 @@ pragma solidity 0.8.26;
 interface IStrategy {
     
     function deposit(address token, uint256 amount) external;
-    function withdraw() external returns (uint256);
+    function withdraw(address token, uint256 amount) external returns (uint256);
     function redeemYeild() external;
 }
\ No newline at end of file