from mexc_sdk import Spot
from dotenv import load_dotenv
import logging
import os


# Set up logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

# Load the environment variables from the .env file
load_dotenv()


class TradeBot(object):
    def __init__(self, api_key, api_secret):
        self.client = Spot(api_key=api_key, api_secret=api_secret)
        self.symbol = "BTCUSDT"

    def change_symbol(self, new_symbol):
        self.symbol = new_symbol

    def time(self):
        # Format is in dict
        return f"Server Time: {self.client.time()['serverTime']}"

    def get_market_price(self, symbol=None):
        """Get current market price for a symbol"""
        try:
            symbol = symbol or self.symbol
            ticker = self.client.ticker_price(symbol)
            price = float(ticker["price"])
            logger.info(f"Current {symbol} price: {price}")
            return price
        except Exception as e:
            logger.error(f"Error getting market price: {str(e)}")
            return None

    def get_account_balance(self):
        """Get account balance for all assets"""
        try:
            account_info = self.client.account_info()
            balances = account_info["balances"]

            logger.info("Account balances:")
            for balance in balances:
                if float(balance["free"]) > 0 or float(balance["locked"]) > 0:
                    logger.info(
                        f"{balance['asset']}: Free={balance['free']}, Locked={balance['locked']}"
                    )

            return balances
        except Exception as e:
            logger.error(f"Error getting account balance: {str(e)}")
            return None


if __name__ == "__main__":
    # Access the variables
    api_key = os.getenv("API_KEY")
    api_secret = os.getenv("API_SECRET")
    tb = TradeBot(api_key, api_secret)
    print(tb.time())
    tb.get_market_price()
    tb.get_account_balance()