SimpleDEX Stage 2

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

import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import "@openzeppelin/contracts/token/ERC20/ERC20.sol";

/**
 * @title SimpleDEX - Stage 2: Додаємо функціонал
 * @dev Додаємо removeLiquidity(), getPrice(), getAmountOut()
 *
 * Етап 2 (20 хвилин): Покращуємо функціонал DEX
 * - Видалення ліквідності
 * - Розрахунок цін
 * - Попередній розрахунок swap
 * - ~140 рядків коду
 */
contract SimpleDEX_Stage2 is ERC20 {
    // Токени пари
    IERC20 public immutable token0;
    IERC20 public immutable token1;

    // Резерви токенів
    uint256 public reserve0;
    uint256 public reserve1;

    // Комісія 0.3%
    uint256 public constant FEE_PERCENT = 30;
    uint256 public constant FEE_DENOMINATOR = 10000;

    // Мінімальна ліквідність для запобігання маніпуляціям
    uint256 private constant MINIMUM_LIQUIDITY = 1000;

    // Events
    event LiquidityAdded(address indexed provider, uint256 amount0, uint256 amount1, uint256 liquidity);
    event LiquidityRemoved(address indexed provider, uint256 amount0, uint256 amount1, uint256 liquidity);
    event Swap(address indexed user, address indexed tokenIn, uint256 amountIn, uint256 amountOut);

    constructor(address _token0, address _token1) ERC20("SimpleDEX LP", "SDEX-LP") {
        require(_token0 != address(0) && _token1 != address(0), "Invalid tokens");
        require(_token0 != _token1, "Identical tokens");

        token0 = IERC20(_token0);
        token1 = IERC20(_token1);
    }

    /**
     * @dev Додати ліквідність в пул
     */
    function addLiquidity(uint256 amount0, uint256 amount1) external returns (uint256 liquidity) {
        require(amount0 > 0 && amount1 > 0, "Invalid amounts");

        // Переказуємо токени від користувача
        token0.transferFrom(msg.sender, address(this), amount0);
        token1.transferFrom(msg.sender, address(this), amount1);

        // Розраховуємо LP токени
        uint256 _totalSupply = totalSupply();

        if (_totalSupply == 0) {
            // Перший постачальник - блокуємо MINIMUM_LIQUIDITY
            liquidity = sqrt(amount0 * amount1);
            _mint(address(0), MINIMUM_LIQUIDITY);
            liquidity = liquidity - MINIMUM_LIQUIDITY;
        } else {
            // Наступні постачальники
            liquidity = min(
                (amount0 * _totalSupply) / reserve0,
                (amount1 * _totalSupply) / reserve1
            );
        }

        require(liquidity > 0, "Insufficient liquidity");

        // Мінтимо LP токени
        _mint(msg.sender, liquidity);

        // Оновлюємо резерви
        reserve0 += amount0;
        reserve1 += amount1;

        emit LiquidityAdded(msg.sender, amount0, amount1, liquidity);
    }

    /**
     * @dev Видалити ліквідність з пулу
     */
    function removeLiquidity(uint256 liquidity) external returns (uint256 amount0, uint256 amount1) {
        require(liquidity > 0, "Invalid liquidity");
        require(balanceOf(msg.sender) >= liquidity, "Insufficient LP tokens");

        uint256 _totalSupply = totalSupply();

        // Розраховуємо суму токенів для повернення
        amount0 = (liquidity * reserve0) / _totalSupply;
        amount1 = (liquidity * reserve1) / _totalSupply;

        require(amount0 > 0 && amount1 > 0, "Insufficient amounts");

        // Спалюємо LP токени
        _burn(msg.sender, liquidity);

        // Переказуємо токени назад користувачу
        token0.transfer(msg.sender, amount0);
        token1.transfer(msg.sender, amount1);

        // Оновлюємо резерви
        reserve0 -= amount0;
        reserve1 -= amount1;

        emit LiquidityRemoved(msg.sender, amount0, amount1, liquidity);
    }

    /**
     * @dev Обміняти токени
     */
    function swap(address tokenIn, uint256 amountIn, uint256 minAmountOut)
        external
        returns (uint256 amountOut)
    {
        require(amountIn > 0, "Invalid amount");
        require(tokenIn == address(token0) || tokenIn == address(token1), "Invalid token");

        // Визначаємо напрямок свапу
        bool isToken0 = tokenIn == address(token0);
        (IERC20 tokenInContract, IERC20 tokenOutContract, uint256 reserveIn, uint256 reserveOut) =
            isToken0 ? (token0, token1, reserve0, reserve1) : (token1, token0, reserve1, reserve0);

        // Переказуємо вхідні токени
        tokenInContract.transferFrom(msg.sender, address(this), amountIn);

        // Розраховуємо вихід з комісією (x*y=k формула)
        uint256 amountInWithFee = amountIn * (FEE_DENOMINATOR - FEE_PERCENT);
        amountOut = (amountInWithFee * reserveOut) / (reserveIn * FEE_DENOMINATOR + amountInWithFee);

        require(amountOut >= minAmountOut, "Slippage exceeded");
        require(amountOut > 0, "Insufficient output");

        // Переказуємо вихідні токени
        tokenOutContract.transfer(msg.sender, amountOut);

        // Оновлюємо резерви
        if (isToken0) {
            reserve0 += amountIn;
            reserve1 -= amountOut;
        } else {
            reserve1 += amountIn;
            reserve0 -= amountOut;
        }

        emit Swap(msg.sender, tokenIn, amountIn, amountOut);
    }

    /**
     * @dev Розрахувати скільки токенів отримаєте (без виконання свапу)
     */
    function getAmountOut(address tokenIn, uint256 amountIn)
        external
        view
        returns (uint256 amountOut)
    {
        require(amountIn > 0, "Invalid amount");
        require(tokenIn == address(token0) || tokenIn == address(token1), "Invalid token");

        bool isToken0 = tokenIn == address(token0);
        (uint256 reserveIn, uint256 reserveOut) = isToken0 ? (reserve0, reserve1) : (reserve1, reserve0);

        uint256 amountInWithFee = amountIn * (FEE_DENOMINATOR - FEE_PERCENT);
        amountOut = (amountInWithFee * reserveOut) / (reserveIn * FEE_DENOMINATOR + amountInWithFee);
    }

    /**
     * @dev Отримати ціну token0 в одиницях token1
     */
    function getPrice0() external view returns (uint256) {
        require(reserve0 > 0 && reserve1 > 0, "No liquidity");
        return (reserve1 * 1e18) / reserve0;
    }

    /**
     * @dev Отримати ціну token1 в одиницях token0
     */
    function getPrice1() external view returns (uint256) {
        require(reserve0 > 0 && reserve1 > 0, "No liquidity");
        return (reserve0 * 1e18) / reserve1;
    }

    // Helper functions
    function sqrt(uint256 y) internal pure returns (uint256 z) {
        if (y > 3) {
            z = y;
            uint256 x = y / 2 + 1;
            while (x < z) {
                z = x;
                x = (y / x + x) / 2;
            }
        } else if (y != 0) {
            z = 1;
        }
    }

    function min(uint256 a, uint256 b) internal pure returns (uint256) {
        return a < b ? a : b;
    }
}