// SPDX-License-Identifier: GPL-3.0-or-later
pragma solidity ^0.8.19;

import { Kraiken } from "./Kraiken.sol";
import { IERC20 } from "@openzeppelin/token/ERC20/ERC20.sol";
import { ERC20Permit } from "@openzeppelin/token/ERC20/extensions/ERC20Permit.sol";
import { IERC20Metadata } from "@openzeppelin/token/ERC20/extensions/IERC20Metadata.sol";
import { SafeERC20 } from "@openzeppelin/token/ERC20/utils/SafeERC20.sol";
import { Math } from "@openzeppelin/utils/math/Math.sol";

error ExceededAvailableStake(address receiver, uint256 stakeWanted, uint256 availableStake);
error TooMuchSnatch(address receiver, uint256 stakeWanted, uint256 availableStake, uint256 smallestShare);

/**
 * @title Stake Contract for Kraiken Token
 * @notice This contract manages the staking positions for the Kraiken token, allowing users to stake tokens
 *         in exchange for a share of the total supply. Stakers can set and adjust tax rates on their stakes,
 *         which affect the Universal Basic Income (UBI) paid from the tax pool.
 *
 *         The contract handles:
 *         - Creation of staking positions with specific tax rates.
 *         - Snatching of existing positions under certain conditions to consolidate stakes.
 *         - Calculation and payment of taxes based on stake duration and tax rate.
 *         - Adjustment of tax rates with protections against griefing through rapid changes.
 *         - Exiting of positions, either partially or fully, returning the staked assets to the owner.
 *
 *         Tax rates and staking positions are adjustable, with a mechanism to prevent snatch-grieving by
 *         enforcing a minimum tax payment duration.
 *
 * @dev Self-assessed tax implementation:
 *      - Continuous auction mechanism
 *      - Self-assessed valuations create prediction market
 *      - Tax collection and redistribution through UBI
 */
contract Stake {
    using Math for uint256;

    // the offset between the "precision" of the representation of shares and assets
    // see https://docs.openzeppelin.com/contracts/4.x/erc4626 for reason and details
    uint256 internal DECIMAL_OFFSET = 5 + 2;
    // only 20% of the total KRAIKEN supply can be staked.
    uint256 internal constant MAX_STAKE = 20; // 20% of KRAIKEN supply
    uint256 internal constant TAX_FLOOR_DURATION = 60 * 60 * 24 * 3; //this duration is the minimum basis for fee calculation, regardless of actual holding time.
    // the tax rates are discrete to prevent users from snatching by micro incroments of tax
    uint256[] public TAX_RATES =
        [1, 3, 5, 8, 12, 18, 24, 30, 40, 50, 60, 80, 100, 130, 180, 250, 320, 420, 540, 700, 920, 1200, 1600, 2000, 2600, 3400, 4400, 5700, 7500, 9700];
    // this is the base for the values in the array above: e.g. 1/100 = 1%
    uint256 internal constant TAX_RATE_BASE = 100;
    /**
     * @dev Attempted to deposit more assets than the max amount for `receiver`.
     */

    error TaxTooLow(address receiver, uint64 taxRateWanted, uint64 taxRateMet, uint256 positionId);
    error StakeTooLow(address receiver, uint256 assets, uint256 minStake);
    error NoPermission(address requester, address owner);
    error PositionNotFound(uint256 positionId, address requester);

    event PositionCreated(uint256 indexed positionId, address indexed owner, uint256 kraikenDeposit, uint256 share, uint32 taxRate);
    event PositionTaxPaid(uint256 indexed positionId, address indexed owner, uint256 taxPaid, uint256 newShares, uint256 taxRate);
    event PositionRateHiked(uint256 indexed positionId, address indexed owner, uint256 newTaxRate);
    event PositionShrunk(uint256 indexed positionId, address indexed owner, uint256 newShares, uint256 kraikenPayout);
    event PositionRemoved(uint256 indexed positionId, address indexed owner, uint256 kraikenPayout);

    struct StakingPosition {
        uint256 share;
        address owner;
        uint32 creationTime;
        uint32 lastTaxTime;
        uint32 taxRate; // index of TAX_RATES array
    }

    Kraiken private immutable kraiken;
    address private immutable taxReceiver;

    uint256 public immutable totalSupply;
    uint256 public outstandingStake;
    uint256 public nextPositionId;

    mapping(uint256 => StakingPosition) public positions;

    // Array to keep track of total shares at each tax rate
    uint256[] public totalSharesAtTaxRate;

    /// @notice Initializes the stake contract with references to the Kraiken contract and sets the initial position ID.
    /// @param _kraiken Address of the Kraiken contract which this Stake contract interacts with.
    /// @dev Sets up the total supply based on the decimals of the Kraiken token plus a fixed offset.
    constructor(address _kraiken, address _taxReceiver) {
        kraiken = Kraiken(_kraiken);
        taxReceiver = _taxReceiver;
        totalSupply = 10 ** (kraiken.decimals() + DECIMAL_OFFSET);
        // start counting somewhere
        nextPositionId = 654_321;
        // Initialize totalSharesAtTaxRate array
        totalSharesAtTaxRate = new uint256[](TAX_RATES.length);
    }

    function authorizedStake() private view returns (uint256) {
        return totalSupply * MAX_STAKE / 100;
    }

    /// @dev Internal function to calculate and pay taxes for a position, adjusting shares and handling position liquidation if necessary.
    function _payTax(uint256 positionId, StakingPosition storage pos, uint256 taxFloorDuration) private {
        // existance of position should be checked before
        // ihet = Implied Holding Expiry Timestamp
        uint256 ihet = (block.timestamp - pos.creationTime < taxFloorDuration) ? pos.creationTime + taxFloorDuration : block.timestamp;
        uint256 elapsedTime = ihet - pos.lastTaxTime;
        uint256 assetsBefore = sharesToAssets(pos.share);
        uint256 taxAmountDue = assetsBefore * TAX_RATES[pos.taxRate] * elapsedTime / (365 * 24 * 60 * 60) / TAX_RATE_BASE;
        if (taxAmountDue >= assetsBefore) {
            // can not pay more tax than value of position
            taxAmountDue = assetsBefore;
        }
        if (assetsBefore - taxAmountDue > 0) {
            // if something left over, update storage
            uint256 shareAfterTax = assetsToShares(assetsBefore - taxAmountDue);
            uint256 deltaShare = pos.share - shareAfterTax;
            totalSharesAtTaxRate[pos.taxRate] -= deltaShare;
            outstandingStake -= deltaShare;
            pos.share = shareAfterTax;
            pos.lastTaxTime = uint32(block.timestamp);
            emit PositionTaxPaid(positionId, pos.owner, taxAmountDue, shareAfterTax, pos.taxRate);
        } else {
            // if nothing left over, liquidate position
            totalSharesAtTaxRate[pos.taxRate] -= pos.share;
            outstandingStake -= pos.share;
            emit PositionTaxPaid(positionId, pos.owner, taxAmountDue, 0, pos.taxRate);
            emit PositionRemoved(positionId, pos.owner, 0);
            delete pos.owner;
            delete pos.creationTime;
            delete pos.share;
        }
        SafeERC20.safeTransfer(kraiken, taxReceiver, taxAmountDue);
    }

    /// @dev Internal function to close a staking position, transferring the remaining Kraiken tokens back to the owner after tax payment.
    function _exitPosition(uint256 positionId, StakingPosition storage pos) private {
        totalSharesAtTaxRate[pos.taxRate] -= pos.share;
        outstandingStake -= pos.share;
        address owner = pos.owner;
        uint256 assets = sharesToAssets(pos.share);
        emit PositionRemoved(positionId, owner, assets);
        delete pos.owner;
        delete pos.creationTime;
        delete pos.share;
        SafeERC20.safeTransfer(kraiken, owner, assets);
    }

    /// @dev Internal function to reduce the size of a staking position by a specified number of shares, transferring the corresponding Kraiken tokens to the owner.
    function _shrinkPosition(uint256 positionId, StakingPosition storage pos, uint256 sharesToTake) private {
        require(sharesToTake < pos.share, "position too small");
        uint256 assets = sharesToAssets(sharesToTake);
        pos.share -= sharesToTake;
        totalSharesAtTaxRate[pos.taxRate] -= sharesToTake;
        outstandingStake -= sharesToTake;
        emit PositionShrunk(positionId, pos.owner, pos.share, assets);
        SafeERC20.safeTransfer(kraiken, pos.owner, assets);
    }

    /// @notice Converts Kraiken token assets to shares of the total staking pool.
    /// @param assets Number of Kraiken tokens to convert.
    /// @return Number of shares corresponding to the input assets based on the current total supply of Kraiken tokens.
    function assetsToShares(uint256 assets) public view returns (uint256) {
        return assets.mulDiv(totalSupply, kraiken.totalSupply(), Math.Rounding.Down);
    }

    /// @notice Converts shares of the total staking pool back to Kraiken token assets.
    /// @param shares Number of shares to convert.
    /// @return The equivalent number of Kraiken tokens for the given shares.
    function sharesToAssets(uint256 shares) public view returns (uint256) {
        return shares.mulDiv(kraiken.totalSupply(), totalSupply, Math.Rounding.Down);
    }

    /// @notice Creates a new staking position by potentially snatching shares from existing positions.
    /// @param assets Amount of Kraiken tokens to convert into a staking position.
    /// @param receiver Address that will own the new staking position.
    /// @param taxRate The initial tax rate for the new staking position.
    /// @param positionsToSnatch Array of position IDs that the new position will replace by snatching.
    /// @return positionId The ID of the newly created staking position.
    /// @dev Handles staking logic, including tax rate validation and position merging or dissolving.
    function snatch(uint256 assets, address receiver, uint32 taxRate, uint256[] calldata positionsToSnatch) public returns (uint256 positionId) {
        // check lower boundary
        uint256 sharesWanted = assetsToShares(assets);
        {
            // check that position size is at least minStake
            // to prevent excessive fragmentation, increasing snatch cost
            uint256 minStake = kraiken.minStake();
            if (assets < minStake) {
                revert StakeTooLow(receiver, assets, minStake);
            }
        }
        require(taxRate < TAX_RATES.length, "tax rate out of bounds");

        uint256 smallestPositionShare = totalSupply;
        uint256 availableStake = authorizedStake() - outstandingStake;

        if (positionsToSnatch.length >= 2) {
            // run through all but last positions to snatch
            for (uint256 i = 0; i < positionsToSnatch.length - 1; i++) {
                StakingPosition storage pos = positions[positionsToSnatch[i]];
                if (pos.creationTime == 0) {
                    revert PositionNotFound(positionsToSnatch[i], receiver);
                }
                // check that tax lower
                if (taxRate <= pos.taxRate) {
                    revert TaxTooLow(receiver, taxRate, pos.taxRate, positionsToSnatch[i]);
                }
                if (pos.share < smallestPositionShare) {
                    smallestPositionShare = pos.share;
                }
                // dissolve position
                _payTax(positionsToSnatch[i], pos, 0);
                _exitPosition(positionsToSnatch[i], pos);
            }
        }
        availableStake = authorizedStake() - outstandingStake;

        if (positionsToSnatch.length > 0) {
            // handle last position, either shrink or snatch
            uint256 index = positionsToSnatch.length - 1;
            StakingPosition storage lastPos = positions[positionsToSnatch[index]];
            if (lastPos.creationTime == 0) {
                revert PositionNotFound(positionsToSnatch[index], receiver);
            }
            // check that tax lower
            if (taxRate <= lastPos.taxRate) {
                revert TaxTooLow(receiver, taxRate, lastPos.taxRate, positionsToSnatch[index]);
            }
            if (lastPos.share < smallestPositionShare) {
                smallestPositionShare = lastPos.share;
            }
            // dissolve position
            _payTax(positionsToSnatch[index], lastPos, TAX_FLOOR_DURATION);
            if (availableStake > sharesWanted) {
                revert TooMuchSnatch(receiver, sharesWanted, availableStake, smallestPositionShare);
            }
            uint256 lastSharesNeeded = sharesWanted - availableStake;
            if (lastSharesNeeded > lastPos.share * 80 / 100) {
                _exitPosition(positionsToSnatch[index], lastPos);
            } else {
                _shrinkPosition(positionsToSnatch[index], lastPos, lastSharesNeeded);
            }
        }
        availableStake = authorizedStake() - outstandingStake;

        if (sharesWanted > availableStake) {
            revert ExceededAvailableStake(receiver, sharesWanted, availableStake);
        }

        // avoid greeving where more positions are freed than needed.
        if (availableStake - sharesWanted > smallestPositionShare) {
            revert TooMuchSnatch(receiver, sharesWanted, availableStake, smallestPositionShare);
        }

        // transfer
        SafeERC20.safeTransferFrom(kraiken, msg.sender, address(this), assets);

        // mint
        positionId = nextPositionId++;
        StakingPosition storage sp = positions[positionId];
        sp.share = sharesWanted;
        sp.owner = receiver;
        sp.lastTaxTime = uint32(block.timestamp);
        sp.creationTime = uint32(block.timestamp);
        sp.taxRate = taxRate;

        totalSharesAtTaxRate[taxRate] += sharesWanted;
        outstandingStake += sharesWanted;
        emit PositionCreated(positionId, sp.owner, assets, sp.share, sp.taxRate);
    }

    /// @notice Combines an ERC20 permit operation with the snatch function, allowing a staking position creation in one transaction.
    /// @param assets Number of Kraiken tokens to stake.
    /// @param receiver Address that will own the new staking position.
    /// @param taxRate The initial tax rate for the new staking position.
    /// @param positionsToSnatch Array of position IDs that the new position will replace by snatching.
    /// @param deadline Time until which the permit is valid.
    /// @param v, r, s Components of the signature for the permit.
    /// @return positionId The ID of the newly created staking position.
    function permitAndSnatch(
        uint256 assets,
        address receiver,
        uint32 taxRate,
        uint256[] calldata positionsToSnatch,
        // address owner,
        // address spender,
        // uint256 value,
        uint256 deadline,
        uint8 v,
        bytes32 r,
        bytes32 s
    )
        external
        returns (uint256 positionId)
    {
        ERC20Permit(address(kraiken)).permit(receiver, address(this), assets, deadline, v, r, s);
        return snatch(assets, receiver, taxRate, positionsToSnatch);
    }

    /// @notice Changes the tax rate of an existing staking position.
    /// @param positionId The ID of the staking position to update.
    /// @param taxRate The new tax rate to apply to the position.
    /// @dev Ensures that the tax rate change is valid and applies the minimum tax based on the TAX_FLOOR_DURATION.
    function changeTax(uint256 positionId, uint32 taxRate) external {
        require(taxRate < TAX_RATES.length, "tax rate out of bounds");
        StakingPosition storage pos = positions[positionId];
        if (pos.creationTime == 0) {
            revert PositionNotFound(positionId, msg.sender);
        }
        if (pos.owner != msg.sender) {
            revert NoPermission(msg.sender, pos.owner);
        }
        // to prevent snatch-and-change grieving attack, pay TAX_FLOOR_DURATION
        require(taxRate > pos.taxRate, "tax too low to snatch");
        _payTax(positionId, pos, 0);
        totalSharesAtTaxRate[pos.taxRate] -= pos.share;
        totalSharesAtTaxRate[taxRate] += pos.share;
        pos.taxRate = taxRate;
        emit PositionRateHiked(positionId, pos.owner, taxRate);
    }

    /// @notice Allows the owner of a staking position to exit, returning the staked assets.
    /// @param positionId The ID of the staking position to exit.
    /// @dev Pays the due taxes based on the TAX_FLOOR_DURATION and returns the remaining assets to the position owner.
    function exitPosition(uint256 positionId) external {
        StakingPosition storage pos = positions[positionId];
        if (pos.owner != msg.sender) {
            revert NoPermission(msg.sender, pos.owner);
        }
        if (pos.creationTime == 0) {
            revert PositionNotFound(positionId, msg.sender);
        }
        // to prevent snatch-and-exit grieving attack, pay TAX_FLOOR_DURATION
        _payTax(positionId, pos, TAX_FLOOR_DURATION);
        _exitPosition(positionId, pos);
    }

    /// @notice Manually triggers the tax payment for a specified staking position.
    /// @param positionId The ID of the staking position for which to pay taxes.
    /// @dev Calculates and pays the tax due, possibly adjusting the position's share count.
    function payTax(uint256 positionId) external {
        StakingPosition storage pos = positions[positionId];
        if (pos.creationTime == 0) {
            revert PositionNotFound(positionId, msg.sender);
        }
        _payTax(positionId, pos, 0);
    }

    /// @notice Calculates the Tax that is due to be paid on specific positoin
    /// @param positionId The ID of the staking position for which to pay taxes.
    /// @param taxFloorDuration if a minimum holding duration is applied to the position this value is > 0 in seconds.
    /// @dev Calculates the tax due.
    function taxDue(uint256 positionId, uint256 taxFloorDuration) public view returns (uint256 amountDue) {
        StakingPosition storage pos = positions[positionId];
        // ihet = Implied Holding Expiry Timestamp
        uint256 ihet = (block.timestamp - pos.creationTime < taxFloorDuration) ? pos.creationTime + taxFloorDuration : block.timestamp;
        uint256 elapsedTime = ihet - pos.lastTaxTime;
        uint256 assetsBefore = sharesToAssets(pos.share);
        amountDue = assetsBefore * TAX_RATES[pos.taxRate] * elapsedTime / (365 * 24 * 60 * 60) / TAX_RATE_BASE;
    }

    /// @return averageTaxRate A number between 0 and 1e18 indicating the average tax rate.
    function getAverageTaxRate() external view returns (uint256 averageTaxRate) {
        // Compute average tax rate weighted by shares
        averageTaxRate = 0;
        if (outstandingStake > 0) {
            for (uint256 i = 0; i < TAX_RATES.length; i++) {
                averageTaxRate += TAX_RATES[i] * totalSharesAtTaxRate[i];
            }
            averageTaxRate = averageTaxRate / outstandingStake;
            // normalize tax rate
            averageTaxRate = averageTaxRate * 1e18 / TAX_RATES[TAX_RATES.length - 1];
        }
    }

    /// @notice Computes the percentage of Kraiken staked from outstanding Stake and authorized Stake.
    /// @return percentageStaked A number between 0 and 1e18 indicating the percentage of Kraiken supply staked.
    function getPercentageStaked() external view returns (uint256 percentageStaked) {
        percentageStaked = (outstandingStake * 1e18) / authorizedStake();
    }
}