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

import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";

interface IAavePoolLike {
    /// @notice Standard V3 flash loan (single-asset via length-1 arrays). Use this on Aave V3.2+ pools
    ///         where `flashLoanSimple` may revert with `NotActivated()`.
    function flashLoan(
        address receiverAddress,
        address[] calldata assets,
        uint256[] calldata amounts,
        uint256[] calldata interestRateModes,
        address onBehalfOf,
        bytes calldata params,
        uint16 referralCode
    ) external;

    /// @dev Retained for compatibility with older pool deployments; prefer `flashLoan` for new integrations.
    function flashLoanSimple(
        address receiverAddress,
        address asset,
        uint256 amount,
        bytes calldata params,
        uint16 referralCode
    ) external;

    function supply(address asset, uint256 amount, address onBehalfOf, uint16 referralCode) external;

    function setUserUseReserveAsCollateral(address asset, bool useAsCollateral) external;
}

/// @dev Callback for `flashLoan` (multi-asset API).
interface IAaveFlashLoanReceiver {
    function executeOperation(
        address[] calldata assets,
        uint256[] calldata amounts,
        uint256[] calldata premiums,
        address initiator,
        bytes calldata params
    ) external returns (bool);
}

interface IAaveFlashLoanSimpleReceiver {
    function executeOperation(
        address asset,
        uint256 amount,
        uint256 premium,
        address initiator,
        bytes calldata params
    ) external returns (bool);
}

interface IAaveDODOQuotePushSwapExactIn {
    function swapExactIn(address pool, address tokenIn, uint256 amountIn, uint256 minAmountOut)
        external
        returns (uint256 amountOut);
}

interface IAaveExternalUnwinder {
    function unwind(address tokenIn, address tokenOut, uint256 amountIn, uint256 minAmountOut, bytes calldata data)
        external
        returns (uint256 amountOut);
}

interface IAaveAtomicBridgeCoordinator {
    struct CreateIntentParams {
        uint64 sourceChain;
        uint64 destinationChain;
        address assetIn;
        address assetOut;
        uint256 amountIn;
        uint256 minAmountOut;
        address recipient;
        uint256 deadline;
        bytes32 routeId;
    }

    function createIntent(CreateIntentParams calldata p) external returns (bytes32 obligationId);

    function submitCommitment(bytes32 obligationId, bytes32 settlementMode) external;

    function obligationEscrow() external view returns (address);
}

/**
 * @title AaveQuotePushFlashReceiver
 * @notice Aave V3 flash-loan receiver for the quote-push workflow:
 *         flash borrow quote (`flashLoan` single-asset) -> optional Aave collateral supply/toggle hooks ->
 *         buy PMM base -> optional atomic bridge -> unwind base externally -> repay lender, retaining any surplus.
 * @dev Collateral hooks apply to **this contract's** Aave position (`onBehalfOf` = address(this)). Pre-fund the
 *      receiver with collateral tokens before `supplyBeforeSwap` steps, or supply from flash proceeds in later txs.
 */
contract AaveQuotePushFlashReceiver is IAaveFlashLoanSimpleReceiver, IAaveFlashLoanReceiver, Ownable {
    using SafeERC20 for IERC20;

    address public immutable pool;

    enum CollateralHook {
        BeforeSwap,
        AfterSwap,
        BeforeUnwind
    }

    struct CollateralSupplyStep {
        address asset;
        uint256 amount;
    }

    struct CollateralToggleStep {
        address asset;
        bool useAsCollateral;
    }

    struct CollateralParams {
        CollateralSupplyStep[] supplyBeforeSwap;
        CollateralToggleStep[] toggleBeforeSwap;
        CollateralToggleStep[] toggleAfterSwap;
        CollateralToggleStep[] toggleBeforeUnwind;
    }

    struct QuotePushParams {
        address integration;
        address pmmPool;
        address baseToken;
        address externalUnwinder;
        uint256 minOutPmm;
        uint256 minOutUnwind;
        bytes unwindData;
        AtomicBridgeParams atomicBridge;
        CollateralParams collateral;
    }

    struct AtomicBridgeParams {
        address coordinator;
        uint64 sourceChain;
        uint64 destinationChain;
        address destinationAsset;
        uint256 bridgeAmount;
        uint256 minDestinationAmount;
        address destinationRecipient;
        uint256 destinationDeadline;
        bytes32 routeId;
        bytes32 settlementMode;
        bool submitCommitment;
    }

    error UntrustedPool();
    error UntrustedInitiator();
    error BadParams();
    error InsufficientToRepay();
    error InvalidAtomicBridge();
    error NothingToSweep();

    event QuotePushExecuted(
        address indexed quoteToken,
        address indexed baseToken,
        uint256 borrowedAmount,
        uint256 premium,
        uint256 baseOut,
        uint256 unwindOut,
        uint256 surplus
    );
    event AtomicBridgeTriggered(
        bytes32 indexed obligationId,
        address indexed coordinator,
        address indexed destinationRecipient,
        uint256 bridgeAmount,
        uint256 minDestinationAmount
    );
    event TokenSwept(address indexed token, address indexed to, uint256 amount);
    event SurplusSwept(address indexed token, address indexed to, uint256 amount, uint256 reserveRetained);
    event CollateralSupplied(address indexed asset, uint256 amount);
    event CollateralToggled(address indexed asset, bool useAsCollateral, CollateralHook hook);

    /// @notice Empty `collateral` block for quote-push scripts that do not use Aave position hooks.
    function emptyCollateralParams() external pure returns (CollateralParams memory params) {
        return params;
    }

    /// @notice Non-zero when this deployment supports `QuotePushParams.collateral` atomic hooks.
    function collateralHooksVersion() external pure returns (uint256) {
        return 1;
    }

    constructor(address pool_, address initialOwner) Ownable(initialOwner) {
        if (pool_ == address(0) || initialOwner == address(0)) revert BadParams();
        pool = pool_;
    }

    function flashQuotePush(address asset, uint256 amount, QuotePushParams calldata params) external {
        address[] memory assets = new address[](1);
        uint256[] memory amts = new uint256[](1);
        uint256[] memory modes = new uint256[](1);
        assets[0] = asset;
        amts[0] = amount;
        modes[0] = 0;
        // `onBehalfOf` must not be zero on some Aave V3.2+ deployments (mapping lookups revert NotActivated).
        IAavePoolLike(pool).flashLoan(
            address(this), assets, amts, modes, address(this), abi.encode(address(this), params), 0
        );
    }

    function quoteSurplusBalance(address quoteToken, uint256 reserveRetained) public view returns (uint256 surplus) {
        uint256 quoteBal = IERC20(quoteToken).balanceOf(address(this));
        if (quoteBal > reserveRetained) {
            surplus = quoteBal - reserveRetained;
        }
    }

    function sweepQuoteSurplus(address quoteToken, address to, uint256 reserveRetained)
        external
        onlyOwner
        returns (uint256 amount)
    {
        if (quoteToken == address(0) || to == address(0)) revert BadParams();
        amount = quoteSurplusBalance(quoteToken, reserveRetained);
        if (amount == 0) revert NothingToSweep();
        IERC20(quoteToken).safeTransfer(to, amount);
        emit SurplusSwept(quoteToken, to, amount, reserveRetained);
    }

    function sweepToken(address token, address to, uint256 amount) external onlyOwner {
        if (token == address(0) || to == address(0) || amount == 0) revert BadParams();
        IERC20(token).safeTransfer(to, amount);
        emit TokenSwept(token, to, amount);
    }

    function executeOperation(
        address[] calldata assets,
        uint256[] calldata amounts,
        uint256[] calldata premiums,
        address initiator,
        bytes calldata params
    ) external returns (bool) {
        if (msg.sender != pool) revert UntrustedPool();
        if (assets.length != 1 || amounts.length != 1 || premiums.length != 1) revert BadParams();
        _executeQuotePush(assets[0], amounts[0], premiums[0], initiator, params);
        return true;
    }

    function executeOperation(
        address asset,
        uint256 amount,
        uint256 premium,
        address initiator,
        bytes calldata params
    ) external returns (bool) {
        if (msg.sender != pool) revert UntrustedPool();
        _executeQuotePush(asset, amount, premium, initiator, params);
        return true;
    }

    function _executeQuotePush(
        address asset,
        uint256 amount,
        uint256 premium,
        address initiator,
        bytes calldata params
    ) internal {
        (address expectedInitiator, QuotePushParams memory p) = abi.decode(params, (address, QuotePushParams));
        if (initiator != expectedInitiator) revert UntrustedInitiator();
        if (
            p.integration == address(0) || p.pmmPool == address(0) || p.baseToken == address(0)
                || p.externalUnwinder == address(0)
        ) revert BadParams();
        if (p.baseToken == asset) revert BadParams();

        _applyCollateralSupplies(p.collateral.supplyBeforeSwap);
        _applyCollateralToggles(p.collateral.toggleBeforeSwap, CollateralHook.BeforeSwap);

        uint256 baseOut = _swapQuoteForBase(asset, amount, p.integration, p.pmmPool, p.minOutPmm);

        _applyCollateralToggles(p.collateral.toggleAfterSwap, CollateralHook.AfterSwap);

        uint256 baseBal = IERC20(p.baseToken).balanceOf(address(this));
        if (p.atomicBridge.coordinator != address(0)) {
            _triggerAtomicBridge(p.baseToken, baseBal, p.atomicBridge);
        }

        _applyCollateralToggles(p.collateral.toggleBeforeUnwind, CollateralHook.BeforeUnwind);

        uint256 unwindOut = _unwindBaseIntoQuote(p.baseToken, asset, p.externalUnwinder, p.minOutUnwind, p.unwindData);
        uint256 surplus = _approveRepayment(asset, amount + premium);
        emit QuotePushExecuted(asset, p.baseToken, amount, premium, baseOut, unwindOut, surplus);
    }

    function _swapQuoteForBase(address asset, uint256 amount, address integration, address pmmPool, uint256 minOutPmm)
        internal
        returns (uint256 baseOut)
    {
        IERC20(asset).forceApprove(integration, amount);
        baseOut = IAaveDODOQuotePushSwapExactIn(integration).swapExactIn(pmmPool, asset, amount, minOutPmm);
    }

    function _unwindBaseIntoQuote(
        address baseToken,
        address quoteToken,
        address externalUnwinder,
        uint256 minOutUnwind,
        bytes memory unwindData
    ) internal returns (uint256 unwindOut) {
        uint256 remainingBase = IERC20(baseToken).balanceOf(address(this));
        IERC20(baseToken).forceApprove(externalUnwinder, remainingBase);
        unwindOut =
            IAaveExternalUnwinder(externalUnwinder).unwind(baseToken, quoteToken, remainingBase, minOutUnwind, unwindData);
    }

    function _applyCollateralSupplies(CollateralSupplyStep[] memory steps) internal {
        uint256 n = steps.length;
        for (uint256 i = 0; i < n; ++i) {
            CollateralSupplyStep memory step = steps[i];
            if (step.asset == address(0) || step.amount == 0) continue;
            IERC20(step.asset).forceApprove(pool, step.amount);
            IAavePoolLike(pool).supply(step.asset, step.amount, address(this), 0);
            emit CollateralSupplied(step.asset, step.amount);
        }
    }

    function _applyCollateralToggles(CollateralToggleStep[] memory steps, CollateralHook hook) internal {
        uint256 n = steps.length;
        for (uint256 i = 0; i < n; ++i) {
            CollateralToggleStep memory step = steps[i];
            if (step.asset == address(0)) continue;
            IAavePoolLike(pool).setUserUseReserveAsCollateral(step.asset, step.useAsCollateral);
            emit CollateralToggled(step.asset, step.useAsCollateral, hook);
        }
    }

    function _approveRepayment(address quoteToken, uint256 need) internal returns (uint256 surplus) {
        IERC20 quote = IERC20(quoteToken);
        uint256 quoteBal = quote.balanceOf(address(this));
        if (quoteBal < need) revert InsufficientToRepay();
        surplus = quoteBal - need;
        quote.forceApprove(pool, need);
    }

    function _triggerAtomicBridge(address baseToken, uint256 baseBal, AtomicBridgeParams memory atomicBridge) internal {
        if (
            atomicBridge.destinationAsset == address(0) || atomicBridge.destinationRecipient == address(0)
                || atomicBridge.bridgeAmount == 0 || atomicBridge.bridgeAmount > baseBal
                || atomicBridge.destinationDeadline <= block.timestamp
        ) revert InvalidAtomicBridge();

        address escrowAddress = IAaveAtomicBridgeCoordinator(atomicBridge.coordinator).obligationEscrow();
        IERC20(baseToken).forceApprove(escrowAddress, atomicBridge.bridgeAmount);
        bytes32 obligationId = IAaveAtomicBridgeCoordinator(atomicBridge.coordinator).createIntent(
            IAaveAtomicBridgeCoordinator.CreateIntentParams({
                sourceChain: atomicBridge.sourceChain,
                destinationChain: atomicBridge.destinationChain,
                assetIn: baseToken,
                assetOut: atomicBridge.destinationAsset,
                amountIn: atomicBridge.bridgeAmount,
                minAmountOut: atomicBridge.minDestinationAmount,
                recipient: atomicBridge.destinationRecipient,
                deadline: atomicBridge.destinationDeadline,
                routeId: atomicBridge.routeId
            })
        );
        if (atomicBridge.submitCommitment) {
            IAaveAtomicBridgeCoordinator(atomicBridge.coordinator).submitCommitment(
                obligationId, atomicBridge.settlementMode
            );
        }
        emit AtomicBridgeTriggered(
            obligationId,
            atomicBridge.coordinator,
            atomicBridge.destinationRecipient,
            atomicBridge.bridgeAmount,
            atomicBridge.minDestinationAmount
        );
    }
}