import asyncio
import nest_asyncio
import json
import logging

# TODO обавить обработку ошибок и try

from okx.websocket.WsPublicAsync import WsPublicAsync

nest_asyncio.apply()

logger = logging.getLogger()
logger.setLevel(logging.CRITICAL)

file1 = open('okx_trades.logs', 'w')
file2 = open('okx_depth.logs', 'w')

def publicCallback(msg): # snapshot and update data
    msg = json.loads(msg)
    ev = msg.get('event')
    data = msg.get('data')

    if msg['arg']['channel'] == 'trades':
        if data:
            data = data[0]
            renamed_fields = {'ts': 'ts', 'price': 'px', 'size': 'sz', 'side': 'side'}
            file1.write(json.dumps({field: data[renamed_fields[field]] for field in renamed_fields}))
            file1.write('\n')
            file1.flush()

    elif msg['arg']['channel'] == 'books':
        if data:
            data = data[0]
            file2.write('asks len: ' + str(len(data['asks'])) + '\n')
            file2.write('bids len: ' + str(len(data['bids'])) + '\n')

            depth_lvl = {
                'asks':    data['asks'][:min(len(data['asks']), 50)],
                'bids':    data['bids'][:min(len(data['bids']), 50)],
                'ts':      data['ts'],
                'action':  msg['action']
            }

            file2.write(json.dumps(depth_lvl))
            file2.write('\n')
            file2.flush()


async def main():
    url = "wss://wspap.okx.com:8443/ws/v5/public?brokerId=9999"
    ws = WsPublicAsync(url=url)
    await ws.start()

    args = []
    arg1 = {"channel": "trades", "instId": "BTC-USDT"}
    args.append(arg1)
    arg2 = {"channel": "books", "instId": "BTC-USDT"}
    args.append(arg2)

    await ws.subscribe(args, publicCallback)

    while True:
        await asyncio.sleep(20)


if __name__ == "__main__":
    asyncio.run(main())