import asyncio
import logging
from functools import cached_property, lru_cache
from typing import Any, List, Optional, Union

from async_lru import alru_cache
from brownie import chain
from multicall import Call
from multicall.utils import await_awaitable
from y import convert
from y.classes.common import ERC20, ContractBase
from y.classes.singleton import Singleton
from y.contracts import Contract
from y.datatypes import AddressOrContract, AnyAddressType, Block, UsdPrice
from y.exceptions import ContractNotVerified
from y.networks import Network
from y.utils.logging import yLazyLogger
from y.utils.multicall import fetch_multicall
from y.utils.raw_calls import raw_call_async

logger = logging.getLogger(__name__)


v1_pools = {
    Network.Mainnet: [
        "0x398eC7346DcD622eDc5ae82352F02bE94C62d119"
    ],
}.get(chain.id, [])

v2_pools = {
    Network.Mainnet: [
        "0x7d2768dE32b0b80b7a3454c06BdAc94A69DDc7A9", # aave
        "0x7937D4799803FbBe595ed57278Bc4cA21f3bFfCB", # aave amm
        "0xcE744a9BAf573167B2CF138114BA32ed7De274Fa", # umee
    ],
    Network.Polygon: [
        "0x8dFf5E27EA6b7AC08EbFdf9eB090F32ee9a30fcf", # aave
    ],
    Network.Avalanche: [
        "0x70BbE4A294878a14CB3CDD9315f5EB490e346163", # blizz
    ]
}.get(chain.id, [])        


class AaveMarketBase(ContractBase):
    def __init__(self, address: AnyAddressType, *args: Any, **kwargs: Any) -> None:
        super().__init__(address, *args, **kwargs)
    
    def __contains__(self, __o: object) -> bool:
        return convert.to_address(__o) in self.atokens


class AaveMarketV1(AaveMarketBase):
    def __init__(self, address: AnyAddressType, *args: Any, **kwargs: Any) -> None:
        super().__init__(address, *args, **kwargs)
    
    def __repr__(self) -> str:
        return f"<AaveMarketV1 '{self.address}'>"
    
    #yLazyLogger(logger)
    @lru_cache
    def underlying(self, token_address: AddressOrContract) -> ERC20:
        return await_awaitable(self.underlying_async(token_address))
    
    #yLazyLogger(logger)
    @alru_cache
    async def underlying_async(self, token_address: AddressOrContract) -> ERC20:
        underlying = await raw_call_async(token_address, 'underlyingAssetAddress()',output='address')
        return ERC20(underlying)
    
    @cached_property
    #yLazyLogger(logger)
    def atokens(self) -> List[ERC20]:
        reserves_data = Call(self.address, ['getReserves()(address[])'], [[self.address,None]])()[self.address]
        reserves_data = fetch_multicall(*[[self.contract, 'getReserveData', reserve] for reserve in reserves_data])
        atokens = [ERC20(reserve['aTokenAddress']) for reserve in reserves_data]
        logger.info(f'loaded {len(atokens)} v1 atokens for {self.__repr__()}')
        return atokens


class AaveMarketV2(AaveMarketBase):
    def __init__(self, address: AnyAddressType, *args: Any, **kwargs: Any) -> None:
        super().__init__(address, *args, **kwargs)
    
    def __repr__(self) -> str:
        return f"<AaveMarketV2 '{self.address}'>"
    
    #yLazyLogger(logger)
    @lru_cache
    def underlying(self, token_address: AddressOrContract) -> ERC20:
        return await_awaitable(self.underlying_async(token_address))
    
    #yLazyLogger(logger)
    @alru_cache
    async def underlying_async(self, token_address: AddressOrContract) -> ERC20:
        underlying = await raw_call_async(token_address, 'UNDERLYING_ASSET_ADDRESS()',output='address')
        print(underlying)
        return ERC20(underlying)

    @cached_property
    #yLazyLogger(logger)
    def atokens(self) -> List[ERC20]:
        reserves = Call(self.address, ['getReservesList()(address[])'], [[self.address,None]])()[self.address]
        reserves_data = [
            Call(
                self.address,
                ['getReserveData(address)((uint256,uint128,uint128,uint128,uint128,uint128,uint40,address,address,address,address,uint8))',reserve],
            )()
            for reserve in reserves
        ]

        try:
            atokens = [ERC20(reserve_data[7]) for reserve_data in reserves_data]
            logger.info(f'loaded {len(atokens)} v2 atokens for {self.__repr__()}')
            return atokens
        except TypeError: # TODO figure out what to do about non verified aave markets
            logger.error(f'failed to load tokens for {self.__repr__()}')
            return []


class AaveRegistry(metaclass = Singleton):
    def __init__(self) -> None:
        pass

    @cached_property
    def pools(self) -> List[Union[AaveMarketV1, AaveMarketV2]]:
        return self.pools_v1 + self.pools_v2
    
    @cached_property
    def pools_v1(self) -> List[AaveMarketV1]:
        return [AaveMarketV1(pool) for pool in v1_pools]
    
    @cached_property
    def pools_v2(self) -> List[AaveMarketV2]:
        return [AaveMarketV2(pool) for pool in v2_pools]
    
    #yLazyLogger(logger)
    def pool_for_atoken(self, token_address: AnyAddressType) -> Optional[Union[AaveMarketV1, AaveMarketV2]]:
        for pool in self.pools:
            if token_address in pool:
                return pool

    def __contains__(self, __o: object) -> bool:
        return any(__o in pool for pool in self.pools)

    #yLazyLogger(logger)
    @lru_cache
    def is_atoken(self, token_address: AnyAddressType) -> bool:
        return token_address in self
    
    async def is_wrapped_atoken(self, token_address: AnyAddressType) -> bool:
        try:
            contract = await Contract.coroutine(token_address)
            attrs = "ATOKEN", "STATIC_ATOKEN_LM_REVISION", "staticToDynamicAmount"
            return all(hasattr(contract, attr) for attr in attrs)
        except ContractNotVerified:
            return False
    
    #yLazyLogger(logger)
    @lru_cache
    def underlying(self, token_address: AddressOrContract) -> ERC20:
        return await_awaitable(self.underlying_async(token_address))
    
    #yLazyLogger(logger)
    @alru_cache
    async def underlying_async(self, token_address: AddressOrContract) -> ERC20:
        pool = self.pool_for_atoken(token_address)
        return await pool.underlying_async(token_address)
    
    #yLazyLogger(logger)
    def get_price(self, token_address: AddressOrContract, block: Optional[Block] = None) -> UsdPrice:
        return await_awaitable(self.get_price_async(token_address, block=block))
    
    #yLazyLogger(logger)
    async def get_price_async(self, token_address: AddressOrContract, block: Optional[Block] = None) -> UsdPrice:
        underlying = await self.underlying_async(token_address)
        return await underlying.price_async(block)
    
    async def get_price_wrapped_async(self, token_address: AddressOrContract, block: Optional[Block] = None) -> UsdPrice:
        contract, scale = await asyncio.gather(Contract.coroutine(token_address), ERC20(token_address).scale)
        underlying, price_per_share = await asyncio.gather(
            contract.ATOKEN.coroutine(block_identifier=block),
            contract.staticToDynamicAmount.coroutine(scale, block_identifier=block),
        )
        price_per_share /= scale
        return price_per_share * await ERC20(underlying).price_async(block)


aave = AaveRegistry()
