from mexc_sdk import Spot
from dotenv import load_dotenv
import yaml
from pymongo import MongoClient
from pymongo.errors import ConnectionFailure, OperationFailure
import pytz

import logging
import os
from datetime import datetime
import time
from urllib.parse import quote_plus


# 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 MongoDBConn(object):
    def __init__(self):
        """Initialize MongoDB connection parameters"""
        self.username = quote_plus(os.getenv("DB_USER"))
        self.password = quote_plus(os.getenv("DB_PWD"))
        self.host = os.getenv("DB_HOST")
        self.port = os.getenv("DB_PORT")
        self.db_name = os.getenv("DB_NAME")
        self.client = None
        self.db = None

    def connect(self):
        """Establish connection to MongoDB"""
        try:
            # Create connection URI
            uri = f"mongodb://{self.username}:{self.password}@{self.host}:{self.port}"

            # Connect to MongoDB
            self.client = MongoClient(uri)

            # Test the connection
            self.client.admin.command("ping")

            # Get database reference
            self.db = self.client[self.db_name]

            logger.info("Successfully connected to MongoDB")
            return True

        except ConnectionFailure as e:
            logger.error(f"Could not connect to MongoDB: {e}")
            return False
        except OperationFailure as e:
            logger.error(f"Authentication failed: {e}")
            return False
        except Exception as e:
            logger.error(f"An error occurred: {e}")
            return False

    def close(self):
        """Close the MongoDB connection"""
        if self.client:
            self.client.close()
            logger.info("MongoDB connection closed")

    def insert_coin_price(self, collection_name, document):
        """Insert a single document into a collection"""
        try:
            collection = self.db[collection_name]
            result = collection.insert_one(document)
            logger.info(f"Document inserted with ID: {result.inserted_id}")
            return result.inserted_id
        except Exception as e:
            logger.error(f"Error inserting document: {e}")
            return None


class TradeBot(object):
    def __init__(self, api_key, api_secret, mdb: MongoDBConn):
        self.client = Spot(api_key=api_key, api_secret=api_secret)
        self._load_coin()
        self.timezone = pytz.timezone("Asia/Seoul")
        self.mdb = mdb

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

    def _get_current_time(self):
        return datetime.now(self.timezone)

    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_all_interset_market_price(self):
        for coin in self.interest_coins:
            data = {}
            data["price"] = self._get_market_price(symbol=coin + self.base)
            data["timestamp"] = self._get_current_time()
            self.mdb.insert_coin_price(collection_name=coin, document=data)

    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

    def _load_coin(self):
        data = None
        with open("coin.yaml", "r") as file:
            data = yaml.safe_load(file)
        self.interest_coins = data["interest"]
        self.base = data["base"][0]


def main(tb: TradeBot):
    while True:
        tb.get_all_interset_market_price()
        time.sleep(5)


if __name__ == "__main__":
    mdb = MongoDBConn()
    mdb.connect()
    api_key = os.getenv("API_KEY")
    api_secret = os.getenv("API_SECRET")
    tb = TradeBot(api_key, api_secret, mdb)
    main(tb)