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

# %% auto 0
__all__ = ['ABI_DEFINITION', 'CONTRACT_BALANCE_FOR_V3_SWAPS', 'CommandType', 'RoutePlanner', 'setup_planner']

# %% ../src/swap.ipynb 3
from .quote import Quote, pack_path
from .helpers import apply_slippage
from .pool import LiquidityPoolForSwap
from enum import IntEnum
from typing import List, Any, Tuple
from eth_abi import encode

# %% ../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

# Define ABI types for each command
ABI_DEFINITION = {
    CommandType.V3_SWAP_EXACT_IN: [
        "address",
        "uint256",
        "uint256",
        "bytes",
        "bool"
    ],
    CommandType.V2_SWAP_EXACT_IN: [
        "address",
        "uint256",
        "uint256",
        "(address,address,bool)[]",
        "bool"
    ],
    CommandType.V2_SWAP_EXACT_OUT: [
        "address",
        "uint256",
        "uint256",
        "(address,address,bool)[]",
        "bool"
    ],
    CommandType.V3_SWAP_EXACT_OUT: [
        "address",
        "uint256",
        "uint256",
        "bytes",
        "bool"
    ],
    CommandType.WRAP_ETH: [
        "address",
        "uint256"
    ],
    CommandType.UNWRAP_WETH: [
        "address",
        "uint256"
    ],
    CommandType.SWEEP: [
        "address",
        "address",
        "uint256"
    ]
}

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

    def add_command(self, command_type: CommandType, parameters: List[Any]) -> 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_hex = format(command_type, '02x')
        self.commands = self.commands + 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 5
# Constants
CONTRACT_BALANCE_FOR_V3_SWAPS = int("0x8000000000000000000000000000000000000000000000000000000000000000", 16)

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

    # By default money comes from contract
    tokens_come_from_contract = False
    
    # 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]]] = []
    
    # helpers for getting from and to token addresses
    def from_token_address(node: Tuple[LiquidityPoolForSwap, bool]) -> str:
        pool, reversed = node
        return pool.token0_address if not reversed else pool.token1_address
    def to_token_address(node: Tuple[LiquidityPoolForSwap, bool]) -> str:
        pool, reversed = node
        return pool.token1_address if not reversed else pool.token0_address


    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,
                [
                    # from, to, stable
                    (from_token_address(n), to_token_address(n), n[0].is_stable) for n in nodes
                ] if is_v2_pool else pack_path(nodes).encoded,
                not tokens_come_from_contract,
            ]
        )
    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
                [
                    # from, to, stable
                    (from_token_address(n), to_token_address(n), n[0].is_stable) for n in first_batch
                ] if is_first_batch_v2 else pack_path(first_batch).encoded,
                not tokens_come_from_contract,
            ]
        )
        
        # 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_FOR_V3_SWAPS,
                    0,  # No expectations for middle batches
                    [
                        (from_token_address(n), to_token_address(n), n[0].is_stable) for n in batch
                    ] if is_batch_v2 else pack_path(batch).encoded,
                    False,  # Money comes from contract
                ]
            )
        
        # # 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_FOR_V3_SWAPS,
                min_amount_out,
                [
                    (from_token_address(n), to_token_address(n), n[0].is_stable) for n in last_batch
                ] if is_last_batch_v2 else pack_path(last_batch).encoded,
                False,  # Money comes from contract
            ]
        )
    
    # 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
