"""
Token constants for supported networks.

This module provides well-known token addresses for each supported chain,
making it easy to reference tokens without hardcoding addresses.

Example:
    >>> from xccy.tokens import PolygonTokens
    >>> 
    >>> # Use USDC as collateral
    >>> client.margin.deposit(account, PolygonTokens.USDC, 1000 * 10**6)
    >>> 
    >>> # Build pool key with yield-bearing token
    >>> pool = client.pool.build_pool_key(
    ...     underlying=PolygonTokens.USDC,
    ...     compound_token=PolygonTokens.A_USDC,
    ...     ...
    ... )
"""

from dataclasses import dataclass
from typing import Optional


@dataclass(frozen=True)
class TokenInfo:
    """
    Information about a token.
    
    Attributes:
        address: Contract address (checksummed).
        symbol: Token symbol (e.g., "USDC").
        decimals: Number of decimals.
        name: Full token name.
    """
    
    address: str
    symbol: str
    decimals: int
    name: str = ""
    
    def __str__(self) -> str:
        return self.address


class PolygonTokens:
    """
    Well-known token addresses on Polygon mainnet (chain ID 137).
    
    Includes stablecoins, yield-bearing tokens (aTokens), and ETH variants.
    
    Example:
        >>> from xccy.tokens import PolygonTokens
        >>> 
        >>> # Access address directly
        >>> usdc_address = PolygonTokens.USDC  # "0x3c499c..."
        >>> 
        >>> # Or get full token info
        >>> usdc = PolygonTokens.get_info("USDC")
        >>> print(usdc.decimals)  # 6
    """
    
    # ═══════════════════════════════════════════════════════════════════════════
    # Stablecoins
    # ═══════════════════════════════════════════════════════════════════════════
    
    USDC: str = "0x3c499c542cEF5E3811e1192ce70d8cC03d5c3359"
    """Native USDC on Polygon (6 decimals)."""
    
    USDC_E: str = "0x2791Bca1f2de4661ED88A30C99A7a9449Aa84174"
    """Bridged USDC.e on Polygon (6 decimals)."""
    
    USDT: str = "0xc2132D05D31c914a87C6611C10748AEb04B58e8F"
    """Tether USD on Polygon (6 decimals)."""
    
    DAI: str = "0x8f3Cf7ad23Cd3CaDbD9735AFf958023239c6A063"
    """DAI Stablecoin on Polygon (18 decimals)."""
    
    # ═══════════════════════════════════════════════════════════════════════════
    # Aave v3 aTokens (Yield-bearing)
    # ═══════════════════════════════════════════════════════════════════════════
    
    A_USDC: str = "0x625E7708f30cA75bfd92586e17077590C60eb4cD"
    """Aave v3 aUSDC (interest-bearing USDC)."""
    
    A_USDT: str = "0x6ab707Aca953eDAeFBc4fD23bA73294241490620"
    """Aave v3 aUSDT (interest-bearing USDT)."""
    
    A_DAI: str = "0x82E64f49Ed5EC1bC6e43DAD4FC8Af9bb3A2312EE"
    """Aave v3 aDAI (interest-bearing DAI)."""
    
    A_WETH: str = "0xe50fA9b3c56FfB159cB0FCA61F5c9D750e8128c8"
    """Aave v3 aWETH (interest-bearing WETH)."""
    
    A_WMATIC: str = "0x6d80113e533a2C0fe82EaBD35f1875DcEA89Ea97"
    """Aave v3 aWMATIC (interest-bearing WMATIC)."""
    
    # ═══════════════════════════════════════════════════════════════════════════
    # ETH Variants
    # ═══════════════════════════════════════════════════════════════════════════
    
    WETH: str = "0x7ceB23fD6bC0adD59E62ac25578270cFf1b9f619"
    """Wrapped Ether on Polygon (18 decimals)."""
    
    WMATIC: str = "0x0d500B1d8E8eF31E21C99d1Db9A6444d3ADf1270"
    """Wrapped MATIC (18 decimals)."""
    
    # ═══════════════════════════════════════════════════════════════════════════
    # Liquid Staking
    # ═══════════════════════════════════════════════════════════════════════════
    
    STMATIC: str = "0x3A58a54C066FdC0f2D55FC9C89F0415C92eBf3C4"
    """Lido stMATIC (18 decimals)."""
    
    WSTETH: str = "0x03b54A6e9a984069379fae1a4fC4dBAE93B3bCCD"
    """Wrapped stETH on Polygon (18 decimals)."""
    
    # ═══════════════════════════════════════════════════════════════════════════
    # Token Info Registry
    # ═══════════════════════════════════════════════════════════════════════════
    
    _TOKEN_INFO: dict[str, TokenInfo] = {
        "USDC": TokenInfo("0x3c499c542cEF5E3811e1192ce70d8cC03d5c3359", "USDC", 6, "USD Coin"),
        "USDC.e": TokenInfo("0x2791Bca1f2de4661ED88A30C99A7a9449Aa84174", "USDC.e", 6, "Bridged USD Coin"),
        "USDT": TokenInfo("0xc2132D05D31c914a87C6611C10748AEb04B58e8F", "USDT", 6, "Tether USD"),
        "DAI": TokenInfo("0x8f3Cf7ad23Cd3CaDbD9735AFf958023239c6A063", "DAI", 18, "Dai Stablecoin"),
        "aUSDC": TokenInfo("0x625E7708f30cA75bfd92586e17077590C60eb4cD", "aUSDC", 6, "Aave Polygon USDC"),
        "aUSDT": TokenInfo("0x6ab707Aca953eDAeFBc4fD23bA73294241490620", "aUSDT", 6, "Aave Polygon USDT"),
        "aDAI": TokenInfo("0x82E64f49Ed5EC1bC6e43DAD4FC8Af9bb3A2312EE", "aDAI", 18, "Aave Polygon DAI"),
        "aWETH": TokenInfo("0xe50fA9b3c56FfB159cB0FCA61F5c9D750e8128c8", "aWETH", 18, "Aave Polygon WETH"),
        "aWMATIC": TokenInfo("0x6d80113e533a2C0fe82EaBD35f1875DcEA89Ea97", "aWMATIC", 18, "Aave Polygon WMATIC"),
        "WETH": TokenInfo("0x7ceB23fD6bC0adD59E62ac25578270cFf1b9f619", "WETH", 18, "Wrapped Ether"),
        "WMATIC": TokenInfo("0x0d500B1d8E8eF31E21C99d1Db9A6444d3ADf1270", "WMATIC", 18, "Wrapped Matic"),
        "stMATIC": TokenInfo("0x3A58a54C066FdC0f2D55FC9C89F0415C92eBf3C4", "stMATIC", 18, "Lido Staked MATIC"),
        "wstETH": TokenInfo("0x03b54A6e9a984069379fae1a4fC4dBAE93B3bCCD", "wstETH", 18, "Wrapped stETH"),
    }
    
    @classmethod
    def get_info(cls, symbol: str) -> Optional[TokenInfo]:
        """
        Get token information by symbol.
        
        Args:
            symbol: Token symbol (e.g., "USDC", "aUSDC").
            
        Returns:
            TokenInfo if found, None otherwise.
            
        Example:
            >>> info = PolygonTokens.get_info("USDC")
            >>> print(info.decimals)  # 6
        """
        return cls._TOKEN_INFO.get(symbol)
    
    @classmethod
    def get_decimals(cls, address: str) -> int:
        """
        Get decimals for a token by address.
        
        Args:
            address: Token contract address.
            
        Returns:
            Number of decimals, or 18 as default.
        """
        address_lower = address.lower()
        for info in cls._TOKEN_INFO.values():
            if info.address.lower() == address_lower:
                return info.decimals
        return 18  # Default to 18 decimals


class ArbitrumTokens:
    """
    Well-known token addresses on Arbitrum One (chain ID 42161).
    
    Includes stablecoins, yield-bearing tokens (aTokens), ETH variants,
    and ARB governance token.
    
    Example:
        >>> from xccy.tokens import ArbitrumTokens
        >>> 
        >>> usdc_address = ArbitrumTokens.USDC  # "0xaf88d065..."
        >>> 
        >>> info = ArbitrumTokens.get_info("USDC")
        >>> print(info.decimals)  # 6
    """
    
    # ═══════════════════════════════════════════════════════════════════════════
    # Stablecoins
    # ═══════════════════════════════════════════════════════════════════════════
    
    USDC: str = "0xaf88d065e77c8cC2239327C5EDb3A432268e5831"
    """Native USDC on Arbitrum (6 decimals)."""
    
    USDC_E: str = "0xFF970A61A04b1cA14834A43f5dE4533eBDDB5CC8"
    """Bridged USDC.e on Arbitrum (6 decimals)."""
    
    USDT: str = "0xFd086bC7CD5C481DCC9C85ebE478A1C0b69FCbb9"
    """Tether USD on Arbitrum (6 decimals)."""
    
    DAI: str = "0xDA10009cBd5D07dd0CeCc66161FC93D7c9000da1"
    """DAI Stablecoin on Arbitrum (18 decimals)."""
    
    # ═══════════════════════════════════════════════════════════════════════════
    # Aave v3 aTokens (Yield-bearing)
    # ═══════════════════════════════════════════════════════════════════════════
    
    A_USDC: str = "0x724dc807b04555b71ed48a6896b6F41593b8C637"
    """Aave v3 aUSDC (interest-bearing native USDC)."""
    
    A_USDC_E: str = "0x625E7708f30cA75bfd92586e17077590C60eb4cD"
    """Aave v3 aUSDC.e (interest-bearing bridged USDC)."""
    
    A_USDT: str = "0x6ab707Aca953eDAeFBc4fD23bA73294241490620"
    """Aave v3 aUSDT (interest-bearing USDT)."""
    
    A_DAI: str = "0x82E64f49Ed5EC1bC6e43DAD4FC8Af9bb3A2312EE"
    """Aave v3 aDAI (interest-bearing DAI)."""
    
    A_WETH: str = "0xe50fA9b3c56FfB159cB0FCA61F5c9D750e8128c8"
    """Aave v3 aWETH (interest-bearing WETH)."""
    
    # ═══════════════════════════════════════════════════════════════════════════
    # ETH Variants
    # ═══════════════════════════════════════════════════════════════════════════
    
    WETH: str = "0x82aF49447D8a07e3bd95BD0d56f35241523fBab1"
    """Wrapped Ether on Arbitrum (18 decimals)."""
    
    WBTC: str = "0x2f2a2543B76A4166549F7aaB2e75Bef0aefC5B0f"
    """Wrapped Bitcoin on Arbitrum (8 decimals)."""
    
    # ═══════════════════════════════════════════════════════════════════════════
    # Liquid Staking
    # ═══════════════════════════════════════════════════════════════════════════
    
    WSTETH: str = "0x5979D7b546E38E414F7E9822514be443A4800529"
    """Wrapped stETH on Arbitrum (18 decimals)."""
    
    # ═══════════════════════════════════════════════════════════════════════════
    # Governance
    # ═══════════════════════════════════════════════════════════════════════════
    
    ARB: str = "0x912CE59144191C1204E64559FE8253a0e49E6548"
    """ARB governance token (18 decimals)."""
    
    # ═══════════════════════════════════════════════════════════════════════════
    # Token Info Registry
    # ═══════════════════════════════════════════════════════════════════════════
    
    _TOKEN_INFO: dict[str, TokenInfo] = {
        "USDC": TokenInfo("0xaf88d065e77c8cC2239327C5EDb3A432268e5831", "USDC", 6, "USD Coin"),
        "USDC.e": TokenInfo("0xFF970A61A04b1cA14834A43f5dE4533eBDDB5CC8", "USDC.e", 6, "Bridged USD Coin"),
        "USDT": TokenInfo("0xFd086bC7CD5C481DCC9C85ebE478A1C0b69FCbb9", "USDT", 6, "Tether USD"),
        "DAI": TokenInfo("0xDA10009cBd5D07dd0CeCc66161FC93D7c9000da1", "DAI", 18, "Dai Stablecoin"),
        "aUSDC": TokenInfo("0x724dc807b04555b71ed48a6896b6F41593b8C637", "aUSDC", 6, "Aave Arbitrum USDC"),
        "aUSDC.e": TokenInfo("0x625E7708f30cA75bfd92586e17077590C60eb4cD", "aUSDC.e", 6, "Aave Arbitrum USDC.e"),
        "aUSDT": TokenInfo("0x6ab707Aca953eDAeFBc4fD23bA73294241490620", "aUSDT", 6, "Aave Arbitrum USDT"),
        "aDAI": TokenInfo("0x82E64f49Ed5EC1bC6e43DAD4FC8Af9bb3A2312EE", "aDAI", 18, "Aave Arbitrum DAI"),
        "aWETH": TokenInfo("0xe50fA9b3c56FfB159cB0FCA61F5c9D750e8128c8", "aWETH", 18, "Aave Arbitrum WETH"),
        "WETH": TokenInfo("0x82aF49447D8a07e3bd95BD0d56f35241523fBab1", "WETH", 18, "Wrapped Ether"),
        "WBTC": TokenInfo("0x2f2a2543B76A4166549F7aaB2e75Bef0aefC5B0f", "WBTC", 8, "Wrapped Bitcoin"),
        "wstETH": TokenInfo("0x5979D7b546E38E414F7E9822514be443A4800529", "wstETH", 18, "Wrapped stETH"),
        "ARB": TokenInfo("0x912CE59144191C1204E64559FE8253a0e49E6548", "ARB", 18, "Arbitrum"),
    }
    
    @classmethod
    def get_info(cls, symbol: str) -> Optional[TokenInfo]:
        """
        Get token information by symbol.
        
        Args:
            symbol: Token symbol (e.g., "USDC", "aUSDC").
            
        Returns:
            TokenInfo if found, None otherwise.
        """
        return cls._TOKEN_INFO.get(symbol)
    
    @classmethod
    def get_decimals(cls, address: str) -> int:
        """
        Get decimals for a token by address.
        
        Args:
            address: Token contract address.
            
        Returns:
            Number of decimals, or 18 as default.
        """
        address_lower = address.lower()
        for info in cls._TOKEN_INFO.values():
            if info.address.lower() == address_lower:
                return info.decimals
        return 18


# Alias for convenience (defaults to Polygon for backward compatibility)
Tokens = PolygonTokens

# All token registries for cross-chain lookups
_ALL_REGISTRIES: list[type] = [PolygonTokens, ArbitrumTokens]


def _resolve_decimals(token: str) -> int:
    """Resolve decimals by searching all token registries."""
    for registry in _ALL_REGISTRIES:
        info = registry.get_info(token)
        if info:
            return info.decimals
    for registry in _ALL_REGISTRIES:
        decimals = registry.get_decimals(token)
        if decimals != 18:
            return decimals
    return 18


# ═══════════════════════════════════════════════════════════════════════════════
# Helper functions for decimal conversion
# ═══════════════════════════════════════════════════════════════════════════════

def parse_amount(amount: float | int, token: str | TokenInfo) -> int:
    """
    Convert human-readable amount to raw token units.
    
    Searches all supported chains (Polygon, Arbitrum) when resolving symbols.
    
    Args:
        amount: Human-readable amount (e.g., 100.5 for 100.5 USDC).
        token: Token symbol (str) or TokenInfo object.
        
    Returns:
        Raw amount in smallest token units.
        
    Example:
        >>> from xccy.tokens import parse_amount
        >>> 
        >>> # 100 USDC (6 decimals) -> 100_000_000
        >>> raw = parse_amount(100, "USDC")
        >>> 
        >>> # 0.5 WETH (18 decimals) -> 500_000_000_000_000_000
        >>> raw = parse_amount(0.5, "WETH")
    """
    if isinstance(token, TokenInfo):
        decimals = token.decimals
    elif isinstance(token, str):
        decimals = _resolve_decimals(token)
    else:
        decimals = 18
    
    return int(amount * (10 ** decimals))


def format_amount(raw_amount: int, token: str | TokenInfo, precision: int = 4) -> str:
    """
    Convert raw token units to human-readable string.
    
    Searches all supported chains (Polygon, Arbitrum) when resolving symbols.
    
    Args:
        raw_amount: Amount in smallest token units.
        token: Token symbol (str) or TokenInfo object.
        precision: Decimal places in output. Defaults to 4.
        
    Returns:
        Formatted string with appropriate decimals.
        
    Example:
        >>> from xccy.tokens import format_amount
        >>> 
        >>> # 100_000_000 USDC units -> "100.0000"
        >>> formatted = format_amount(100_000_000, "USDC")
        >>> 
        >>> # -500_000_000_000_000_000 WETH -> "-0.5000"
        >>> formatted = format_amount(-500_000_000_000_000_000, "WETH")
    """
    if isinstance(token, TokenInfo):
        decimals = token.decimals
    elif isinstance(token, str):
        decimals = _resolve_decimals(token)
    else:
        decimals = 18
    
    value = raw_amount / (10 ** decimals)
    return f"{value:.{precision}f}"


def get_decimals(token: str) -> int:
    """
    Get decimals for a token by symbol or address.
    
    Searches all supported chains (Polygon, Arbitrum).
    
    Args:
        token: Token symbol (e.g., "USDC") or address.
        
    Returns:
        Number of decimals.
        
    Example:
        >>> from xccy.tokens import get_decimals
        >>> get_decimals("USDC")  # 6
        >>> get_decimals("WETH")  # 18
    """
    return _resolve_decimals(token)
