import aiohttp
import time
import asyncio
import zlib
import json, ujson
import zlib
import hashlib
import hmac
import base64
import traceback
import random, csv, sys
import logging, logging.handlers
import utils
import model

def empty_call(msg):
    pass


def inflate(data):
    '''
        解压缩数据
    '''
    decompress = zlib.decompressobj(-zlib.MAX_WBITS)
    inflated = decompress.decompress(data)
    inflated += decompress.flush()
    return inflated

class BinanceCoinSwapWs:

    def __init__(self, params: model.ClientParams, colo=0, is_print=0):
        if colo:
            print('不支持colo高速线路')
            self.URL = 'wss://dstream.binance.com/ws/'
        else:
            self.URL = 'wss://dstream.binance.com/ws/'
        self.params = params
        self.name = self.params.name
        self.base = self.params.pair.split('_')[0].upper()
        self.quote = self.params.pair.split('_')[1].upper()
        self.symbol = self.base + self.quote
        if len(self.params.pair.split('_')) > 2:
            self.delivery = self.params.pair.split('_')[2] # 210924
            self.symbol += f"_{self.delivery}"
        else:
            self.symbol += '_PERP'
        self.callback = {
            "onMarket":self.save_market,
            "onPosition":empty_call,
            "onEquity":empty_call,
            "onOrder":empty_call,
            "onTicker":empty_call,
            "onDepth":empty_call,
            "onExit":empty_call,
            }
        self.is_print = is_print
        self.proxy = None
        if 'win' in sys.platform:
            self.proxy = self.params.proxy
        self.logger = self.get_logger()
        self.multiplier = None
        self.stop_flag = 0
        self.max_buy = 0.0
        self.min_sell = 0.0
        self.buy_v = 0.0
        self.buy_q = 0.0
        self.sell_v = 0.0
        self.sell_q = 0.0
        self.ticker_info = {"name":self.name,'bp':0.0,'ap':0.0}
        self.depth = []
        #### 指定发包ip
        iplist = utils.get_local_ip_list()
        self.ip = iplist[int(self.params.ip)]
    
    def get_logger(self):
        logger = logging.getLogger(__name__)
        logger.setLevel(logging.DEBUG)
        # log to txt
        formatter = logging.Formatter('[%(asctime)s] - %(levelname)s - %(message)s')
        handler = logging.handlers.RotatingFileHandler(f"log.log",maxBytes=1024*1024)
        handler.setLevel(logging.DEBUG)
        handler.setFormatter(formatter)
        logger.addHandler(handler)
        return logger

    def save_market(self, msg):
        date = time.strftime('%Y-%m-%d',time.localtime())
        interval = self.params.interval
        if msg:
            exchange = msg['name']
            if len(msg['data']) > 1:
                with open(f'./history/{exchange}_{self.symbol}_{interval}_{date}.csv',
                            'a',
                            newline='',
                            encoding='utf-8') as f:
                    writer = csv.writer(f, delimiter=',')
                    writer.writerow(msg['data'])
                if self.is_print:print(f'写入行情 {self.symbol}')

    async def get_sign(self):
        headers = {}
        headers['Content-Type'] = 'application/json'
        headers['X-MBX-APIKEY'] = self.params.access_key
        params = {
            'timestamp':int(time.time())*1000,
            'recvWindow':5000,
        }
        query_string = "&".join(["{}={}".format(k, params[k]) for k in sorted(params.keys())])
        signature = hmac.new(self.params.secret_key.encode(), msg=query_string.encode(), digestmod=hashlib.sha256).hexdigest()
        params['signature']=signature
        url = 'https://dapi.binance.com/dapi/v1/listenKey'
        session = aiohttp.ClientSession()
        response = await session.post(
            url, 
            params=params,
            headers=headers, 
            timeout=5, 
            proxy=self.proxy
            )
        login_str = await response.text()
        await session.close()
        return ujson.loads(login_str)['listenKey']

    async def get_depth_flash(self):
        headers = {}
        headers['Content-Type'] = 'application/json'
        headers['X-MBX-APIKEY'] = self.params.access_key
        url = f'https://dapi.binance.com/dapi/v1/depth?symbol={self.symbol}&limit=1000'
        session = aiohttp.ClientSession()
        response = await session.get(
            url, 
            headers=headers, 
            timeout=5, 
            proxy=self.proxy
            )
        depth_flash = await response.text()
        await session.close()
        return ujson.loads(depth_flash)

    def _update_ticker(self, msg):
        msg = ujson.loads(msg)
        bp = float(msg['b'])
        bq = float(msg['B'])
        ap = float(msg['a'])
        aq = float(msg['A'])
        self.ticker_info['bp'] = bp
        self.ticker_info['ap'] = ap
        self.callback['onTicker'](self.ticker_info)
        #### 标准化深度
        self.depth = [bp,bq,ap,aq]
        self.callback['onDepth']({'name':self.name,'data':self.depth})
        
    def _update_depth(self, msg):
        msg = ujson.loads(msg)
        ##### on ticker event
        self.ticker_info['bp'] = float(msg['b'][0][0])
        self.ticker_info['ap'] = float(msg['a'][0][0])
        self.callback['onTicker'](self.ticker_info)
        ##### 标准化深度
        mp = (self.ticker_info["bp"] + self.ticker_info["ap"])*0.5
        step = mp * utils.EFF_RANGE / utils.LEVEL
        bp = []
        ap = []
        bv = [0 for _ in range(utils.LEVEL)]
        av = [0 for _ in range(utils.LEVEL)]
        for i in range(utils.LEVEL):
            bp.append(self.ticker_info["bp"]-step*i)
        for i in range(utils.LEVEL):
            ap.append(self.ticker_info["ap"]+step*i)
        # 
        price_thre = self.ticker_info["bp"] - step
        index = 0
        for bid in msg['b']:
            price = float(bid[0])
            amount = float(bid[1])
            if price > price_thre:
                bv[index] += amount
            else:
                price_thre -= step
                index += 1
                if index == utils.LEVEL:
                    break
                bv[index] += amount
        price_thre = self.ticker_info["ap"] + step
        index = 0
        for ask in msg['a']:
            price = float(ask[0])
            amount = float(ask[1])
            if price < price_thre:
                av[index] += amount
            else:
                price_thre += step
                index += 1
                if index == utils.LEVEL:
                    break
                av[index] += amount
        self.depth = bp + bv + ap + av
        self.callback['onDepth']({'name':self.name,'data':self.depth})

    def _update_trade(self, msg):
        msg = ujson.loads(msg)
        side = 'sell' if msg['m'] else 'buy'
        price = float(msg['p'])
        amount = float(msg['q'])
        if price > self.max_buy or self.max_buy == 0.0:
            self.max_buy = price
        if price < self.min_sell or self.min_sell == 0.0:
            self.min_sell = price
        if side == 'buy':
            self.buy_q += amount
            self.buy_v += amount*price
        elif side == 'sell':
            self.sell_q += amount
            self.sell_v += amount*price

    def _update_account(self, msg):
        msg = ujson.loads(msg)
        for i in msg['a']['B']:
            if i['a'].lower() == self.base.lower(): 
                self.callback['onEquity']({self.quote:float(i['wb'])})
    
    def _update_order(self, msg):
        msg = ujson.loads(msg)
        i = msg['o']
        if i['s'] == self.symbol:
            if i['X'] == 'NEW':  # 新增订单
                pass
            if i['X'] == 'FILLED':  # 删除订单
                self.callback['onOrder']({"deleteOrder":i['i']})
            if i['X'] == 'CANCELED':  # 删除订单
                self.callback['onOrder']({"deleteOrder":i['i']})

    def _update_position(self, msg):
        long_pos, short_pos = 0, 0
        long_avg, short_avg = 0, 0 
        msg = ujson.loads(msg)
        for i in msg['a']['P']:
            if i['s'] == self.symbol:
                if i['ps'] == 'LONG':
                    long_pos += abs(float(i['pa']))
                    long_avg = abs(float(i['ep']))
                if i['ps'] == 'SHORT':
                    short_pos += abs(float(i['pa']))
                    short_avg = abs(float(i['ep']))
        pos = model.Position()
        pos.longPos = long_pos
        pos.shortPos = short_pos
        pos.longAvg = long_avg
        pos.shortAvg = short_avg
        self.callback['onPosition'](pos)

    def _get_data(self):
        market_data = self.depth + [self.max_buy, self.min_sell]
        self.max_buy = 0
        self.min_sell = 0
        self.buy_v = 0.0
        self.buy_q = 0.0
        self.sell_v = 0.0
        self.sell_q = 0.0
        return {'name': self.name,'data':market_data}

    async def go(self):
        interval = float(self.params.interval)
        if self.is_print:print(f'Ws循环器启动 interval {interval}')
        ### onTrade
        while 1:
            try:
                # 更新市场信息
                market_data = self._get_data()
                self.callback['onMarket'](market_data)
            except:
                traceback.print_exc()
            await asyncio.sleep(interval)

    async def before_trade(self):
        session = aiohttp.ClientSession()
        response = await session.get(
            "https://dapi.binance.com/dapi/v1/exchangeInfo", 
            proxy=self.proxy
            )
        response = await response.json()
        for i in response['symbols']:
            if self.symbol in i['symbol'].upper():
                self.multiplier = float(i['contractSize'])
        

    async def run(self, is_auth=0, sub_trade=0, sub_fast=0):
        session = aiohttp.ClientSession()
        while True:
            try:
                # 尝试连接
                print(f'{self.name} 尝试连接ws')
                # 登陆
                if is_auth:
                    listenKey = await self.get_sign()
                else:
                    listenKey = 'qqlh'
                # 更新
                await self.before_trade()
                async with session.ws_connect(
                        self.URL+listenKey,
                        proxy=self.proxy,
                        timeout=30,
                        receive_timeout=30,
                        ) as _ws:
                    print(f'{self.name} ws连接成功')
                    self.logger.info(f'{self.name} ws连接成功')
                    # 订阅
                    symbol = self.symbol.lower()
                    if sub_fast:
                        channels=[f"{symbol}@bookTicker",]
                    else:
                        channels=[f"{symbol}@depth20@100ms",]
                    if sub_trade:
                        channels.append(f"{symbol}@aggTrade")                        
                    sub_str = ujson.dumps({"method": "SUBSCRIBE", "params": channels, "id":random.randint(1,1000)})
                    await _ws.send_str(sub_str)
                    while True:
                        # 停机信号
                        if self.stop_flag:return
                        # 接受消息
                        try:
                            msg = await _ws.receive(timeout=30)
                        except:
                            print(f'{self.name} ws长时间没有收到消息 准备重连...')
                            self.logger.error(f'{self.name} ws长时间没有收到消息 准备重连...')
                            break
                        msg = msg.data
                        # 处理消息
                        if 'depthUpdate' in msg:self._update_depth(msg)
                        elif 'aggTrade' in msg:self._update_trade(msg)
                        elif 'bookTicker' in msg:self._update_ticker(msg)
                        elif 'ACCOUNT_UPDATE' in msg:self._update_position(msg)
                        elif 'ACCOUNT_UPDATE' in msg:self._update_account(msg)
                        elif 'ORDER_TRADE_UPDATE' in msg:self._update_order(msg)
                        elif 'ping' in msg:await _ws.send_str('pong')
                        elif 'listenKeyExpired' in msg:raise Exception('key过期重连')
            except:
                _ws = None
                traceback.print_exc()
                print(f'{self.name} ws连接失败 开始重连...')
                self.logger.error(f'{self.name} ws连接失败 开始重连...')