import MetaTrader5 as mt5
import pandas as pd
from enum import Enum


class MT5Period(Enum):
    M1 = mt5.TIMEFRAME_M1
    M2 = mt5.TIMEFRAME_M2
    M3 = mt5.TIMEFRAME_M3
    M4 = mt5.TIMEFRAME_M4
    M5 = mt5.TIMEFRAME_M5
    M6 = mt5.TIMEFRAME_M6
    M10 = mt5.TIMEFRAME_M10
    M12 = mt5.TIMEFRAME_M12
    M15 = mt5.TIMEFRAME_M15
    M20 = mt5.TIMEFRAME_M20
    M30 = mt5.TIMEFRAME_M30
    H1 = mt5.TIMEFRAME_H1
    H2 = mt5.TIMEFRAME_H2
    H3 = mt5.TIMEFRAME_H3
    H4 = mt5.TIMEFRAME_H4
    H6 = mt5.TIMEFRAME_H6
    H8 = mt5.TIMEFRAME_H8
    H12 = mt5.TIMEFRAME_H12
    D1 = mt5.TIMEFRAME_D1
    W1 = mt5.TIMEFRAME_W1
    MN1 = mt5.TIMEFRAME_MN1
    
    @classmethod
    def get_name(cls, value):
        try:
            return cls(value).name
        except ValueError:
            return str(value)

class MT5Config:
    def __init__(self, login,password,server):
        self.login = login
        self.password = password
        self.server = server

class MT5Connector:
    
    def __init__(self, config: MT5Config):
        self.config = config
        self.login = config.login
        self.password = config.password
        self.server = config.server
        self.is_connected = False
        self.is_authorized = False


    def connect(self):
        try:
            if not mt5.initialize():
                print("MT5初始化失败")
                return False
            
            self.is_connected = True
            
            if self.login and self.password:
                return self.login_account()
            else:
                return True
                
        except Exception as e:
            print(f"连接MT5失败: {e}")
            return False
    
    def login_account(self):
        try:
            authorized = mt5.login(self.login, self.password, self.server)
            if authorized:
                self.is_authorized = True
                return True
            else:
                print(f"账户登录失败: {mt5.last_error()}")
                return False
        except Exception as e:
            print(f"登录失败: {e}")
            return False
    
    def disconnect(self):
        try:
            mt5.shutdown()
            self.is_connected = False
            self.is_authorized = False
        except Exception as e:
            print(f"断开连接失败: {e}")
    
    def get_symbols(self, limit=10):
        if not self.is_connected:
            print("未连接到MT5")
            return []
        
        try:
            symbols = mt5.symbols_get()
            return symbols
        except Exception as e:
            print(f"获取交易品种失败: {e}")
            return []
    
    def get_symbol_info(self, symbol):
        if not self.is_connected:
            print("未连接到MT5")
            return None
        
        try:
            symbol_info = mt5.symbol_info(symbol)
            if symbol_info is not None:
                print(f"{symbol}点差: {symbol_info.spread}")
                print(f"合约大小: {symbol_info.trade_contract_size}")
                return symbol_info
            else:
                print(f"未找到品种: {symbol}")
                return None
        except Exception as e:
            print(f"获取品种信息失败: {e}")
            return None
    
    def enable_symbols(self, symbols):
        if not self.is_connected:
            print("未连接到MT5")
            return False
        
        try:
            enabled_count = 0
            for symbol in symbols:
                symbol_info = mt5.symbol_info(symbol)
                if symbol_info is None or not symbol_info.visible:
                    if mt5.symbol_select(symbol, True):
                        enabled_count += 1
                    else:
                        print(f"启用品种失败: {symbol}")
                else:
                    enabled_count += 1

            print(f"成功启用 {enabled_count}/{len(symbols)} 个品种")
            return enabled_count == len(symbols)
        except Exception as e:
            print(f"启用品种失败: {e}")
            return False
    
    def get_symbol_tick(self, symbols):
        if not self.is_connected:
            print("未连接到MT5")
            return {}
        
        try:
            rates = {}
            for symbol in symbols:
                tick = mt5.symbol_info_tick(symbol)
                
                if tick is not None:
                    rates[symbol] = {
                        'bid': tick.bid,
                        'ask': tick.ask,
                        'last': tick.last,
                        'volume': tick.volume,
                        'time': pd.to_datetime(tick.time, unit='s'),
                        'spread': tick.ask - tick.bid
                    }
                else:
                    print(f"获取 {symbol} 报价失败")
            return rates
        except Exception as e:
            print(f"获取实时报价失败: {e}")
            return {}
    
    def get_kline_data(self, symbol, period, count=1000, start_pos=0):
        if not self.is_connected:
            print("未连接到MT5")
            return None
        
        try:
            rates = mt5.copy_rates_from_pos(symbol, period, start_pos, count)
            
            if rates is None or len(rates) == 0:
                print(f"获取 {symbol} K线数据失败 (返回为空)")
                return None
                
            df = pd.DataFrame(rates)
            df['time'] = pd.to_datetime(df['time'], unit='s')
            df['period'] = MT5Period.get_name(period)
            
            return df
        except Exception as e:
            print(f"获取K线数据异常: {e}")
            return None

    def get_kline_range(self, symbol, period, date_from, date_to):
        if not self.is_connected:
            print("未连接到MT5")
            return None
            
        try:
            rates = mt5.copy_rates_range(symbol, period, date_from, date_to)
            
            if rates is None or len(rates) == 0:
                print(f"获取 {symbol} 范围K线数据失败")
                return None
                
            df = pd.DataFrame(rates)
            df['time'] = pd.to_datetime(df['time'], unit='s')
            df['period'] = MT5Period.get_name(period)
            return df
        except Exception as e:
            print(f"获取范围K线数据异常: {e}")
            return None
    
    def __enter__(self):
        self.connect()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.disconnect()


def main():
    with MT5Connector() as conn:
        symbols = ["EURUSD", "GBPUSD", "XAUUSD", "BTCUSD", "ETHUSD"]
        
        conn.enable_symbols(symbols)
        
        rates = conn.get_symbol_tick(symbols)
        
        if rates:
            df = pd.DataFrame(rates).T
            print("\n实时报价汇总:")
            print(df)
            
        print("\n获取 XAUUSD 最近10根M1 K线:")
        klines = conn.get_kline_data("XAUUSD", period=mt5.TIMEFRAME_M1, count=10)
        if klines is not None:
            print(klines)

if __name__ == '__main__':
    main()
