# AUTOGENERATED! DO NOT EDIT! File to edit: ../src/swap.ipynb.

# %% auto 0
__all__ = ['FLAG_ALLOW_REVERT', 'ABI_DEFINITION', 'CONTRACT_BALANCE', 'CommandType', 'BridgeType', 'RoutePlanner',
           'setup_planner', 'SuperSwapDataInput', 'SuperSwapData', 'build_super_swap_data']

# %% ../src/swap.ipynb 3
import copy
from dataclasses import dataclass
from .config import XCHAIN_GAS_LIMIT_UPPERBOUND
from .quote import Quote, SuperswapQuote, pack_path
from .helpers import apply_slippage, ICACallData, hash_ICA_calls, to_bytes32, to_bytes32_str, MAX_UINT256
from .helpers import ADDRESS_ZERO
from .token import Token
from .pool import LiquidityPoolForSwap
from enum import IntEnum
from typing import List, Any, Tuple, Optional
from eth_abi import encode
from eth_abi.packed import encode_packed
from web3 import Web3
from .abi import get_abi, erc20_abi


# %% ../src/swap.ipynb 4
class CommandType(IntEnum):
    V3_SWAP_EXACT_IN = 0x00
    V3_SWAP_EXACT_OUT = 0x01
    SWEEP = 0x04
    V2_SWAP_EXACT_IN = 0x08
    V2_SWAP_EXACT_OUT = 0x09
    WRAP_ETH = 0x0b
    UNWRAP_WETH = 0x0c
    TRANSFER_FROM = 0x07
    BRIDGE_TOKEN = 0x12
    EXECUTE_CROSS_CHAIN = 0x13
    EXECUTE_SUB_PLAN = 0x21


class BridgeType(IntEnum):
    """
    Enum for different bridge types.
    """
    HYP_XERC20 = 0x01
    XVELO = 0x02

FLAG_ALLOW_REVERT = 0x80

# Define ABI types for each command
ABI_DEFINITION = {
    CommandType.V3_SWAP_EXACT_IN: [
        "address", # recipient
        "uint256", # amountIn
        "uint256", # amountOutMin
        "bytes", # path
        "bool", # payerIsUser
        "bool"  # isUni
    ],
    CommandType.V2_SWAP_EXACT_IN: [
        "address",
        "uint256",
        "uint256",
        "bytes",
        "bool",
        "bool"  # isUni
    ],
    CommandType.V2_SWAP_EXACT_OUT: [
        "address",
        "uint256",
        "uint256",
        "bytes",
        "bool",
        "bool"  # isUni
    ],
    CommandType.V3_SWAP_EXACT_OUT: [
        "address",
        "uint256",
        "uint256",
        "bytes",
        "bool",
        "bool"  # isUni
    ],
    CommandType.WRAP_ETH: [
        "address",
        "uint256"
    ],
    CommandType.UNWRAP_WETH: [
        "address",
        "uint256"
    ],
    CommandType.SWEEP: [
        "address",
        "address",
        "uint256"
    ],
    CommandType.TRANSFER_FROM: [
        "address",  # token
        "address",  # recipient
        "uint256"   # amount
    ],
    CommandType.BRIDGE_TOKEN: [
        "uint8",    # bridgeType
        "address",  # recipient
        "address",  # token
        "address",  # bridge
        "uint256",  # amount
        "uint256",  # msgFee
        "uint32",   # domain
        "bool"      # payerIsUser
    ],
    CommandType.EXECUTE_SUB_PLAN: [
        "bytes",    
        "bytes[]"
    ],
    CommandType.EXECUTE_CROSS_CHAIN: [
        "uint32",   # domain
        "address",  # icaRouter
        "bytes32",  # remoteRouter
        "bytes32",  # ism
        "bytes32",  # commitment
        "uint256",  # msgFee
        "address",  # hook
        "bytes"
    ]
}

class RoutePlanner:
    def __init__(self):
        """Initialize a new RoutePlanner"""
        self.commands = "0x"
        self.inputs: List[bytes] = []
        self.bytes_commands: List[bytes] = []

    def add_command(self, command_type: CommandType, parameters: List[Any], allow_revert = False) -> None:
        """
        Add a command to the route planner
        
        Args:
            command_type: Type of command to add
            parameters: Parameters for the command
        """
        # Get the ABI definition for this command
        abi_types = ABI_DEFINITION[command_type]
        self.inputs.append(encode(abi_types, parameters))
        # Add command byte to commands
        command_type = command_type | FLAG_ALLOW_REVERT if allow_revert else command_type
        command_hex = format(command_type, '02x')
        self.commands = self.commands + command_hex
        # TODO: figure out why we keep this alongside command_hex
        self.bytes_commands.append(command_hex)

    def get_encoded_commands(self) -> str: return self.commands
    
    def get_encoded_inputs(self) -> List[bytes]: return self.inputs

    # using this for testing
    def get_pretty_encoded_inputs(self) -> List[str]: return list(map(lambda i: "0x" + i.hex(), self.get_encoded_inputs())) 

# %% ../src/swap.ipynb 6
# Constants
CONTRACT_BALANCE = int("0x8000000000000000000000000000000000000000000000000000000000000000", 16)

def setup_planner(quote: Quote, slippage: float, account: str, router_address: str) -> RoutePlanner:
    """Setup route planner with the given quote and chain"""

    route_planner = RoutePlanner()
    min_amount_out = apply_slippage(quote.amount_out, slippage)

    tokens_come_from_contract = quote.input.amount_in == CONTRACT_BALANCE
    
    # Handle wrapped native token if needed
    if quote.from_token.wrapped_token_address:
        # When trading from native token, wrap token first
        route_planner.add_command(CommandType.WRAP_ETH, [router_address, quote.amount_in])
        tokens_come_from_contract = True
    
    # Group nodes by pool type (v2 or v3)
    grouped_nodes: List[List[Tuple[LiquidityPoolForSwap, bool]]] = []

    for node in quote.path:
        if not grouped_nodes: grouped_nodes.append([node])
        elif node[0].type < 1:
            # Current node is a v2 pool
            if float(grouped_nodes[-1][0][0].type) < 1: grouped_nodes[-1].append(node)
            else: grouped_nodes.append([node])
        else:
            # Current node is a v3 pool
            if grouped_nodes[-1][0][0].type >= 1: grouped_nodes[-1].append(node)
            else: grouped_nodes.append([node])
    
    if len(grouped_nodes) == 1:
        # All nodes belong to the same pool type
        nodes = grouped_nodes[0]
        is_v2_pool = float(nodes[0][0].type) < 1
        
        route_planner.add_command(
            CommandType.V2_SWAP_EXACT_IN if is_v2_pool else CommandType.V3_SWAP_EXACT_IN,
            [
                # Where should money go?
                router_address if quote.to_token.wrapped_token_address else account,
                quote.amount_in,
                min_amount_out,
                pack_path(nodes, for_swap=True).encoded,
                not tokens_come_from_contract,
                False, # isUni
            ]
        )
    else:
        # Mixed v2 and v3 pools
        first_batch = grouped_nodes[0]
        last_batch = grouped_nodes[-1]
        rest = grouped_nodes[1:-1]
        
        # Handle first batch
        is_first_batch_v2 = not first_batch[0][0].is_cl
        next_batch = rest[0] if rest else last_batch
        
        route_planner.add_command(
            CommandType.V2_SWAP_EXACT_IN if is_first_batch_v2 else CommandType.V3_SWAP_EXACT_IN,
            [
                router_address if is_first_batch_v2 else next_batch[0][0].lp,
                quote.amount_in,
                0,  # No expectations on min amount out for first batch
                pack_path(first_batch, for_swap=True).encoded,
                not tokens_come_from_contract,
                False,  # isUni
            ]
        )
        
        # Handle middle batches
        for idx, batch in enumerate(rest):
            is_batch_v2 = not batch[0][0].is_cl
            next_batch = rest[idx + 1] if idx + 1 < len(rest) else last_batch
            
            route_planner.add_command(
                CommandType.V2_SWAP_EXACT_IN if is_batch_v2 else CommandType.V3_SWAP_EXACT_IN,
                [
                    router_address if is_batch_v2 else next_batch[0][0].lp,
                    0 if is_batch_v2 else CONTRACT_BALANCE,
                    0,  # No expectations for middle batches
                    pack_path(batch, for_swap=True).encoded,
                    False,  # Money comes from contract
                    False,  # isUni
                ]
            )
        
        # # Handle last batch
        is_last_batch_v2 = not last_batch[0][0].is_cl
        
        route_planner.add_command(
            CommandType.V2_SWAP_EXACT_IN if is_last_batch_v2 else CommandType.V3_SWAP_EXACT_IN,
            [
                router_address if quote.to_token.wrapped_token_address else account,
                0 if is_last_batch_v2 else CONTRACT_BALANCE,
                min_amount_out,
                pack_path(last_batch, for_swap=True).encoded,
                False,  # Money comes from contract
                False,  # isUni
            ]
        )
    
    # Handle unwrapping WETH if needed
    if quote.to_token.wrapped_token_address: route_planner.add_command(CommandType.UNWRAP_WETH, [account, min_amount_out])
    
    return route_planner

# %% ../src/swap.ipynb 8
@dataclass(frozen=True)
class SuperSwapDataInput:
    from_token: Token
    to_token: Token
    from_bridge_token: Token
    to_bridge_token: Token
    account: str
    user_ICA: str
    user_ICA_balance: int
    origin_domain: int
    origin_bridge: str
    origin_hook: str
    origin_ICA_router: str
    destination_ICA_router: str
    destination_router: str
    destination_domain: int
    slippage: float
    swapper_contract_addr: str
    salt: str
    bridge_fee: int 
    xchain_fee: int
    destination_quote: Optional[Quote]

    @staticmethod
    def build(
            quote: SuperswapQuote, account: str,
            user_ICA: str, user_ICA_balance: int,
            origin_domain: int, origin_bridge: str, origin_hook: str, origin_ICA_router: str,
            destination_ICA_router: str, destination_router: str, destination_domain: int,
            slippage: float,
            swapper_contract_addr: str, salt: str,
            bridge_fee: int, xchain_fee: int
        ) -> 'SuperSwapDataInput':
        return SuperSwapDataInput(
            from_token=quote.from_token, to_token=quote.to_token, from_bridge_token=quote.from_bridge_token, to_bridge_token=quote.to_bridge_token,
            account=account, user_ICA=user_ICA, user_ICA_balance=user_ICA_balance,
            origin_domain=origin_domain, origin_bridge=origin_bridge, origin_hook=origin_hook, origin_ICA_router=origin_ICA_router,
            destination_ICA_router=destination_ICA_router, destination_router=destination_router, destination_domain=destination_domain,
            slippage=slippage,
            swapper_contract_addr=swapper_contract_addr, salt=salt,
            bridge_fee=bridge_fee, xchain_fee=xchain_fee,
            destination_quote=quote.destination_quote
        )

@dataclass(frozen=True)
class SuperSwapData:
    destination_planner: RoutePlanner
    calls: List[ICACallData]
    origin_domain: int
    salt: str
    needs_relay: bool
    commitment_hash: Optional[bytes] 

# %% ../src/swap.ipynb 10
def build_super_swap_data(i: SuperSwapDataInput) -> SuperSwapData:
    d_quote, account, slippage, swap_contract_addr = i.destination_quote, i.account, i.slippage, i.swapper_contract_addr

    # TODO: figure out if destination quote should come with tweaked amount
    if d_quote: 
        d_quote_with_max_amount_in = copy.deepcopy(d_quote)
        d_quote_with_max_amount_in.input.amount_in = CONTRACT_BALANCE
    
    destination_chain_swap_plan = setup_planner(d_quote_with_max_amount_in, slippage, account, swap_contract_addr) if d_quote else None

    swap_subplan_cmds = None

    if destination_chain_swap_plan:
        swap_subplan_cmds = destination_chain_swap_plan.get_encoded_commands()
        swap_subplan_cmds = "0x" + format(CommandType.TRANSFER_FROM, '02x') + swap_subplan_cmds.replace('0x', '')
    
    swap_subplan_cmds = bytes.fromhex(swap_subplan_cmds.replace('0x', '')) if swap_subplan_cmds else None

    destination_transfer_args = encode(ABI_DEFINITION[CommandType.TRANSFER_FROM], [
        i.to_bridge_token.token_address, i.destination_router, CONTRACT_BALANCE
    ])
    # Encode fallback transfer
    fallback_transfer_cmd = encode_packed(["bytes1"], [bytes([CommandType.TRANSFER_FROM])])
    fallback_transfer_args = encode(
        ABI_DEFINITION[CommandType.TRANSFER_FROM],
        [i.to_bridge_token.token_address, account, CONTRACT_BALANCE]
    )
    destination_inputs = [
        encode(ABI_DEFINITION[CommandType.EXECUTE_SUB_PLAN], [swap_subplan_cmds, [destination_transfer_args] + destination_chain_swap_plan.inputs]),
        encode(ABI_DEFINITION[CommandType.EXECUTE_SUB_PLAN], [fallback_transfer_cmd, [fallback_transfer_args]]),
    ] if swap_subplan_cmds else None

    calls = []

    if destination_inputs:
        # Encode Sub Plan commands
        destination_cmds = encode_packed(
            ["bytes", "bytes"],
            [
                bytes([CommandType.EXECUTE_SUB_PLAN | FLAG_ALLOW_REVERT]), 
                bytes([CommandType.EXECUTE_SUB_PLAN | FLAG_ALLOW_REVERT])
            ]
        )
        # ICA approves router to withdraw a very high amount
        calls.append(ICACallData(
            to=to_bytes32_str(i.to_bridge_token.token_address),
            value=0,
            data=Web3().eth.contract(abi=erc20_abi).encode_abi("approve", args=[i.destination_router, MAX_UINT256])
        ))
        # destination chain swap
        calls.append(ICACallData(
            to=to_bytes32_str(i.destination_router),
            value=0,
            data=Web3().eth.contract(abi=get_abi("swapper")).encode_abi("execute", args=[destination_cmds, destination_inputs])
        ))
        # reset approval back to zero
        calls.append(ICACallData(
            to=to_bytes32_str(i.to_bridge_token.token_address),
            value=0,
            data=Web3().eth.contract(abi=erc20_abi).encode_abi("approve", args=[i.destination_router, 0])
        ))

    commitment_hash = hash_ICA_calls(calls, i.salt)
    needs_relay = i.to_token.token_address != i.to_bridge_token.token_address

    destination_chain_planner = RoutePlanner()

    # TODO: confirm this is OK (destination quote check)
    # i think this works because bridge -> bridge is handled in its own way
    starts_with_bridge_token = (i.from_token.token_address == i.from_bridge_token.token_address) and d_quote is not None

    # bridge command
    destination_chain_planner.add_command(CommandType.BRIDGE_TOKEN, [
        BridgeType.HYP_XERC20,
        i.user_ICA if needs_relay else i.account,
        i.from_bridge_token.token_address,
        i.origin_bridge,
        d_quote.amount_in if starts_with_bridge_token else CONTRACT_BALANCE,
        i.bridge_fee,
        i.destination_domain,
        starts_with_bridge_token,
    ])

    if needs_relay:
        variant = 1
        hook_metadata = encode_packed(["uint16", "uint256", "uint256", "address"], [
            variant, i.xchain_fee, XCHAIN_GAS_LIMIT_UPPERBOUND, i.account 
        ])
        # destination swap command
        destination_chain_planner.add_command(CommandType.EXECUTE_CROSS_CHAIN, [
            i.destination_domain,
            i.origin_ICA_router,
            to_bytes32(i.destination_ICA_router),
            to_bytes32(ADDRESS_ZERO),
            commitment_hash,
            i.xchain_fee,
            i.origin_hook,
            hook_metadata
        ])

    return SuperSwapData(
        destination_planner=destination_chain_planner,
        calls=calls,
        origin_domain=i.origin_domain,
        salt=i.salt,
        needs_relay=needs_relay,
        commitment_hash=commitment_hash if needs_relay else None
    )
