import logging
import json
import os
from decimal import Decimal
from typing import Optional, Union, Dict, Any, List

from web3 import Web3
from web3.contract import Contract
from eth_account import Account
from eth_account.signers.local import LocalAccount

# Try to import optional dependencies
try:
    from dotenv import load_dotenv
    DOTENV_AVAILABLE = True
except ImportError:
    DOTENV_AVAILABLE = False

try:
    from google.colab import userdata  # type: ignore[import-untyped]
    COLAB_AVAILABLE = True
except ImportError:
    COLAB_AVAILABLE = False

from . import constants
from . import utils
from .exceptions import (
    FXProtocolError,
    TransactionFailedError,
    InsufficientBalanceError,
    ContractCallError,
    ConfigurationError
)

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("fx_sdk")

class ProtocolClient:
    """
    The main client for interacting with the f(x) Protocol.
    
    This class abstracts Web3.py boilerplate and provides easy-to-use
    methods for reading from and writing to the f(x) Protocol smart contracts.
    
    Wallet Authentication Priority:
    1. Explicit private_key parameter
    2. Environment variable: FX_PROTOCOL_PRIVATE_KEY
    3. .env file: FX_PROTOCOL_PRIVATE_KEY
    4. Google Colab secret: 'fx_protocol_private_key'
    5. Browser-injected wallet (window.ethereum)
    """

    def __init__(
        self,
        rpc_url: str,
        private_key: Optional[str] = None,
        abi_dir: Optional[str] = None,
        log_level: int = logging.INFO,
        use_browser_wallet: bool = False
    ):
        """
        Initialize the ProtocolClient.
        
        Args:
            rpc_url: The RPC URL for the Ethereum network.
            private_key: Optional private key for signing transactions. If not provided,
                       the client will attempt to discover credentials from environment
                       variables, .env files, Colab secrets, or browser wallets.
            abi_dir: Optional directory where contract ABIs are stored. Defaults to internal package directory.
            log_level: Logging level (default logging.INFO).
            use_browser_wallet: If True, attempt to connect to a browser-injected wallet (MetaMask, etc.).
                              Requires running in a browser environment with Web3 wallet extension.
        """
        logger.setLevel(log_level)
        
        # Load .env file if available
        if DOTENV_AVAILABLE:
            load_dotenv()
        
        # Discover wallet credentials
        discovered_key = self._discover_wallet_credentials(private_key, use_browser_wallet)
        
        # Initialize Web3 connection
        if use_browser_wallet and discovered_key is None:
            # Try to use browser-injected provider
            try:
                # This will work in browser environments (Jupyter with ipywidgets, etc.)
                # For Node.js-like environments, you'd use window.ethereum
                self.w3 = Web3()  # Will be set by browser provider
                logger.warning("Browser wallet connection requires additional setup. Falling back to RPC provider.")
                self.w3 = Web3(Web3.HTTPProvider(rpc_url))
            except Exception as e:
                logger.warning(f"Browser wallet not available: {e}. Using RPC provider.")
                self.w3 = Web3(Web3.HTTPProvider(rpc_url))
        else:
            self.w3 = Web3(Web3.HTTPProvider(rpc_url))
        
        try:
            is_connected = self.w3.is_connected()
        except AttributeError:
            is_connected = self.w3.isConnected()
            
        if not is_connected:
            raise ConfigurationError(f"Failed to connect to RPC at {rpc_url}")

        # Set up account
        self.account: Optional[LocalAccount] = None
        self.use_browser_wallet = use_browser_wallet and discovered_key is None
        
        if discovered_key:
            self.account = Account.from_key(discovered_key)
            self.address = self.account.address
            logger.info(f"Initialized client for address: {self.address}")
        elif use_browser_wallet:
            # Browser wallet will be used for signing, but we need to get address from provider
            try:
                accounts = self.w3.eth.accounts
                if accounts:
                    self.address = accounts[0]
                    logger.info(f"Connected to browser wallet: {self.address}")
                else:
                    self.address = None
                    logger.warning("Browser wallet connected but no accounts available")
            except Exception as e:
                logger.warning(f"Could not get browser wallet address: {e}")
                self.address = None
        else:
            self.address = None
            logger.warning("Initialized client without private key (Read-Only Mode)")

        # Default abi_dir to the internal package directory
        if abi_dir is None:
            self.abi_dir = os.path.join(os.path.dirname(__file__), "abis")
        else:
            self.abi_dir = abi_dir
            
        self.contracts: Dict[str, Contract] = {}
        self._load_contracts()

    def _discover_wallet_credentials(
        self, 
        explicit_key: Optional[str] = None,
        use_browser: bool = False
    ) -> Optional[str]:
        """
        Discover wallet credentials from multiple secure sources.
        
        Priority order:
        1. Explicit private_key parameter
        2. Environment variable: FX_PROTOCOL_PRIVATE_KEY
        3. .env file: FX_PROTOCOL_PRIVATE_KEY (loaded via dotenv)
        4. Google Colab secret: 'fx_protocol_private_key'
        5. Browser wallet (if use_browser=True)
        
        Returns:
            Optional[str]: Private key if discovered, None otherwise.
        """
        # 1. Explicit parameter (highest priority)
        if explicit_key:
            logger.debug("Using explicitly provided private key")
            return explicit_key
        
        # 2. Environment variable
        env_key = os.getenv("FX_PROTOCOL_PRIVATE_KEY")
        if env_key:
            logger.debug("Using private key from environment variable")
            return env_key
        
        # 3. .env file (already loaded by load_dotenv if available)
        # This is the same as env var, but loaded from .env file
        
        # 4. Google Colab secret
        if COLAB_AVAILABLE:
            try:
                colab_key = userdata.get('fx_protocol_private_key')
                if colab_key:
                    logger.debug("Using private key from Google Colab secret")
                    return colab_key
            except Exception as e:
                logger.debug(f"Colab secret not available: {e}")
        
        # 5. Browser wallet (handled separately in __init__)
        if use_browser:
            logger.debug("Browser wallet will be used for signing")
            return None
        
        return None

    def _load_contracts(self):
        """Internal method to pre-instantiate contract objects."""
        # Core Contracts
        self.fxusd = self._get_contract("fxusd", constants.FXUSD)
        self.fxusd_base_pool = self._get_contract("fxusd_base_pool", constants.FXUSD_BASE_POOL)
        self.diamond_router = self._get_contract("diamond", constants.DIAMOND_ROUTER)
        
        # Governance
        self.fxn = self._get_contract("fxn", constants.FXN)
        self.vefxn = self._get_contract("vefxn", constants.VEFXN)
        self.gauge_controller = self._get_contract("gauge_controller", constants.GAUGE_CONTROLLER)
        
        # V1 Market (for fETH/xETH)
        self.v1_market = self._get_contract("market", constants.MARKET_PROXY)
        self.v1_rebalance_registry = self._get_contract("rebalance_pool_registry", constants.REBALANCE_POOL_REGISTRY)
        
        # Supporting
        self.multipath_converter = self._get_contract("multipath_converter", constants.MULTI_PATH_CONVERTER)
        self.steth_gateway = self._get_contract("steth_gateway", constants.STETH_GATEWAY)

    def _get_contract(self, name: str, address: str) -> Contract:
        """
        Load a contract by its name and address.
        
        Args:
            name: The name of the contract (used to find the ABI file).
            address: The Ethereum address of the contract.
            
        Returns:
            Contract: The Web3 contract object.
        """
        abi_path = os.path.join(self.abi_dir, f"{name.lower()}.json")
        checksum_address = utils.to_checksum_address(address)
        try:
            if os.path.exists(abi_path) and os.path.getsize(abi_path) > 0:
                with open(abi_path, "r") as f:
                    abi = json.load(f)
            else:
                abi = []  # Fallback to empty ABI
            return self.w3.eth.contract(address=checksum_address, abi=abi)
        except Exception as e:
            # For now, we return a contract with empty ABI if file loading fails
            # This prevents initialization errors while ABIs are being added
            return self.w3.eth.contract(address=checksum_address, abi=[])

    # --- Generic Read Methods ---

    def get_token_balance(self, token_address: str, account_address: Optional[str] = None) -> Decimal:
        """
        Get the human-readable balance of a token for an account.
        
        Args:
            token_address: The address of the token contract (ERC20).
            account_address: Optional account address (defaults to client's address).
            
        Returns:
            Decimal: The human-readable balance.
        """
        target_address = account_address or self.address
        if not target_address:
            raise FXProtocolError("No account address provided or available in client.")
        
        contract = self.w3.eth.contract(
            address=utils.to_checksum_address(token_address),
            abi=[
                {"constant": True, "inputs": [{"name": "_owner", "type": "address"}], "name": "balanceOf", "outputs": [{"name": "balance", "type": "uint256"}], "type": "function"},
                {"constant": True, "inputs": [], "name": "decimals", "outputs": [{"name": "", "type": "uint8"}], "type": "function"}
            ]
        )
        
        try:
            raw_balance = contract.functions.balanceOf(target_address).call()
            decimals = contract.functions.decimals().call()
            return utils.wei_to_decimal(raw_balance, decimals)
        except Exception as e:
            raise ContractCallError(f"Failed to get balance: {str(e)}")

    def get_token_total_supply(self, token_address: str) -> Decimal:
        """Get the total supply of a token."""
        contract = self.w3.eth.contract(
            address=utils.to_checksum_address(token_address),
            abi=[
                {"constant": True, "inputs": [], "name": "totalSupply", "outputs": [{"name": "", "type": "uint256"}], "type": "function"},
                {"constant": True, "inputs": [], "name": "decimals", "outputs": [{"name": "", "type": "uint8"}], "type": "function"}
            ]
        )
        try:
            raw_supply = contract.functions.totalSupply().call()
            decimals = contract.functions.decimals().call()
            return utils.wei_to_decimal(raw_supply, decimals)
        except Exception as e:
            raise ContractCallError(f"Failed to get total supply: {str(e)}")

    def get_allowance(self, token_address: str, owner: str, spender: str) -> Decimal:
        """Get the allowance of a spender for a token owner."""
        contract = self.w3.eth.contract(
            address=utils.to_checksum_address(token_address),
            abi=[
                {"constant": True, "inputs": [{"name": "_owner", "type": "address"}, {"name": "_spender", "type": "address"}], "name": "allowance", "outputs": [{"name": "", "type": "uint256"}], "type": "function"},
                {"constant": True, "inputs": [], "name": "decimals", "outputs": [{"name": "", "type": "uint8"}], "type": "function"}
            ]
        )
        try:
            raw_allowance = contract.functions.allowance(
                utils.to_checksum_address(owner),
                utils.to_checksum_address(spender)
            ).call()
            decimals = contract.functions.decimals().call()
            return utils.wei_to_decimal(raw_allowance, decimals)
        except Exception as e:
            raise ContractCallError(f"Failed to get allowance: {str(e)}")

    # --- V2 Product-Specific Read Methods ---

    def get_fxusd_total_supply(self) -> Decimal:
        """Get the total supply of fxUSD."""
        return self.get_token_total_supply(constants.FXUSD)

    def get_steth_price(self) -> Decimal:
        """Get the current stETH price from the oracle."""
        contract = self.w3.eth.contract(
            address=utils.to_checksum_address(constants.STETH_PRICE_ORACLE),
            abi=[{"constant": True, "inputs": [], "name": "getPrice", "outputs": [{"name": "", "type": "uint256"}], "type": "function"}]
        )
        try:
            raw_price = contract.functions.getPrice().call()
            return utils.wei_to_decimal(raw_price, 18)
        except Exception as e:
            raise ContractCallError(f"Failed to get stETH price: {str(e)}")

    def get_v2_pool_info(self) -> Dict[str, Any]:
        """
        Get information about the V2 fxUSD Base Pool.
        Note: Requires Base Pool ABI.
        """
        try:
            info = self.fxusd_base_pool.functions.getPoolInfo().call()
            return {
                "base_pool_address": info[0],
                "total_assets": utils.wei_to_decimal(info[1]),
                "total_supply": utils.wei_to_decimal(info[2]),
                # ... add more fields as needed based on ABI
            }
        except Exception as e:
            raise ContractCallError(f"Failed to get V2 pool info: {str(e)}")

    # --- V1 Legacy Read Methods ---

    def get_v1_nav(self) -> Dict[str, Decimal]:
        """
        Get the Net Asset Value (NAV) for V1 fETH and xETH.
        Note: Requires Market ABI.
        """
        try:
            nav = self.v1_market.functions.getNav().call()
            return {
                "fETH_NAV": utils.wei_to_decimal(nav[0]),
                "xETH_NAV": utils.wei_to_decimal(nav[1])
            }
        except Exception as e:
            raise ContractCallError(f"Failed to get V1 NAV: {str(e)}")

    def get_v1_collateral_ratio(self) -> Decimal:
        """
        Get the current collateral ratio of the V1 market.
        Note: Requires Market ABI.
        """
        try:
            cr = self.v1_market.functions.collateralRatio().call()
            return utils.wei_to_decimal(cr, 18)
        except Exception as e:
            raise ContractCallError(f"Failed to get V1 collateral ratio: {str(e)}")

    def get_v1_rebalance_pools(self) -> List[str]:
        """
        Get all registered V1 rebalance pools.
        Note: Requires RebalancePoolRegistry ABI.
        """
        try:
            # We assume a standard 'getPools' or similar function
            pools = self.v1_rebalance_registry.functions.getPools().call()
            return [utils.to_checksum_address(p) for p in pools]
        except Exception as e:
            logger.warning(f"Failed to fetch rebalance pools: {str(e)}")
            return []

    def get_v1_rebalance_pool_balances(self, pool_address: str, account_address: Optional[str] = None) -> Dict[str, Decimal]:
        """
        Get all balances for an account in a V1 rebalance pool.
        
        Returns:
            Dict: {staked, unlocked, unlocking} balances.
        """
        target_address = account_address or self.address
        pool = self._get_contract("rebalance_pool", pool_address)
        try:
            staked = pool.functions.balanceOf(target_address).call()
            unlocked = pool.functions.unlockedBalanceOf(target_address).call()
            unlocking = pool.functions.unlockingBalanceOf(target_address).call()
            return {
                "staked": utils.wei_to_decimal(staked, 18),
                "unlocked": utils.wei_to_decimal(unlocked, 18),
                "unlocking": utils.wei_to_decimal(unlocking, 18)
            }
        except Exception as e:
            raise ContractCallError(f"Failed to get rebalance pool balances: {str(e)}")

    # --- Infrastructure Read Methods ---

    def get_pool_manager_info(self, pool_address: str) -> Dict[str, Any]:
        """Get information from the Pool Manager for a specific pool."""
        contract = self._get_contract("pool_manager", constants.POOL_MANAGER)
        try:
            info = contract.functions.getPoolInfo(utils.to_checksum_address(pool_address)).call()
            return {
                "collateral_capacity": utils.wei_to_decimal(info[0]),
                "collateral_balance": utils.wei_to_decimal(info[1]),
                "debt_capacity": utils.wei_to_decimal(info[2]),
                "debt_balance": utils.wei_to_decimal(info[3]),
            }
        except Exception as e:
            raise ContractCallError(f"Failed to get pool manager info: {str(e)}")

    def get_reserve_pool_bonus_ratio(self, token_address: str) -> Decimal:
        """Get the bonus ratio for a token in the Reserve Pool."""
        contract = self._get_contract("reserve_pool", constants.RESERVE_POOL)
        try:
            ratio = contract.functions.bonusRatio(utils.to_checksum_address(token_address)).call()
            return utils.wei_to_decimal(ratio, 18)
        except Exception as e:
            raise ContractCallError(f"Failed to get reserve pool bonus ratio: {str(e)}")

    def get_steth_treasury_info(self) -> Dict[str, Any]:
        """Get information from the stETH Treasury."""
        contract = self._get_contract("steth_treasury", constants.STETH_TREASURY_PROXY)
        try:
            return {
                "total_base_token": utils.wei_to_decimal(contract.functions.totalBaseToken().call()),
                "collateral_ratio": utils.wei_to_decimal(contract.functions.collateralRatio().call(), 18),
                "leverage_ratio": utils.wei_to_decimal(contract.functions.leverageRatio().call(), 18),
            }
        except Exception as e:
            raise ContractCallError(f"Failed to get stETH treasury info: {str(e)}")

    def get_treasury_nav(self) -> Dict[str, Decimal]:
        """Get Net Asset Values from the treasury."""
        contract = self._get_contract("steth_treasury", constants.STETH_TREASURY_PROXY)
        try:
            nav = contract.functions.getCurrentNav().call()
            return {
                "base_nav": utils.wei_to_decimal(nav[0]),
                "f_nav": utils.wei_to_decimal(nav[1]),
                "x_nav": utils.wei_to_decimal(nav[2])
            }
        except Exception as e:
            raise ContractCallError(f"Failed to get treasury NAV: {str(e)}")

    def get_market_info(self, market_address: str) -> Dict[str, Any]:
        """
        Get info for a specific market (V1 or V2).
        """
        contract = self._get_contract("market", market_address)
        try:
            return {
                "collateral_ratio": utils.wei_to_decimal(contract.functions.collateralRatio().call(), 18),
                "total_collateral": utils.wei_to_decimal(contract.functions.totalCollateral().call()),
            }
        except Exception as e:
            raise ContractCallError(f"Failed to get market info: {str(e)}")

    def get_fxusd_balance(self, account_address: Optional[str] = None) -> Decimal:
        """Get the fxUSD balance of an account."""
        return self.get_token_balance(constants.FXUSD, account_address)

    def get_feth_balance(self, account_address: Optional[str] = None) -> Decimal:
        """Get the fETH balance of an account."""
        return self.get_token_balance(constants.FETH, account_address)

    def get_rusd_balance(self, account_address: Optional[str] = None) -> Decimal:
        """Get the rUSD balance of an account."""
        return self.get_token_balance(constants.RUSD, account_address)

    def get_arusd_balance(self, account_address: Optional[str] = None) -> Decimal:
        """Get the arUSD balance of an account."""
        # Note: ARUSD address needs to be verified and added to constants
        if not hasattr(constants, 'ARUSD'):
            raise ConfigurationError("arUSD address not yet configured in constants.py")
        return self.get_token_balance(constants.ARUSD, account_address)

    def get_btcusd_balance(self, account_address: Optional[str] = None) -> Decimal:
        """Get the btcUSD balance of an account."""
        return self.get_token_balance(constants.BTCUSD, account_address)

    def get_cvxusd_balance(self, account_address: Optional[str] = None) -> Decimal:
        """Get the cvxUSD balance of an account."""
        return self.get_token_balance(constants.CVXUSD, account_address)

    def get_xeth_balance(self, account_address: Optional[str] = None) -> Decimal:
        """Get the xETH balance of an account."""
        return self.get_token_balance(constants.XETH, account_address)

    def get_xcvx_balance(self, account_address: Optional[str] = None) -> Decimal:
        """Get the xCVX balance of an account."""
        return self.get_token_balance(constants.XCVX, account_address)

    def get_xwbtc_balance(self, account_address: Optional[str] = None) -> Decimal:
        """Get the xWBTC balance of an account."""
        return self.get_token_balance(constants.XWBTC, account_address)

    def get_xeeth_balance(self, account_address: Optional[str] = None) -> Decimal:
        """Get the xeETH balance of an account."""
        return self.get_token_balance(constants.XEETH, account_address)

    def get_xezeth_balance(self, account_address: Optional[str] = None) -> Decimal:
        """Get the xezETH balance of an account."""
        return self.get_token_balance(constants.XEZETH, account_address)

    def get_xsteth_balance(self, account_address: Optional[str] = None) -> Decimal:
        """Get the xstETH balance of an account."""
        return self.get_token_balance(constants.XSTETH, account_address)

    def get_xfrxeth_balance(self, account_address: Optional[str] = None) -> Decimal:
        """Get the xfrxETH balance of an account."""
        return self.get_token_balance(constants.XFRXETH, account_address)

    # --- Savings & Stability Pool Read Methods ---

    def get_fxsave_balance(self, account_address: Optional[str] = None) -> Decimal:
        """Get the fxSAVE (Saving fxUSD) balance of an account."""
        return self.get_token_balance(constants.SAVING_FXUSD, account_address)

    def get_fxsp_balance(self, account_address: Optional[str] = None) -> Decimal:
        """Get the fxSP (Stability Pool) balance of an account."""
        # This usually returns the staked amount
        return self.get_token_balance(constants.FXSP, account_address)

    def get_savings_apr(self) -> Decimal:
        """Get the current APR for fxSAVE."""
        contract = self._get_contract("saving_fxusd", constants.SAVING_FXUSD)
        try:
            # Get current APR from the savings contract
            apr = contract.functions.currentAPR().call()
            return utils.wei_to_decimal(apr, 18)
        except Exception:
            return Decimal(0)

    # --- Governance Read Methods ---

    def get_fxn_balance(self, account_address: Optional[str] = None) -> Decimal:
        """Get the FXN balance of an account."""
        return self.get_token_balance(constants.FXN, account_address)

    def get_vefxn_balance(self, account_address: Optional[str] = None) -> Decimal:
        """Get the veFXN balance of an account."""
        return self.get_token_balance(constants.VEFXN, account_address)

    def get_vefxn_locked_info(self, account_address: Optional[str] = None) -> Dict[str, Any]:
        """Get locked FXN info in veFXN."""
        target_address = account_address or self.address
        if not target_address:
            raise FXProtocolError("No account address provided or available in client.")
        
        try:
            locked = self.vefxn.functions.locked(utils.to_checksum_address(target_address)).call()
            return {
                "amount": utils.wei_to_decimal(locked[0], 18),
                "end": locked[1]
            }
        except Exception as e:
            raise ContractCallError(f"Failed to get veFXN locked info: {str(e)}")

    def get_gauge_weight(self, gauge_address: str) -> Decimal:
        """Get the relative weight of a gauge in the controller."""
        try:
            # We use the pre-loaded gauge_controller if available
            weight = self.gauge_controller.functions.get_gauge_weight(
                utils.to_checksum_address(gauge_address)
            ).call()
            # Weights are typically returned with 18 decimals in GaugeController
            return utils.wei_to_decimal(weight, 18)
        except Exception as e:
            # Fallback for when ABI is empty or call fails
            raise ContractCallError(f"Failed to get gauge weight: {str(e)}")

    def get_gauge_relative_weight(self, gauge_address: str) -> Decimal:
        """Get the relative weight of a gauge."""
        try:
            weight = self.gauge_controller.functions.gauge_relative_weight(
                utils.to_checksum_address(gauge_address)
            ).call()
            return utils.wei_to_decimal(weight, 18)
        except Exception as e:
            raise ContractCallError(f"Failed to get gauge relative weight: {str(e)}")

    def get_claimable_rewards(self, gauge_address: str, token_address: str, account_address: Optional[str] = None) -> Decimal:
        """Get claimable rewards from a gauge."""
        target_address = account_address or self.address
        if not target_address:
            raise FXProtocolError("No account address provided or available in client.")
        
        gauge = self._get_contract("liquidity_gauge", gauge_address)
        try:
            # LiquidityGauge has claimable(address, address)
            amount = gauge.functions.claimable(
                utils.to_checksum_address(target_address),
                utils.to_checksum_address(token_address)
            ).call()
            # We need to know token decimals, assuming 18 for now
            return utils.wei_to_decimal(amount, 18)
        except Exception as e:
            raise ContractCallError(f"Failed to get claimable rewards: {str(e)}")

    def get_all_balances(self, account_address: Optional[str] = None) -> Dict[str, Decimal]:
        """
        Get all protocol token balances for an account.
        
        Includes all protocol tokens: fxUSD, fETH, rUSD, btcUSD, cvxUSD, arUSD,
        and all x tokens (xETH, xCVX, xWBTC, xeETH, xezETH, xstETH, xfrxETH).
        
        Args:
            account_address: Optional account address.
            
        Returns:
            Dict[str, Decimal]: Map of token names to balances.
        """
        tokens = {
            "fxUSD": constants.FXUSD,
            "fETH": constants.FETH,
            "rUSD": constants.RUSD,
            "btcUSD": constants.BTCUSD,
            "cvxUSD": constants.CVXUSD,
            "xETH": constants.XETH,
            "xCVX": constants.XCVX,
            "xWBTC": constants.XWBTC,
            "xeETH": constants.XEETH,
            "xezETH": constants.XEZETH,
            "xstETH": constants.XSTETH,
            "xfrxETH": constants.XFRXETH,
        }
        
        balances = {}
        for name, address in tokens.items():
            try:
                balances[name] = self.get_token_balance(address, account_address)
            except Exception:
                balances[name] = Decimal(0)
                
        # Handle arUSD separately since it might be missing
        if hasattr(constants, 'ARUSD'):
            try:
                balances["arUSD"] = self.get_token_balance(constants.ARUSD, account_address)
            except Exception:
                balances["arUSD"] = Decimal(0)
                
        return balances

    def get_all_gauge_balances(self, account_address: Optional[str] = None) -> Dict[str, Decimal]:
        """
        Get all liquidity gauge stakes for an account.
        """
        balances = {}
        for name, address in constants.GAUGES.items():
            try:
                balances[name] = self.get_token_balance(address, account_address)
            except Exception:
                balances[name] = Decimal(0)
        return balances

    def mint_via_treasury(self, base_in: Union[int, float, Decimal, str], recipient: Optional[str] = None, option: int = 0) -> str:
        """
        Mint fToken and/or xToken via the Treasury.
        
        Args:
            base_in: human-readable amount of base token.
            recipient: recipient address.
            option: MintOption (0: Both, 1: fToken, 2: xToken).
        """
        contract = self._get_contract("steth_treasury", constants.STETH_TREASURY_PROXY)
        target_recipient = recipient or self.address
        raw_base_in = utils.decimal_to_wei(base_in, 18)
        
        return self._build_and_send_transaction(
            contract.functions.mint(raw_base_in, target_recipient, option)
        )

    def redeem_via_treasury(self, f_token_in: Union[int, float, Decimal, str] = 0, x_token_in: Union[int, float, Decimal, str] = 0, owner: Optional[str] = None) -> str:
        """
        Redeem fToken and/or xToken for base collateral via Treasury.
        """
        contract = self._get_contract("steth_treasury", constants.STETH_TREASURY_PROXY)
        target_owner = owner or self.address
        raw_f_in = utils.decimal_to_wei(f_token_in, 18)
        raw_x_in = utils.decimal_to_wei(x_token_in, 18)
        
        return self._build_and_send_transaction(
            contract.functions.redeem(raw_f_in, raw_x_in, target_owner)
        )

    def harvest_treasury(self) -> str:
        """
        Harvest rewards from the Treasury.
        """
        contract = self._get_contract("steth_treasury", constants.STETH_TREASURY_PROXY)
        return self._build_and_send_transaction(contract.functions.harvest())

    def initialize_v2_treasury(self, sample_interval: int) -> str:
        """
        Initialize V2 features for the Treasury.
        """
        contract = self._get_contract("steth_treasury", constants.STETH_TREASURY_PROXY)
        return self._build_and_send_transaction(contract.functions.initializeV2(sample_interval))

    # --- Write Methods ---

    def _build_and_send_transaction(self, contract_function, value: int = 0) -> str:
        """
        Internal helper to build, sign, and send a transaction.
        
        Supports both private key signing and browser wallet signing.
        
        Args:
            contract_function: The contract function to call.
            value: Optional ETH value to send with the transaction (in Wei).
            
        Returns:
            str: The transaction hash.
        """
        if not self.account and not self.use_browser_wallet:
            raise ConfigurationError(
                "Private key or browser wallet required for write operations. "
                "Provide a private key, set FX_PROTOCOL_PRIVATE_KEY environment variable, "
                "or use use_browser_wallet=True with a browser wallet extension."
            )

        if not self.address:
            raise ConfigurationError("No account address available for transaction.")

        nonce = self.w3.eth.get_transaction_count(self.address)
        
        tx_params = {
            'from': self.address,
            'nonce': nonce,
            'value': value,
            'gasPrice': self.w3.eth.gas_price
        }

        try:
            # Estimate gas
            tx_params['gas'] = contract_function.estimate_gas(tx_params)
            
            # Build transaction
            built_tx = contract_function.build_transaction(tx_params)
            
            # Sign and send based on wallet type
            if self.use_browser_wallet:
                # Browser wallet: send_transaction prompts user in browser
                # The provider (MetaMask, etc.) handles signing
                tx_hash = self.w3.eth.send_transaction(built_tx)
            else:
                # Private key: sign locally and send raw transaction
                signed_tx = self.w3.eth.account.sign_transaction(built_tx, self.account.key)
                tx_hash = self.w3.eth.send_raw_transaction(signed_tx.rawTransaction)
            
            # Wait for receipt
            receipt = self.w3.eth.wait_for_transaction_receipt(tx_hash)
            
            if receipt.status != 1:
                raise TransactionFailedError(f"Transaction failed: {tx_hash.hex()}")
                
            return tx_hash.hex()
            
        except Exception as e:
            if isinstance(e, TransactionFailedError):
                raise
            raise TransactionFailedError(f"Failed to send transaction: {str(e)}")

    def approve(self, token_address: str, spender_address: str, amount: Union[int, float, Decimal, str]) -> str:
        """
        Approve a spender to spend tokens.
        
        Args:
            token_address: The address of the token contract.
            spender_address: The address of the spender.
            amount: Human-readable amount to approve.
            
        Returns:
            str: Transaction hash.
        """
        contract = self.w3.eth.contract(
            address=utils.to_checksum_address(token_address),
            abi=[
                {"constant": False, "inputs": [{"name": "_spender", "type": "address"}, {"name": "_value", "type": "uint256"}], "name": "approve", "outputs": [{"name": "", "type": "bool"}], "type": "function"},
                {"constant": True, "inputs": [], "name": "decimals", "outputs": [{"name": "", "type": "uint8"}], "type": "function"}
            ]
        )
        
        decimals = contract.functions.decimals().call()
        raw_amount = utils.decimal_to_wei(amount, decimals)
        
        return self._build_and_send_transaction(contract.functions.approve(spender_address, raw_amount))

    def transfer(self, token_address: str, recipient_address: str, amount: Union[int, float, Decimal, str]) -> str:
        """
        Transfer tokens to a recipient.
        
        Args:
            token_address: The address of the token contract.
            recipient_address: The address of the recipient.
            amount: Human-readable amount to transfer.
            
        Returns:
            str: Transaction hash.
        """
        contract = self.w3.eth.contract(
            address=utils.to_checksum_address(token_address),
            abi=[
                {"constant": False, "inputs": [{"name": "_to", "type": "address"}, {"name": "_value", "type": "uint256"}], "name": "transfer", "outputs": [{"name": "", "type": "bool"}], "type": "function"},
                {"constant": True, "inputs": [], "name": "decimals", "outputs": [{"name": "", "type": "uint8"}], "type": "function"}
            ]
        )
        
        decimals = contract.functions.decimals().call()
        raw_amount = utils.decimal_to_wei(amount, decimals)
        
        return self._build_and_send_transaction(contract.functions.transfer(recipient_address, raw_amount))

    # --- f(x) Protocol Specific Write Methods ---

    def mint_f_token(self, market_address: str, base_in: Union[int, float, Decimal, str], recipient: Optional[str] = None, min_f_token_out: Union[int, float, Decimal, str] = 0) -> str:
        """
        Mint fToken using base collateral (e.g., stETH).
        
        Args:
            market_address: The address of the market contract.
            base_in: Human-readable amount of base token to provide.
            recipient: Optional recipient address (defaults to client address).
            min_f_token_out: Human-readable minimum fToken to receive.
        """
        contract = self._get_contract("market", market_address)
        target_recipient = recipient or self.address
        
        # We need to know the decimals of the base token, usually 18 for stETH
        # For simplicity, we assume 18 here, but a robust version would check
        raw_base_in = utils.decimal_to_wei(base_in, 18)
        raw_min_out = utils.decimal_to_wei(min_f_token_out, 18)
        
        return self._build_and_send_transaction(
            contract.functions.mintFToken(raw_base_in, target_recipient, raw_min_out)
        )

    def mint_x_token(self, market_address: str, base_in: Union[int, float, Decimal, str], recipient: Optional[str] = None, min_x_token_out: Union[int, float, Decimal, str] = 0) -> str:
        """
        Mint xToken using base collateral.
        """
        contract = self._get_contract("market", market_address)
        target_recipient = recipient or self.address
        
        raw_base_in = utils.decimal_to_wei(base_in, 18)
        raw_min_out = utils.decimal_to_wei(min_x_token_out, 18)
        
        return self._build_and_send_transaction(
            contract.functions.mintXToken(raw_base_in, target_recipient, raw_min_out)
        )

    def mint_both_tokens(self, market_address: str, base_in: Union[int, float, Decimal, str], recipient: Optional[str] = None, min_f_token_out: Union[int, float, Decimal, str] = 0, min_x_token_out: Union[int, float, Decimal, str] = 0) -> str:
        """
        Mint both fToken and xToken proportionally.
        """
        contract = self._get_contract("market", market_address)
        target_recipient = recipient or self.address
        
        raw_base_in = utils.decimal_to_wei(base_in, 18)
        raw_min_f_out = utils.decimal_to_wei(min_f_token_out, 18)
        raw_min_x_out = utils.decimal_to_wei(min_x_token_out, 18)
        
        return self._build_and_send_transaction(
            contract.functions.mint(raw_base_in, target_recipient, raw_min_f_out, raw_min_x_out)
        )

    def redeem(self, market_address: str, f_token_in: Union[int, float, Decimal, str] = 0, x_token_in: Union[int, float, Decimal, str] = 0, recipient: Optional[str] = None, min_base_out: Union[int, float, Decimal, str] = 0) -> str:
        """
        Redeem fToken and/or xToken for base collateral.
        """
        contract = self._get_contract("market", market_address)
        target_recipient = recipient or self.address
        
        raw_f_in = utils.decimal_to_wei(f_token_in, 18)
        raw_x_in = utils.decimal_to_wei(x_token_in, 18)
        raw_min_base_out = utils.decimal_to_wei(min_base_out, 18)
        
        return self._build_and_send_transaction(
            contract.functions.redeem(raw_f_in, raw_x_in, target_recipient, raw_min_base_out)
        )

    # --- Governance Write Methods ---

    def deposit_to_vefxn(self, amount: Union[int, float, Decimal, str], unlock_time: int) -> str:
        """
        Deposit FXN into veFXN to get voting power.
        
        Args:
            amount: Amount of FXN to deposit.
            unlock_time: Unix timestamp for the unlock date.
        """
        raw_amount = utils.decimal_to_wei(amount, 18)
        # Note: 'create_lock' is the standard function name in VotingEscrow
        return self._build_and_send_transaction(
            self.vefxn.functions.create_lock(raw_amount, unlock_time)
        )

    def vote_for_gauge_weight(self, gauge_address: str, user_weight: int) -> str:
        """
        Vote for a gauge's weight in the GaugeController.
        
        Args:
            gauge_address: The address of the gauge.
            user_weight: The weight to assign (out of 10000).
        """
        return self._build_and_send_transaction(
            self.gauge_controller.functions.vote_for_gauge_weights(
                utils.to_checksum_address(gauge_address),
                user_weight
            )
        )

    def claim_gauge_rewards(self, gauge_address: str, account: Optional[str] = None) -> str:
        """
        Claim rewards from a liquidity gauge.
        
        Args:
            gauge_address: The address of the gauge.
            account: Optional account address to claim for.
        """
        gauge = self._get_contract("liquidity_gauge", gauge_address)
        target_account = account or self.address
        if not target_account:
            raise FXProtocolError("No account address provided or available in client.")
            
        return self._build_and_send_transaction(gauge.functions.claim(target_account))

    def claim_all_gauge_rewards(self) -> List[str]:
        """
        Claim rewards from all configured gauges.
        """
        tx_hashes = []
        for name, address in constants.GAUGES.items():
            try:
                logger.info(f"Claiming rewards for gauge: {name}")
                tx_hash = self.claim_gauge_rewards(address)
                tx_hashes.append(tx_hash)
            except Exception as e:
                logger.warning(f"Failed to claim rewards for gauge {name}: {str(e)}")
        return tx_hashes

    def operate_position(self, pool_address: str, position_id: int, new_collateral: Union[int, float, Decimal, str], new_debt: Union[int, float, Decimal, str]) -> str:
        """
        Operate on a position in the Pool Manager.
        
        Args:
            pool_address: The address of the pool.
            position_id: The ID of the position.
            new_collateral: The new human-readable collateral amount (delta).
            new_debt: The new human-readable debt amount (delta).
        """
        contract = self._get_contract("pool_manager", constants.POOL_MANAGER)
        raw_coll = utils.decimal_to_wei(new_collateral, 18)
        raw_debt = utils.decimal_to_wei(new_debt, 18)
        
        return self._build_and_send_transaction(
            contract.functions.operate(
                utils.to_checksum_address(pool_address),
                position_id,
                raw_coll,
                raw_debt
            )
        )

    def rebalance_position(self, pool_address: str, receiver: str, position_id: int, max_fxusd: Union[int, float, Decimal, str], max_stable: Union[int, float, Decimal, str]) -> str:
        """
        Rebalance a position in the Pool Manager.
        """
        contract = self._get_contract("pool_manager", constants.POOL_MANAGER)
        raw_fxusd = utils.decimal_to_wei(max_fxusd, 18)
        raw_stable = utils.decimal_to_wei(max_stable, 18)
        
        return self._build_and_send_transaction(
            contract.functions.rebalance(
                utils.to_checksum_address(pool_address),
                utils.to_checksum_address(receiver),
                position_id,
                raw_fxusd,
                raw_stable
            )
        )

    def liquidate_position(self, pool_address: str, receiver: str, position_id: int, max_fxusd: Union[int, float, Decimal, str], max_stable: Union[int, float, Decimal, str]) -> str:
        """
        Liquidate a position in the Pool Manager.
        """
        contract = self._get_contract("pool_manager", constants.POOL_MANAGER)
        raw_fxusd = utils.decimal_to_wei(max_fxusd, 18)
        raw_stable = utils.decimal_to_wei(max_stable, 18)
        
        return self._build_and_send_transaction(
            contract.functions.liquidate(
                utils.to_checksum_address(pool_address),
                utils.to_checksum_address(receiver),
                position_id,
                raw_fxusd,
                raw_stable
            )
        )

    def harvest_pool_manager(self, pool_address: str) -> str:
        """
        Harvest rewards for a pool in the Pool Manager.
        """
        contract = self._get_contract("pool_manager", constants.POOL_MANAGER)
        return self._build_and_send_transaction(
            contract.functions.harvest(utils.to_checksum_address(pool_address))
        )

    def get_position_info(self, position_id: int) -> Dict[str, Any]:
        """
        Get details for a specific position in the Pool Manager.
        
        Returns:
            Dict: collateral, debt, and owner info.
        """
        contract = self._get_contract("pool_manager", constants.POOL_MANAGER)
        try:
            info = contract.functions.getPosition(position_id).call()
            # Assuming returns (owner, collateral, debt) based on V2 common patterns
            return {
                "owner": info[0],
                "collateral": utils.wei_to_decimal(info[1]),
                "debt": utils.wei_to_decimal(info[2])
            }
        except Exception as e:
            raise ContractCallError(f"Failed to get position info: {str(e)}")

    def get_peg_keeper_info(self) -> Dict[str, Any]:
        """Get the current status from the Peg Keeper."""
        contract = self._get_contract("peg_keeper", constants.PEG_KEEPER)
        try:
            return {
                "is_active": contract.functions.isActive().call(),
                "debt_ceiling": utils.wei_to_decimal(contract.functions.debtCeiling().call()),
                "total_debt": utils.wei_to_decimal(contract.functions.totalDebt().call()),
            }
        except Exception as e:
            raise ContractCallError(f"Failed to get peg keeper info: {str(e)}")

    def flash_loan(self, token_address: str, amount: Union[int, float, Decimal, str], receiver: str, data: bytes = b"") -> str:
        """
        Execute a flash loan from the Pool Manager.
        """
        contract = self._get_contract("pool_manager", constants.POOL_MANAGER)
        raw_amount = utils.decimal_to_wei(amount, 18)
        
        return self._build_and_send_transaction(
            contract.functions.flashLoan(
                utils.to_checksum_address(receiver),
                utils.to_checksum_address(token_address),
                raw_amount,
                data
            )
        )

    def request_reserve_pool_bonus(self, token_address: str, recipient: str, original_amount: Union[int, float, Decimal, str]) -> str:
        """
        Request a bonus from the Reserve Pool.
        """
        contract = self._get_contract("reserve_pool", constants.RESERVE_POOL)
        raw_amount = utils.decimal_to_wei(original_amount, 18)
        
        return self._build_and_send_transaction(
            contract.functions.requestBonus(
                utils.to_checksum_address(token_address),
                utils.to_checksum_address(recipient),
                raw_amount
            )
        )

    # --- V1 Write Methods ---

    def deposit_to_rebalance_pool(self, pool_address: str, amount: Union[int, float, Decimal, str], recipient: Optional[str] = None) -> str:
        """
        Deposit assets into a V1 rebalance pool.
        
        Args:
            pool_address: The address of the rebalance pool.
            amount: Human-readable amount to deposit.
            recipient: Optional recipient address (defaults to client address).
        """
        pool = self._get_contract("rebalance_pool", pool_address)
        target_recipient = recipient or self.address
        raw_amount = utils.decimal_to_wei(amount, 18)
        return self._build_and_send_transaction(pool.functions.deposit(raw_amount, utils.to_checksum_address(target_recipient)))

    def unlock_rebalance_pool_assets(self, pool_address: str, amount: Union[int, float, Decimal, str]) -> str:
        """
        Start the unlock process for assets in a V1 rebalance pool.
        
        Args:
            pool_address: The address of the rebalance pool.
            amount: Human-readable amount to unlock.
        """
        pool = self._get_contract("rebalance_pool", pool_address)
        raw_amount = utils.decimal_to_wei(amount, 18)
        return self._build_and_send_transaction(pool.functions.unlock(raw_amount))

    def withdraw_unlocked_rebalance_pool_assets(self, pool_address: str, claim_rewards: bool = True) -> str:
        """
        Withdraw unlocked assets from a V1 rebalance pool.
        
        Args:
            pool_address: The address of the rebalance pool.
            claim_rewards: Whether to claim rewards during withdrawal.
        """
        pool = self._get_contract("rebalance_pool", pool_address)
        return self._build_and_send_transaction(pool.functions.withdrawUnlocked(claim_rewards))

    def claim_rebalance_pool_rewards(self, pool_address: str, tokens: List[str]) -> str:
        """
        Claim rewards from a V1 rebalance pool.
        
        Args:
            pool_address: The address of the rebalance pool.
            tokens: List of reward token addresses to claim.
        """
        pool = self._get_contract("rebalance_pool", pool_address)
        return self._build_and_send_transaction(pool.functions.claim([utils.to_checksum_address(t) for t in tokens]))

    # --- Supporting Write Methods ---

    def swap(self, token_in: str, amount_in: Union[int, float, Decimal, str], encoding: int, routes: List[int]) -> str:
        """
        Swap tokens using the MultiPathConverter.
        
        Args:
            token_in: The address of the token to swap from.
            amount_in: Human-readable amount to swap.
            encoding: Encoding for the converter.
            routes: List of routes for the swap.
        """
        raw_amount_in = utils.decimal_to_wei(amount_in, 18)  # Adjust decimals as needed
        return self._build_and_send_transaction(
            self.multipath_converter.functions.convert(
                utils.to_checksum_address(token_in),
                raw_amount_in,
                encoding,
                routes
            )
        )

    def mint_f_token_via_gateway(self, amount_eth: Union[int, float, Decimal, str], min_f_token_out: Union[int, float, Decimal, str] = 0) -> str:
        """
        Deposit ETH/stETH through the stETH Gateway to mint fToken.
        
        Args:
            amount_eth: Human-readable amount of ETH to send.
            min_f_token_out: Human-readable minimum fToken to receive.
        """
        raw_amount_eth = utils.decimal_to_wei(amount_eth, 18)
        raw_min_out = utils.decimal_to_wei(min_f_token_out, 18)
        
        return self._build_and_send_transaction(
            self.steth_gateway.functions.mintFToken(raw_min_out),
            value=raw_amount_eth
        )

    def mint_x_token_via_gateway(self, amount_eth: Union[int, float, Decimal, str], min_x_token_out: Union[int, float, Decimal, str] = 0) -> str:
        """
        Deposit ETH/stETH through the stETH Gateway to mint xToken.
        """
        raw_amount_eth = utils.decimal_to_wei(amount_eth, 18)
        raw_min_out = utils.decimal_to_wei(min_x_token_out, 18)
        
        return self._build_and_send_transaction(
            self.steth_gateway.functions.mintXToken(raw_min_out),
            value=raw_amount_eth
        )

    # --- Savings Write Methods ---

    def deposit_fxsave(self, amount: Union[int, float, Decimal, str]) -> str:
        """
        Deposit fxUSD into fxSAVE.
        """
        contract = self._get_contract("saving_fxusd", constants.SAVING_FXUSD)
        raw_amount = utils.decimal_to_wei(amount, 18)
        return self._build_and_send_transaction(contract.functions.deposit(raw_amount))

    def redeem_fxsave(self, amount: Union[int, float, Decimal, str]) -> str:
        """
        Redeem fxUSD from fxSAVE.
        """
        contract = self._get_contract("saving_fxusd", constants.SAVING_FXUSD)
        raw_amount = utils.decimal_to_wei(amount, 18)
        return self._build_and_send_transaction(contract.functions.withdraw(raw_amount))

    # --- Stability Pool Write Methods ---

    def deposit_to_stability_pool(self, amount: Union[int, float, Decimal, str]) -> str:
        """
        Deposit assets into the stability pool (fxSP).
        """
        contract = self._get_contract("fxsp", constants.FXSP)
        raw_amount = utils.decimal_to_wei(amount, 18)
        return self._build_and_send_transaction(contract.functions.deposit(raw_amount))

    def withdraw_from_stability_pool(self, amount: Union[int, float, Decimal, str]) -> str:
        """
        Withdraw assets from the stability pool (fxSP).
        """
        contract = self._get_contract("fxsp", constants.FXSP)
        raw_amount = utils.decimal_to_wei(amount, 18)
        return self._build_and_send_transaction(contract.functions.withdraw(raw_amount))

    # --- Vesting Write Methods ---

    def claim_fxn_vesting(self) -> str:
        """Claim vested FXN tokens."""
        vesting = self._get_contract("vesting", constants.VESTING_FXN)
        return self._build_and_send_transaction(vesting.functions.claim())

    def claim_feth_vesting(self) -> str:
        """Claim vested fETH tokens."""
        vesting = self._get_contract("vesting", constants.VESTING_FETH)
        return self._build_and_send_transaction(vesting.functions.claim())

    def claim_fxusd_vesting(self) -> str:
        """Claim vested fxUSD tokens."""
        vesting = self._get_contract("vesting", constants.VESTING_FXUSD)
        return self._build_and_send_transaction(vesting.functions.claim())

