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

import {Initializable} from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
import {UUPSUpgradeable} from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
import {IPyth} from "@pythnetwork/pyth-sdk-solidity/IPyth.sol";
import {PythStructs} from "@pythnetwork/pyth-sdk-solidity/PythStructs.sol";
import {IHooks} from "@uniswap/v4-core/src/interfaces/IHooks.sol";
import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
import {IUnlockCallback} from "@uniswap/v4-core/src/interfaces/callback/IUnlockCallback.sol";
import {LPFeeLibrary} from "@uniswap/v4-core/src/libraries/LPFeeLibrary.sol";
import {StateLibrary} from "@uniswap/v4-core/src/libraries/StateLibrary.sol";
import {TickMath} from "@uniswap/v4-core/src/libraries/TickMath.sol";
import {FullMath} from "@uniswap/v4-core/src/libraries/FullMath.sol";
import {FixedPoint128} from "@uniswap/v4-core/src/libraries/FixedPoint128.sol";
import {PoolKey} from "@uniswap/v4-core/src/types/PoolKey.sol";
import {PoolId, PoolIdLibrary} from "@uniswap/v4-core/src/types/PoolId.sol";
import {Currency} from "@uniswap/v4-core/src/types/Currency.sol";
import {BalanceDelta, BalanceDeltaLibrary} from "@uniswap/v4-core/src/types/BalanceDelta.sol";
import {BeforeSwapDelta, toBeforeSwapDelta} from "@uniswap/v4-core/src/types/BeforeSwapDelta.sol";
import {ModifyLiquidityParams, SwapParams} from "@uniswap/v4-core/src/types/PoolOperation.sol";
import {IMaasHook} from "../interfaces/IMaasHook.sol";
import {PriceMath} from "../libraries/PriceMath.sol";
import {PerpMath} from "../libraries/PerpMath.sol";

/// @notice UUPS fee hook (all 14 flags): vault free, others base fee + toxic surcharge vs Pyth.
/// Hook takes fees (LP fee 0) and streams them + vault rewards to full-range LPs by liquidity×time.
contract MaasHook is IHooks, IMaasHook, IUnlockCallback, Initializable, UUPSUpgradeable, OwnableUpgradeable {
    using StateLibrary for IPoolManager;
    using SafeERC20 for IERC20;
    using PoolIdLibrary for PoolKey;

    uint256 public constant DURATION = 1 hours;
    uint16 constant MAX_FEE = 1000; // 10% hard ceiling for `cap`
    uint256 constant EXACT_OUT_FEE_SLOT = 0x6d616173686f6f6b2e65786163744f75744665650000000000000000000000; // transient

    struct Params {
        uint16 baseFee; // bps
        uint16 gapShare; // bps of oracle gap added when swap moves pool away from oracle
        uint16 cap; // bps
        uint16 staleFee; // bps when oracle older than maxAge / missing
        uint32 maxAge; // s
    }

    struct Pool {
        bytes32 feed;
        Currency currency0;
        Currency currency1;
        bool baseIs0;
        uint8 baseDecimals;
        uint128 totalLiquidity;
    }

    struct Stream {
        uint256 rate; // reward/s, Q128
        uint256 rewardPerLiquidity; // Q128
        uint64 finish;
        uint64 updated;
    }

    struct Account {
        uint256 paid; // rewardPerLiquidity checkpoint
        uint256 owed;
    }

    IPoolManager public immutable poolManager;

    IPyth public pyth;
    address public vault;
    Params public params;
    mapping(address => bool) public routers;
    mapping(PoolId => Pool) public pools;
    mapping(PoolId => mapping(Currency => Stream)) public streams;
    mapping(PoolId => mapping(address => uint128)) public liquidityOf;
    mapping(PoolId => mapping(address => mapping(Currency => Account))) public accounts;
    mapping(Currency => uint256) public erc20Held; // reward tokens held as ERC20 (rest are PoolManager claims)

    event Reward(PoolId indexed id, Currency indexed currency, uint256 amount);
    event Claimed(PoolId indexed id, address indexed owner, uint256 amount0, uint256 amount1);

    error NotPoolManager();
    error NotVault();
    error BadPool();
    error BadParams();
    error NotFullRange();
    error InsufficientLiquidity();
    error NotOwner();

    modifier onlyPoolManager() {
        if (msg.sender != address(poolManager)) revert NotPoolManager();
        _;
    }

    constructor(IPoolManager _poolManager) {
        poolManager = _poolManager;
        _disableInitializers();
    }

    function initialize(address owner_, IPyth pyth_) external initializer {
        __Ownable_init(owner_);
        pyth = pyth_;
        params = Params({baseFee: 3, gapShare: 5000, cap: 100, staleFee: 10, maxAge: 60});
    }

    // ---------------- owner ----------------

    function setPool(PoolKey calldata key, bytes32 feed, bool baseIs0, uint8 baseDecimals) external onlyOwner {
        if (address(key.hooks) != address(this) || key.fee != LPFeeLibrary.DYNAMIC_FEE_FLAG || feed == 0) {
            revert BadPool();
        }
        Pool storage p = pools[key.toId()];
        (p.feed, p.currency0, p.currency1, p.baseIs0, p.baseDecimals) =
        (feed, key.currency0, key.currency1, baseIs0, baseDecimals);
    }

    function setParams(Params calldata p) external onlyOwner {
        if (p.cap > MAX_FEE || p.gapShare > 10_000) revert BadParams();
        params = p;
    }

    function setVault(address v) external onlyOwner {
        vault = v;
    }

    function setRouter(address router, bool allowed) external onlyOwner {
        routers[router] = allowed;
    }

    function _authorizeUpgrade(address) internal override onlyOwner {}

    // ---------------- fees ----------------

    /// @notice Fee (bps) a non-vault swap pays now in direction `zeroForOne`.
    function feeBps(PoolId id, bool zeroForOne) public view returns (uint256 fee) {
        Params memory c = params;
        Pool storage p = pools[id];
        (bool fresh, uint256 oracle) = _oraclePrice(p.feed, c.maxAge);
        if (!fresh) return c.staleFee < c.cap ? c.staleFee : c.cap;
        (uint160 sqrtPriceX96,,,) = poolManager.getSlot0(id);
        uint256 pool = PriceMath.fromSqrtPriceX96(sqrtPriceX96, p.baseDecimals, p.baseIs0);
        fee = c.baseFee;
        // base price rises when buying base: zeroForOne buys currency1
        bool baseUp = zeroForOne != p.baseIs0;
        if (baseUp ? pool > oracle : pool < oracle) {
            uint256 gap = (pool > oracle ? pool - oracle : oracle - pool) * 10_000 / oracle;
            fee += c.gapShare * gap / 10_000;
        }
        if (fee > c.cap) fee = c.cap;
    }

    function _oraclePrice(bytes32 feed, uint256 maxAge) internal view returns (bool, uint256) {
        try pyth.getPriceUnsafe(feed) returns (PythStructs.Price memory pr) {
            if (pr.price <= 0 || pr.publishTime + maxAge < block.timestamp) return (false, 0);
            return (true, PerpMath.toPrice6(pr.price, pr.expo));
        } catch {
            return (false, 0);
        }
    }

    // vault must call PoolManager directly (sender = PoolManager caller)
    function beforeSwap(address sender, PoolKey calldata key, SwapParams calldata sp, bytes calldata)
        external
        onlyPoolManager
        returns (bytes4, BeforeSwapDelta delta, uint24)
    {
        uint256 bps = sender == vault ? 0 : feeBps(key.toId(), sp.zeroForOne);
        if (sp.amountSpecified < 0) {
            uint256 fee = uint256(-sp.amountSpecified) * bps / 10_000;
            if (fee > 0) {
                _takeFee(key, sp.zeroForOne ? key.currency0 : key.currency1, fee);
                delta = toBeforeSwapDelta(int128(int256(fee)), 0);
            }
        } else {
            assembly ("memory-safe") {
                tstore(EXACT_OUT_FEE_SLOT, bps)
            }
        }
        return (this.beforeSwap.selector, delta, LPFeeLibrary.OVERRIDE_FEE_FLAG);
    }

    function afterSwap(address, PoolKey calldata key, SwapParams calldata sp, BalanceDelta delta, bytes calldata)
        external
        onlyPoolManager
        returns (bytes4, int128)
    {
        if (sp.amountSpecified < 0) return (this.afterSwap.selector, 0);
        uint256 bps;
        assembly ("memory-safe") {
            bps := tload(EXACT_OUT_FEE_SLOT)
            tstore(EXACT_OUT_FEE_SLOT, 0)
        }
        // exact-out: unspecified = input, negative for swapper
        int128 input = sp.zeroForOne ? delta.amount0() : delta.amount1();
        uint256 fee = uint256(uint128(-input)) * bps / 10_000;
        if (fee == 0) return (this.afterSwap.selector, 0);
        _takeFee(key, sp.zeroForOne ? key.currency0 : key.currency1, fee);
        return (this.afterSwap.selector, int128(int256(fee)));
    }

    // hook delta +fee is settled by minting ERC6909 claims to self
    function _takeFee(PoolKey calldata key, Currency c, uint256 fee) internal {
        poolManager.mint(address(this), c.toId(), fee);
        _notify(key.toId(), c, fee);
    }

    // ---------------- rewards ----------------

    function notifyReward(PoolId id, Currency currency, uint256 amount) external {
        if (msg.sender != vault) revert NotVault();
        Pool storage p = pools[id];
        if (p.feed == 0 || !(currency == p.currency0 || currency == p.currency1)) revert BadPool();
        IERC20(Currency.unwrap(currency)).safeTransferFrom(msg.sender, address(this), amount);
        erc20Held[currency] += amount;
        _notify(id, currency, amount);
    }

    function earned(PoolId id, address lp, Currency currency) public view returns (uint256) {
        Account memory a = accounts[id][lp][currency];
        uint256 rpl = _rewardPerLiquidity(streams[id][currency], pools[id].totalLiquidity);
        return a.owed + FullMath.mulDiv(liquidityOf[id][lp], rpl - a.paid, FixedPoint128.Q128);
    }

    /// @notice Pays msg.sender's accrued rewards in both pool currencies.
    function claim(PoolId id) external returns (uint256 amount0, uint256 amount1) {
        Pool storage p = pools[id];
        amount0 = _checkpoint(id, msg.sender, p.currency0, true);
        amount1 = _checkpoint(id, msg.sender, p.currency1, true);
        uint256 r0 = _payErc20(p.currency0, amount0);
        uint256 r1 = _payErc20(p.currency1, amount1);
        if (r0 | r1 != 0) poolManager.unlock(abi.encode(msg.sender, p.currency0, r0, p.currency1, r1));
        emit Claimed(id, msg.sender, amount0, amount1);
    }

    function unlockCallback(bytes calldata data) external onlyPoolManager returns (bytes memory) {
        (address to, Currency c0, uint256 r0, Currency c1, uint256 r1) =
            abi.decode(data, (address, Currency, uint256, Currency, uint256));
        _payClaims(c0, to, r0);
        _payClaims(c1, to, r1);
        return "";
    }

    function _payClaims(Currency c, address to, uint256 amount) internal {
        if (amount == 0) return;
        poolManager.burn(address(this), c.toId(), amount);
        poolManager.take(c, to, amount);
    }

    // pays from ERC20 balance first; returns remainder to pay from claims
    function _payErc20(Currency c, uint256 amount) internal returns (uint256 rest) {
        uint256 held = erc20Held[c];
        uint256 e = amount < held ? amount : held;
        if (e == 0) return amount;
        erc20Held[c] = held - e;
        c.transfer(msg.sender, e);
        return amount - e;
    }

    function _notify(PoolId id, Currency c, uint256 amount) internal {
        Stream storage s = _update(id, c);
        uint256 leftover =
            s.finish > block.timestamp ? FullMath.mulDiv(s.finish - block.timestamp, s.rate, FixedPoint128.Q128) : 0;
        s.rate = FullMath.mulDiv(amount + leftover, FixedPoint128.Q128, DURATION);
        s.finish = uint64(block.timestamp + DURATION);
        emit Reward(id, c, amount);
    }

    function _update(PoolId id, Currency c) internal returns (Stream storage s) {
        s = streams[id][c];
        uint128 total = pools[id].totalLiquidity;
        // no LPs: pause stream (push finish out) so unstreamed reward is kept
        if (total == 0 && s.finish > s.updated) s.finish += uint64(block.timestamp - s.updated);
        else s.rewardPerLiquidity = _rewardPerLiquidity(s, total);
        s.updated = uint64(block.timestamp);
    }

    function _rewardPerLiquidity(Stream memory s, uint128 total) internal view returns (uint256) {
        uint256 end = block.timestamp < s.finish ? block.timestamp : s.finish;
        if (total == 0 || end <= s.updated) return s.rewardPerLiquidity;
        return s.rewardPerLiquidity + FullMath.mulDiv(end - s.updated, s.rate, total);
    }

    // settles stream + owner's account; returns owed (and zeroes it when `take`)
    function _checkpoint(PoolId id, address lp, Currency c, bool take) internal returns (uint256 owed) {
        uint256 rpl = _update(id, c).rewardPerLiquidity;
        Account storage a = accounts[id][lp][c];
        owed = a.owed + FullMath.mulDiv(liquidityOf[id][lp], rpl - a.paid, FixedPoint128.Q128);
        (a.paid, a.owed) = (rpl, take ? 0 : owed);
    }

    // ---------------- liquidity tracking ----------------

    function beforeAddLiquidity(
        address sender,
        PoolKey calldata key,
        ModifyLiquidityParams calldata mp,
        bytes calldata hookData
    ) external onlyPoolManager returns (bytes4) {
        _modify(sender, key, mp, hookData);
        return this.beforeAddLiquidity.selector;
    }

    function beforeRemoveLiquidity(
        address sender,
        PoolKey calldata key,
        ModifyLiquidityParams calldata mp,
        bytes calldata hookData
    ) external onlyPoolManager returns (bytes4) {
        _modify(sender, key, mp, hookData);
        return this.beforeRemoveLiquidity.selector;
    }

    function _modify(address sender, PoolKey calldata key, ModifyLiquidityParams calldata mp, bytes calldata hookData)
        internal
    {
        if (
            mp.tickLower != TickMath.minUsableTick(key.tickSpacing)
                || mp.tickUpper != TickMath.maxUsableTick(key.tickSpacing)
        ) revert NotFullRange();
        address lp = sender;
        if (routers[sender]) {
            // ponytail: open test router trusts hookData; pin to tx.origin. Use an owner-setting router for prod.
            lp = abi.decode(hookData, (address));
            if (lp != tx.origin) revert NotOwner();
        }
        PoolId id = key.toId();
        _checkpoint(id, lp, key.currency0, false);
        _checkpoint(id, lp, key.currency1, false);
        Pool storage p = pools[id];
        uint128 liq = liquidityOf[id][lp];
        if (mp.liquidityDelta >= 0) {
            uint128 d = uint128(uint256(mp.liquidityDelta));
            (liquidityOf[id][lp], p.totalLiquidity) = (liq + d, p.totalLiquidity + d);
        } else {
            uint256 d = uint256(-mp.liquidityDelta);
            if (d > liq) revert InsufficientLiquidity();
            (liquidityOf[id][lp], p.totalLiquidity) = (liq - uint128(d), p.totalLiquidity - uint128(d));
        }
    }

    // ---------------- unused callbacks ----------------

    function beforeInitialize(address, PoolKey calldata key, uint160) external view onlyPoolManager returns (bytes4) {
        if (key.fee != LPFeeLibrary.DYNAMIC_FEE_FLAG || pools[key.toId()].feed == 0) revert BadPool();
        return this.beforeInitialize.selector;
    }

    function afterInitialize(address, PoolKey calldata, uint160, int24) external pure returns (bytes4) {
        return this.afterInitialize.selector;
    }

    function afterAddLiquidity(
        address,
        PoolKey calldata,
        ModifyLiquidityParams calldata,
        BalanceDelta,
        BalanceDelta,
        bytes calldata
    ) external pure returns (bytes4, BalanceDelta) {
        return (this.afterAddLiquidity.selector, BalanceDeltaLibrary.ZERO_DELTA);
    }

    function afterRemoveLiquidity(
        address,
        PoolKey calldata,
        ModifyLiquidityParams calldata,
        BalanceDelta,
        BalanceDelta,
        bytes calldata
    ) external pure returns (bytes4, BalanceDelta) {
        return (this.afterRemoveLiquidity.selector, BalanceDeltaLibrary.ZERO_DELTA);
    }

    function beforeDonate(address, PoolKey calldata, uint256, uint256, bytes calldata) external pure returns (bytes4) {
        return this.beforeDonate.selector;
    }

    function afterDonate(address, PoolKey calldata, uint256, uint256, bytes calldata) external pure returns (bytes4) {
        return this.afterDonate.selector;
    }
}