from importlib import resources as impresources
import json
from typing import TYPE_CHECKING
import enum

CHAINS_FILE = impresources.files("tristero.files") / "chains.json"
CHAINS = json.loads(CHAINS_FILE.read_text())

def _chain_key(chain: str | int) -> str:
    """Resolve a chain key from either a chain key (e.g. 'arbitrum') or a chain id (e.g. 42161/'42161')."""
    s = str(chain)
    if s in CHAINS:
        return s

    try:
        chain_id = int(s)
    except ValueError as exc:
        raise KeyError(f"Unknown chain: {chain}") from exc

    for key, data in CHAINS.items():
        if data.get("chainId") == chain_id:
            return key

    raise KeyError(f"Unknown chain id: {chain}")


def get_address(chain: str | int, id: str) -> str | None:
    key = _chain_key(chain)
    addresses = CHAINS[key].get("addresses") or {}
    return addresses.get(id)

def get_permit2_addr(chain: str | int):
    return get_address(chain, 'permit2')

def get_wrapped_gas_addr(chain: str | int):
    return get_address(chain, 'wrappedGasToken')

def get_gas_addr(chain: str | int):
    return get_address(chain, 'gasToken')

def chain(id: str | int) -> str | None:
    """Return the chain key for a given chain id or chain key."""
    return _chain_key(id)

if TYPE_CHECKING:
    from enum import IntEnum
    class ChainID(IntEnum):
        _value_: int
else:
    ChainID = enum.Enum(
        'ChainID',
        {
            name: data['chainId']
            for name, data in CHAINS.items()
        }
    )
