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

# %% auto 0
__all__ = ['original_format_batched_response', 'T', 'safe_format_batched_response', 'require_context', 'require_async_context',
           'CommonChain', 'AsyncChain', 'Chain', 'OPChainCommon', 'AsyncOPChain', 'OPChain', 'BaseChainCommon',
           'AsyncBaseChain', 'BaseChain']

# %% ../src/chains.ipynb 3
import os, asyncio
from concurrent.futures import ThreadPoolExecutor, as_completed
from functools import wraps, lru_cache
from async_lru import alru_cache
from cachetools import cached, TTLCache
from typing import List, TypeVar, Callable, Optional, Tuple, Dict
from web3 import Web3, HTTPProvider, AsyncWeb3, AsyncHTTPProvider, Account
from web3.eth.async_eth import AsyncContract
from web3.eth import Contract
from web3.manager import RequestManager, RequestBatcher
from .config import ChainSettings, make_op_chain_settings, make_base_chain_settings
from .helpers import normalize_address, MAX_UINT256, float_to_uint256, apply_slippage, get_future_timestamp, ADDRESS_ZERO, chunk, Pair
from .helpers import find_all_paths, time_it, atime_it
from .abi import get_abi
from .token import Token
from .pool import LiquidityPool, LiquidityPoolForSwap, LiquidityPoolEpoch
from .price import Price
from .deposit import Deposit
from .quote import QuoteInput, Quote
from .swap import setup_planner

# %% ../src/chains.ipynb 4
# monkey patching how web3 handles errors in batched requests
# re: https://github.com/ethereum/web3.py/issues/3657
original_format_batched_response = RequestManager._format_batched_response
def safe_format_batched_response(*args):
    try: return original_format_batched_response(*args)
    except Exception as e: return e
RequestManager._format_batched_response = safe_format_batched_response

# %% ../src/chains.ipynb 6
T = TypeVar('T')

def require_context(f: Callable[..., T]) -> Callable[..., T]:
    @wraps(f)
    def wrapper(self: 'CommonChain', *args, **kwargs) -> T:
        if not self._in_context: raise RuntimeError("Chain methods can only be accessed within 'async with' block")
        return f(self, *args, **kwargs)
    return wrapper

def require_async_context(f: Callable[..., T]) -> Callable[..., T]:
    @wraps(f)
    async def wrapper(self: 'CommonChain', *args, **kwargs) -> T:
        if not self._in_context: raise RuntimeError("Chain methods can only be accessed within 'async with' block")
        return await f(self, *args, **kwargs)
    return wrapper


class CommonChain:
    @property
    def account(self) -> Account: return self.web3.eth.account.from_key(os.getenv("SUGAR_PK"))

    pools: Optional[List[LiquidityPool]] = None
    pools_for_swap: Optional[List[LiquidityPoolForSwap]] = None

    def __init__(self, settings: ChainSettings, **kwargs):
        self.settings, self._in_context = settings, False

        if "pools" in kwargs: self.pools = kwargs["pools"]
        if "pools_for_swap" in kwargs: self.pools_for_swap = kwargs["pools_for_swap"] 
    
    def prepare_set_token_allowance_contract(self, token: Token, contract_wrapper):
        ERC20_ABI = [{
            "name": "approve",
            "type": "function",
            "constant": False,
            "inputs": [{"name": "spender", "type": "address"}, {"name": "amount", "type": "uint256"}],
            "outputs": [{"name": "", "type": "bool"}]
        }]
        return contract_wrapper(address=token.wrapped_token_address or token.token_address, abi=ERC20_ABI)
    
    def prepare_tokens(self, tokens: List[Tuple], listed_only: bool) -> List[Token]:
        native = Token.make_native_token(self.settings.native_token_symbol, self.settings.wrapped_native_token_addr, self.settings.native_token_decimals)
        ts = list(map(lambda t: Token.from_tuple(t), tokens))
        return [native] + (list(filter(lambda t: t.listed, ts)) if listed_only else ts)
    
    def _prepare_prices(self, tokens: List[Token], rates: List[int]) -> Dict[str, int]:
        # token_address => normalized rate
        result = {}
        # rates are returned multiplied by eth decimals + the difference in decimals to eth
        # we want them all normalized to 18 decimals
        for cnt, rate in enumerate(rates):
            t, eth_decimals = tokens[cnt], self.settings.native_token_decimals
            if t.decimals == eth_decimals: nr = rate
            elif t.decimals < eth_decimals: nr = rate // (10 ** (eth_decimals - t.decimals))
            else: nr = rate * (10 ** (t.decimals - eth_decimals))
            result[t.token_address] = nr
        return result

    def prepare_prices(self, tokens: List[Token], prices: List[int]) -> List[Price]:
        """Get prices for tokens in target stable token"""
        eth_decimals = self.settings.native_token_decimals
        # all rates in EHT: token => rate
        rates_in_eth = self._prepare_prices(tokens, prices)
        eth_rate, usd_rate = rates_in_eth[self.settings.native_token_symbol], rates_in_eth[self.settings.stable_token_addr]
        # this gives us the price of 1 eth in usd with 18 decimals precision
        eth_usd_price = (eth_rate * 10 ** eth_decimals) // usd_rate
        # finally convert to prices in terms of stable
        return [Price(token=t, price=(rates_in_eth[t.token_address] * eth_usd_price // 10 ** eth_decimals) / 10 ** eth_decimals) for t in tokens]
    
    def prepare_pools(self, pools: List[Tuple], tokens: List[Token], prices: List[Price]) -> List[LiquidityPool]:
        tokens, prices = {t.token_address: t for t in tokens}, {price.token.token_address: price for price in prices}
        return list(filter(lambda p: p is not None, map(lambda p: LiquidityPool.from_tuple(p, tokens, prices), pools)))
    
    def prepare_pools_for_swap(self, pools: List[Tuple]) -> List[LiquidityPoolForSwap]:
        return list(map(lambda p: LiquidityPoolForSwap.from_tuple(p), pools))

    def prepare_pool_epochs(self, epochs: List[Tuple], pools: List[LiquidityPool], tokens: List[Token], prices: List[Price]) -> List[LiquidityPoolEpoch]:
        tokens, prices, pools = {t.token_address: t for t in tokens}, {price.token.token_address: price for price in prices}, {p.lp: p for p in pools}
        return list(map(lambda p: LiquidityPoolEpoch.from_tuple(p, pools, tokens, prices), epochs))
    
    def filter_pools_for_swap(self, pools: List[LiquidityPoolForSwap], from_token: Token, to_token: Token) -> List[LiquidityPoolForSwap]:
        match_tokens = set(self.settings.connector_tokens_addrs + [from_token.token_address, to_token.token_address])
        return list(filter(lambda p: p.token0_address in match_tokens or p.token1_address in match_tokens, pools))
    
    def paths_to_pools(self, pools: List[LiquidityPoolForSwap], paths: List[List[Tuple]]) -> List[LiquidityPoolForSwap]:
        pools_dict = {p.lp: p for p in pools}
        return [list(map(lambda p: pools_dict[p[2]], path)) for path in paths]

    def prepare_quote_batch(self, from_token: Token, to_token: Token, batcher: RequestBatcher, pools: List[List[LiquidityPoolForSwap]], amount_in: int, paths: List[List[Tuple]]):
        inputs = []
        for i, path in enumerate(paths):
            p = [(p, p.token0_address != path[i][0]) for i, p in enumerate(pools[i])]
            q = QuoteInput(from_token=from_token, to_token=to_token, amount_in=amount_in, path=p)
            batcher.add(self.quoter.functions.quoteExactInput(q.route.encoded, amount_in))
            inputs.append(q)
        return batcher, inputs

    def prepare_quotes(self, quote_inputs: List[QuoteInput], responses):
        if len(responses) != len(quote_inputs): raise ValueError(f"Number of responses {len(responses)} does not match number of quote inputs {len(quote_inputs)}")
        quotes = []
        for i, r in enumerate(responses):
            if isinstance(r, Exception): continue
            else: quotes.append(Quote(input=quote_inputs[i], amount_out=r[0]))
        return quotes
    
    def get_paths_for_quote(self, from_token: Token, to_token: Token, pools: List[LiquidityPoolForSwap], exclude_tokens: List[str]) -> List[List[Tuple]]:
        exclude_tokens_set = set(map(lambda t: normalize_address(t), exclude_tokens))

        if from_token.token_address in exclude_tokens: exclude_tokens_set.remove(from_token.token_address)
        if to_token.token_address in exclude_tokens: exclude_tokens_set.remove(to_token.token_address)

        pairs = [Pair(p.token0_address, p.token1_address, p.lp) for p in pools]
        paths = find_all_paths(pairs, from_token.wrapped_token_address or from_token.token_address, to_token.wrapped_token_address or to_token.token_address)
        # filter out paths with excluded tokens
        return list(filter(lambda p: len(set(map(lambda t: t[0], p)) & exclude_tokens_set) == 0, paths))

    def get_pool_paginator(self, batch_size = 5) -> List[List[Tuple]]:
        limit, upper_bound = self.settings.pool_page_size, self.settings.pools_count_upper_bound
        return chunk(list(map(lambda x: (x, limit), list(range(0, upper_bound, limit)))), batch_size)
    
    def prepare_price_batcher(self, tokens: List[Token], batch: RequestBatcher):
        batches = chunk(tokens, self.settings.price_batch_size)
        for b in batches:
            batch.add(self.prices.functions.getManyRatesToEthWithCustomConnectors(
                list(map(lambda t: t.wrapped_token_address or t.token_address, b)),
                False, # use wrappers
                self.settings.connector_tokens_addrs,
                10 # threshold_filter
            ))
        return batch        

# %% ../src/chains.ipynb 8
class AsyncChain(CommonChain):
    web3: AsyncWeb3
    sugar: AsyncContract
    slipstream: AsyncContract
    router: AsyncContract
    prices: AsyncContract
    quoter: AsyncContract
    swapper: AsyncContract


    async def __aenter__(self):
        """Async context manager entry"""
        self._in_context = True
        self.web3 = AsyncWeb3(AsyncHTTPProvider(self.settings.rpc_uri))
        self.sugar = self.web3.eth.contract(address=self.settings.sugar_contract_addr, abi=get_abi("sugar"))
        self.sugar_rewards = self.web3.eth.contract(address=self.settings.sugar_rewards_contract_addr, abi=get_abi("sugar_rewards"))
        self.slipstream = self.web3.eth.contract(address=self.settings.slipstream_contract_addr, abi=get_abi("slipstream"))
        self.prices = self.web3.eth.contract(address=self.settings.price_oracle_contract_addr, abi=get_abi("price_oracle"))
        self.router = self.web3.eth.contract(address=self.settings.router_contract_addr, abi=get_abi("router"))
        self.quoter = self.web3.eth.contract(address=self.settings.quoter_contract_addr, abi=get_abi("quoter"))
        self.swapper = self.web3.eth.contract(address=self.settings.swapper_contract_addr, abi=get_abi("swapper"))

        # set up caching for price oracle
        self._get_prices = alru_cache(ttl=self.settings.pricing_cache_timeout_seconds)(self._get_prices)

        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Async context manager exit"""
        self._in_context = False
        await self.web3.provider.disconnect()
        return None

    async def apaginate(self, f: Callable):
        async def process_batch(batch: List[Tuple]):
            async with self.web3.batch_requests() as batcher:
                for offset, limit in batch: batcher.add(f(limit, offset))
                return sum(await batcher.async_execute(), [])
        return sum(await asyncio.gather(*[process_batch(batch) for batch in self.get_pool_paginator()]), [])

    @require_async_context
    @alru_cache(maxsize=None)
    async def get_all_tokens(self, listed_only: bool = False) -> List[Token]:
        def get_tokens(limit, offset): return self.sugar.functions.tokens(limit, offset, ADDRESS_ZERO, [])
        return self.prepare_tokens(await self.apaginate(get_tokens), listed_only)

    async def _get_prices(self, tokens: Tuple[Token]):
        async with self.web3.batch_requests() as batch:
            batch = self.prepare_price_batcher(tokens=list(tokens), batch=batch)
            return sum(await batch.async_execute(), [])

    @require_async_context
    async def get_prices(self, tokens: List[Token]) -> List[Price]:
        """Get prices for tokens in target stable token"""
        return self.prepare_prices(tokens, await self._get_prices(tuple(tokens)))

    @alru_cache(maxsize=None)
    async def get_raw_pools(self, for_swaps: bool):
        return await self.apaginate(self.sugar.functions.forSwaps if for_swaps else self.sugar.functions.all)
    
    @require_async_context
    async def get_pools(self, for_swaps: bool = False) -> List[LiquidityPool]:
        pools = await self.get_raw_pools(for_swaps)
        if not for_swaps:
            tokens = await self.get_all_tokens()
            return self.prepare_pools(pools, tokens, await self.get_prices(tokens))
        else: return self.prepare_pools_for_swap(pools)
    
    @require_async_context
    @alru_cache(maxsize=None)
    async def get_pool_by_address(self, address: str) -> Optional[LiquidityPool]:
        try:
            p = await self.sugar.functions.byAddress(address).call()
        except: return None
        tokens = await self.get_all_tokens(listed_only=False)
        return self.prepare_pools([p], tokens, await self.get_prices(tokens))[0]

    @require_async_context
    @alru_cache(maxsize=None)
    async def get_pool_epochs(self, lp: str, offset: int = 0, limit: int = 10) -> List[LiquidityPoolEpoch]:
        tokens, pools = await self.get_all_tokens(listed_only=False), await self.get_pools()
        prices = await self.get_prices(tokens)
        r = await self.sugar_rewards.functions.epochsByAddress(limit, offset, normalize_address(lp)).call()
        return self.prepare_pool_epochs(r, pools, tokens, prices)

    @require_async_context
    @alru_cache(maxsize=None)
    async def get_latest_pool_epochs(self) -> List[LiquidityPoolEpoch]:
        tokens, pools = await self.get_all_tokens(listed_only=False), await self.get_pools()
        prices = await self.get_prices(tokens)
        return self.prepare_pool_epochs(await self.apaginate(self.sugar_rewards.functions.epochsLatest), pools, tokens, prices)
    
    @require_async_context
    async def get_pools_for_swaps(self) -> List[LiquidityPoolForSwap]: return await self.get_pools(for_swaps=True)

    @require_async_context
    async def _get_quotes_for_paths(self, from_token: Token, to_token: Token, amount_in: int, pools: List[LiquidityPoolForSwap], paths: List[List[Tuple]]) -> List[Optional[Quote]]:
        path_pools = self.paths_to_pools(pools, paths)
        async with self.web3.batch_requests() as batch:
            batch, inputs = self.prepare_quote_batch(from_token, to_token, batch, path_pools, amount_in, paths)
            return self.prepare_quotes(inputs, await batch.async_execute())

    @require_async_context
    async def get_quote(self, from_token: Token, to_token: Token, amount: float, filter_quotes: Optional[Callable[[Quote], bool]] = None) -> Optional[Quote]:
        amount_in = float_to_uint256(amount, decimals=from_token.decimals)
        pools = self.filter_pools_for_swap(from_token=from_token, to_token=to_token, pools=await self.get_pools_for_swaps())
        paths = self.get_paths_for_quote(from_token, to_token, pools, self.settings.excluded_tokens_addrs)
        quotes = sum(await asyncio.gather(*[self._get_quotes_for_paths(from_token, to_token, amount_in, pools, paths) for paths in chunk(paths, 500)]), [])
        quotes = list(filter(lambda q: q is not None, quotes))
        if filter_quotes is not None: quotes = list(filter(filter_quotes, quotes))
        return max(quotes, key=lambda q: q.amount_out) if len(quotes) > 0 else None
    
    @require_async_context
    async def swap(self, from_token: Token, to_token: Token, amount: float, slippage: Optional[float] = None):
        q = await self.get_quote(from_token, to_token, amount)
        if not q: raise ValueError("No quotes found")
        return await self.swap_from_quote(q, slippage=slippage)
        
    @require_async_context
    async def swap_from_quote(self, quote: Quote, slippage: Optional[float] = None):
        swapper_contract_addr, from_token = self.settings.swapper_contract_addr, quote.from_token
        planner = setup_planner(quote=quote, slippage=slippage if slippage is not None else self.settings.swap_slippage, account=self.account.address, router_address=swapper_contract_addr)
        await self.set_token_allowance(from_token, swapper_contract_addr, quote.input.amount_in)
        value = quote.input.amount_in if from_token.wrapped_token_address else 0
        return await self.sign_and_send_tx(self.swapper.functions.execute(*[planner.commands, planner.inputs]), value=value)

    @require_async_context
    async def set_token_allowance(self, token: Token, addr: str, amount: int):
        token_contract = self.prepare_set_token_allowance_contract(token, self.web3.eth.contract)
        return await self.sign_and_send_tx(token_contract.functions.approve(addr, amount))

    @require_async_context
    async def check_token_allowance(self, token: Token, addr: str) -> int:
        ERC20_ABI = [{
            "name": "allowance",
            "type": "function",
            "constant": True,
            "inputs": [{"name": "owner", "type": "address"}, {"name": "spender", "type": "address"}],
            "outputs": [{"name": "", "type": "uint256"}]
        }]
        token_contract = self.web3.eth.contract(address=token.wrapped_token_address or token.token_address, abi=ERC20_ABI)
        return await token_contract.functions.allowance(self.account.address, addr).call()

    @require_async_context
    async def sign_and_send_tx(self, tx, value: int = 0, wait: bool = True):
        spender = self.account.address
        tx = await tx.build_transaction({ 'from': spender, 'value': value, 'nonce': await self.web3.eth.get_transaction_count(spender) })
        signed_tx = self.account.sign_transaction(tx)
        tx_hash = await self.web3.eth.send_raw_transaction(signed_tx.raw_transaction)
        return await self.web3.eth.wait_for_transaction_receipt(tx_hash) if wait else tx_hash
    
    @require_async_context
    async def deposit(self, deposit: Deposit, delay_in_minutes: float = 30, slippage: float = 0.01):
        amount_token0, pool, router_contract_addr = deposit.amount_token0, deposit.pool, self.settings.router_contract_addr
        print(f"gonna deposit {amount_token0} {pool.token0.symbol} into {pool.symbol} from {self.account.address}")
        [token0_amount, token1_amount, _] = await self.router.functions.quoteAddLiquidity(
            pool.token0.token_address,
            pool.token1.token_address,
            pool.is_stable,
            pool.factory,
            float_to_uint256(amount_token0, pool.token0.decimals),
            MAX_UINT256
        ).call()
        print(f"Quote: {pool.token0.symbol} {token0_amount / 10 ** pool.token0.decimals} -> {pool.token1.symbol} {token1_amount / 10 ** pool.token1.decimals}")

        # set up allowance for both tokens
        print(f"setting up allowance for {pool.token0.symbol}")
        await self.set_token_allowance(pool.token0, router_contract_addr, token0_amount)

        print(f"setting up allowance for {pool.token1.symbol}")
        await self.set_token_allowance(pool.token1, router_contract_addr, token1_amount)

        # check allowances
        token0_allowance = await self.check_token_allowance(pool.token0, router_contract_addr)
        token1_allowance = await self.check_token_allowance(pool.token1, router_contract_addr)

        print(f"allowances: {token0_allowance}, {token1_allowance}")

        # adding liquidity

        # if token 0 is native, use addLiquidityETH instead of standard addLiquidity
        if pool.token0.token_address == self.settings.wrapped_native_token_addr:
            params = [
                pool.token1.token_address,
                pool.is_stable,
                token1_amount,
                apply_slippage(token1_amount, slippage),
                apply_slippage(token0_amount, slippage),
                self.account.address,
                get_future_timestamp(delay_in_minutes)
            ]
            print(f"adding liquidity with params: {params}")
            return await self.sign_and_send_tx(self.router.functions.addLiquidityETH(*params), value=token0_amount)
        
        # token 1 is native, use addLiquidityETH instead of standard addLiquidity
        if pool.token1.token_address == self.settings.wrapped_native_token_addr:
            params = [
                pool.token0.token_address,
                pool.is_stable,
                token0_amount,
                apply_slippage(token0_amount, slippage),
                apply_slippage(token1_amount, slippage),
                self.account.address,
                get_future_timestamp(delay_in_minutes)
            ]
            print(f"adding liquidity with params: {params}")
            return await self.sign_and_send_tx(self.router.functions.addLiquidityETH(*params), value=token1_amount)

        params = [
            pool.token0.token_address,
            pool.token1.token_address,
            pool.is_stable,
            token0_amount,
            token1_amount,
            apply_slippage(token0_amount, slippage),
            apply_slippage(token1_amount, slippage),
            self.account.address,
            get_future_timestamp(delay_in_minutes)
        ]

        print(f"adding liquidity with params: {params}")

        return await self.sign_and_send_tx(self.router.functions.addLiquidity(*params))

# %% ../src/chains.ipynb 10
class Chain(CommonChain):
    web3: Web3
    sugar: Contract
    router: Contract
    slipstream: Contract
    prices: Contract
    quoter: Contract
    swapper: Contract

    def __enter__(self):
        """Sync context manager entry"""
        self._in_context = True
        self.web3 = Web3(HTTPProvider(self.settings.rpc_uri))
        self.sugar = self.web3.eth.contract(address=self.settings.sugar_contract_addr, abi=get_abi("sugar"))
        self.sugar_rewards = self.web3.eth.contract(address=self.settings.sugar_rewards_contract_addr, abi=get_abi("sugar_rewards"))
        self.slipstream = self.web3.eth.contract(address=self.settings.slipstream_contract_addr, abi=get_abi("slipstream"))
        self.prices = self.web3.eth.contract(address=self.settings.price_oracle_contract_addr, abi=get_abi("price_oracle"))
        self.router = self.web3.eth.contract(address=self.settings.router_contract_addr, abi=get_abi("router"))
        self.quoter = self.web3.eth.contract(address=self.settings.quoter_contract_addr, abi=get_abi("quoter"))
        self.swapper = self.web3.eth.contract(address=self.settings.swapper_contract_addr, abi=get_abi("swapper"))

        # set up caching for price oracle
        self._get_prices = cached(TTLCache(ttl=self.settings.pricing_cache_timeout_seconds, maxsize=self.settings.price_batch_size * 10))(self._get_prices)

        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Sync context manager exit"""
        self._in_context = False
        return None
    
    def paginate(self, f: Callable):
        results, batches = [], self.get_pool_paginator()

        def process_batch(batch: List[Tuple]):
            with self.web3.batch_requests() as batcher:
                for offset, limit in batch: batcher.add(f(limit, offset))
                return sum(batcher.execute(), [])

        
        with ThreadPoolExecutor(max_workers=self.settings.threading_max_workers) as executor:
            future_to_batch = {
                executor.submit(process_batch, batch): batch
                for batch in batches
            }
            for future in as_completed(future_to_batch):
                try: results.extend(future.result())
                except Exception as e:
                    print(f"Error processing path chunk: {e}")
                    continue

        return results


    @require_context
    def sign_and_send_tx(self, tx, value: int = 0, wait: bool = True):
        spender = self.account.address
        tx = tx.build_transaction({ 'from': spender, 'value': value, 'nonce': self.web3.eth.get_transaction_count(spender) })
        signed_tx = self.account.sign_transaction(tx)
        tx_hash = self.web3.eth.send_raw_transaction(signed_tx.raw_transaction)
        return self.web3.eth.wait_for_transaction_receipt(tx_hash) if wait else tx_hash
    
    @require_context
    def set_token_allowance(self, token: Token, addr: str, amount: int):
        token_contract = self.prepare_set_token_allowance_contract(token, self.web3.eth.contract)
        return self.sign_and_send_tx(token_contract.functions.approve(addr, amount))

    @require_context
    @lru_cache(maxsize=None)
    def get_all_tokens(self, listed_only: bool = False) -> List[Token]:
        def get_tokens(limit, offset): return self.sugar.functions.tokens(limit, offset, ADDRESS_ZERO, [])
        return self.prepare_tokens(self.paginate(get_tokens), listed_only)
    
    def _get_prices(self, tokens: Tuple[Token]) -> List[int]:
        # token_address => normalized rate
        with self.web3.batch_requests() as batch:
            batch = self.prepare_price_batcher(tokens=list(tokens), batch=batch)
            return sum(batch.execute(), [])

    @require_context
    def get_prices(self, tokens: List[Token]) -> List[Price]:
        """Get prices for tokens in target stable token"""
        return self.prepare_prices(tokens, self._get_prices(tuple(tokens)))
    
    @lru_cache(maxsize=None)
    def get_raw_pools(self, for_swaps: bool):
        return self.paginate(self.sugar.functions.forSwaps if for_swaps else self.sugar.functions.all)

    @require_context
    def get_pools(self, for_swaps: bool = False) -> List[LiquidityPool]:
        pools = self.get_raw_pools(for_swaps)
        if not for_swaps:
            tokens = self.get_all_tokens(listed_only=False)
            return self.prepare_pools(pools, tokens, self.get_prices(tokens))
        else: return self.prepare_pools_for_swap(pools)

    @require_context
    @lru_cache(maxsize=None)
    def get_pool_by_address(self, address: str) -> Optional[LiquidityPool]:
        try:
            p = self.sugar.functions.byAddress(address).call()
        except: return None
        tokens = self.get_all_tokens(listed_only=False)
        return self.prepare_pools([p], tokens, self.get_prices(tokens))[0]
    
    @require_context
    def get_pools_for_swaps(self) -> List[LiquidityPoolForSwap]: return self.get_pools(for_swaps=True)

    @require_context
    @lru_cache(maxsize=None)
    def get_pool_epochs(self, lp: str, offset: int = 0, limit: int = 10) -> List[LiquidityPoolEpoch]:
        tokens, pools = self.get_all_tokens(listed_only=False), self.get_pools()
        prices = self.get_prices(tokens)
        r = self.sugar_rewards.functions.epochsByAddress(limit, offset, normalize_address(lp)).call()
        return self.prepare_pool_epochs(r, pools, tokens, prices)

    @require_context
    @lru_cache(maxsize=None)
    def get_latest_pool_epochs(self) -> List[LiquidityPoolEpoch]:
        tokens, pools = self.get_all_tokens(listed_only=False), self.get_pools()
        prices = self.get_prices(tokens)
        return self.prepare_pool_epochs(self.paginate(self.sugar_rewards.functions.epochsLatest), pools, tokens, prices)

    @require_context
    def _get_quotes_for_paths(self, from_token: Token, to_token: Token, amount_in: int, pools: List[LiquidityPoolForSwap], paths: List[List[Tuple]]) -> List[Optional[Quote]]:
        path_pools = self.paths_to_pools(pools, paths)
        with self.web3.batch_requests() as batch:
            batch, inputs = self.prepare_quote_batch(from_token, to_token, batch, path_pools, amount_in, paths)
            return self.prepare_quotes(inputs, batch.execute())
    
    @require_context
    def get_quote(self, from_token: Token, to_token: Token, amount: float, filter_quotes: Optional[Callable[[Quote], bool]] = None) -> Optional[Quote]:
        amount_in = float_to_uint256(amount, decimals=from_token.decimals)
        pools = self.filter_pools_for_swap(from_token=from_token, to_token=to_token, pools=self.get_pools_for_swaps())
        paths = self.get_paths_for_quote(from_token, to_token, pools, self.settings.excluded_tokens_addrs)
        path_chunks = list(chunk(paths, 500))

        def get_quotes_for_chunk(paths_chunk):
            return self._get_quotes_for_paths(from_token, to_token, amount_in, pools, paths_chunk)
        
        all_quotes = []
        
        with ThreadPoolExecutor(max_workers=self.settings.threading_max_workers) as executor:
            # Submit all chunk processing tasks
            future_to_chunk = {
                executor.submit(get_quotes_for_chunk, chunk_paths): chunk_paths 
                for chunk_paths in path_chunks
            }
            
            # Collect results as they complete
            for future in as_completed(future_to_chunk):
                try:
                    all_quotes.extend(future.result())
                except Exception as e:
                    print(f"Error processing path chunk: {e}")
                    continue

        # Filter out None quotes
        all_quotes = [q for q in all_quotes if q is not None]
    
        if filter_quotes is not None:  all_quotes = list(filter(filter_quotes, all_quotes))

        return max(all_quotes, key=lambda q: q.amount_out) if len(all_quotes) > 0 else None
    
    @require_context
    def swap(self, from_token: Token, to_token: Token, amount: float, slippage: Optional[float] = None):
        q = self.get_quote(from_token, to_token, amount)
        if not q: raise ValueError("No quotes found")
        return self.swap_from_quote(q, slippage=slippage)
        
    @require_context
    def swap_from_quote(self, quote: Quote, slippage: Optional[float] = None):
        swapper_contract_addr, from_token = self.settings.swapper_contract_addr, quote.from_token
        planner = setup_planner(quote=quote, slippage=slippage if slippage is not None else self.settings.swap_slippage, account=self.account.address, router_address=swapper_contract_addr)
        self.set_token_allowance(from_token, swapper_contract_addr, quote.input.amount_in)
        value = quote.input.amount_in if from_token.wrapped_token_address else 0
        return self.sign_and_send_tx(self.swapper.functions.execute(*[planner.commands, planner.inputs]), value=value)

# %% ../src/chains.ipynb 12
class OPChainCommon():
    usdc: Token = Token(token_address='0x0b2C639c533813f4Aa9D7837CAf62653d097Ff85', symbol='USDC', decimals=6, listed=True, wrapped_token_address=None)
    velo: Token = Token(token_address='0x9560e827aF36c94D2Ac33a39bCE1Fe78631088Db', symbol='VELO', decimals=18, listed=True, wrapped_token_address=None)
    eth: Token = Token(token_address='ETH', symbol='ETH', decimals=18, listed=True, wrapped_token_address='0x4200000000000000000000000000000000000006')
    
class AsyncOPChain(AsyncChain, OPChainCommon):
    def __init__(self, **kwargs): super().__init__(make_op_chain_settings(**kwargs), **kwargs)

class OPChain(Chain, OPChainCommon):
    def __init__(self, **kwargs): super().__init__(make_op_chain_settings(**kwargs), **kwargs)


# %% ../src/chains.ipynb 14
class BaseChainCommon():
    usdc: Token = Token(token_address='0x833589fCD6eDb6E08f4c7C32D4f71b54bdA02913', symbol='USDC', decimals=6, listed=True, wrapped_token_address=None)
    aero: Token = Token(token_address='0x940181a94A35A4569E4529A3CDfB74e38FD98631', symbol='AERO', decimals=18, listed=True, wrapped_token_address=None)
    eth: Token = Token(token_address='ETH', symbol='ETH', decimals=18, listed=True, wrapped_token_address='0x4200000000000000000000000000000000000006')
    
class AsyncBaseChain(AsyncChain, BaseChainCommon):
    def __init__(self, **kwargs): super().__init__(make_base_chain_settings(**kwargs), **kwargs)

class BaseChain(Chain, BaseChainCommon):
    def __init__(self, **kwargs): super().__init__(make_base_chain_settings(**kwargs), **kwargs)
