import asyncio
import json
import logging
import socket
import traceback
from datetime import datetime
from typing import AsyncContextManager

import numpy as np
import pandas as pd
import requests.packages.urllib3.util.connection as urllib3_cn  # type: ignore
from sqlalchemy import text
import websockets
from sqlalchemy.ext.asyncio import create_async_engine
import valkey
import os
from dotenv import load_dotenv


### Allow only ipv4 ###
def allowed_gai_family():
    return socket.AF_INET
urllib3_cn.allowed_gai_family = allowed_gai_family

### Database ###
USE_DB: bool = False
USE_VK: bool = True

VK_TICKER = 'fut_ticker_extended'
CON: AsyncContextManager
VAL_KEY: valkey.Valkey

### Logging ###
load_dotenv()
LOG_FILEPATH: str = f'{os.getenv("LOGS_PATH")}/Fund_Rate_Extended_OB.log'

### CONSTANTS ###
SYMBOL: str = 'ETH-USD'

### Globals ###
ALLOW_SYMBOL_CHG: bool = False

### Websocket ###
async def ws_stream():
    global SYMBOL
    
    while True:
        CHANGE_SYMBOL = False
        WSS_URL = f"wss://api.starknet.extended.exchange/stream.extended.exchange/v1/orderbooks/{SYMBOL}?depth=1"
        async for websocket in websockets.connect(WSS_URL):
            if CHANGE_SYMBOL:
                break
            logging.info(f"Connected to {WSS_URL}")
            try:
                async for message in websocket:
                    ### Update Symbol if Algo Outputs Change ###
                    if ALLOW_SYMBOL_CHG:
                        best_symbol_by_exchange: dict = json.loads(s=VAL_KEY.get(name='fr_algo_working_symbol'))  # ty:ignore[invalid-argument-type]
                        best_symbol: str = f'{best_symbol_by_exchange['EXTEND']['lh_asset']}-{best_symbol_by_exchange['EXTEND']['rh_asset']}'
                        if best_symbol != SYMBOL:
                            logging.info(f'Symbol Change: {SYMBOL} -> {best_symbol}')
                            SYMBOL = best_symbol
                            CHANGE_SYMBOL = True
                            await websocket.close()
                            break
                    
                    ts_arrival = round(datetime.now().timestamp()*1000)
                    if isinstance(message, str):
                        try: 
                            data = json.loads(message)
                            if data.get('type', None) is not None:
                                # print(f'OB: {data}')
                                VAL_KEY_OBJ = json.dumps({
                                    'sequence_id': data['seq'],
                                    'timestamp_arrival': ts_arrival,
                                    'timestamp_msg': data['ts'],
                                    'symbol': data['data']['m'],
                                    'best_bid_px': float(data['data']['b'][0]['p']),
                                    'best_bid_qty': float(data['data']['b'][0]['q']),
                                    'best_ask_px': float(data['data']['a'][0]['p']),
                                    'best_ask_qty': float(data['data']['a'][0]['q']),
                                })
                                VAL_KEY.set(VK_TICKER, VAL_KEY_OBJ)
                                continue
                            else:
                                logging.info(f'Initial or unexpected data struct, skipping: {data}')
                                continue
                        except (json.JSONDecodeError, ValueError):
                            logging.warning(f'Message not in JSON format, skipping: {message}')
                            continue
                    else:
                        raise ValueError(f'Type: {type(data)}  not expected: {message}')
            except websockets.ConnectionClosed as e:
                logging.error(f'Connection closed: {e}')
                logging.error(traceback.format_exc())
                continue
            except Exception as e:
                logging.error(f'Connection closed: {e}')
                logging.error(traceback.format_exc())


async def main():
    global VAL_KEY
    global CON
    
    if USE_VK:
        VAL_KEY = valkey.Valkey(host='localhost', port=6379, db=0)
    else:
        logging.warning("VALKEY NOT BEING USED, NO DATA WILL BE PUBLISHED")
        raise NotImplementedError('Cannot run without VK')

    if USE_DB:
        raise NotImplementedError('DB not implemented')
        # engine = create_async_engine('mysql+asyncmy://root:pwd@localhost/fund_rate')
        # async with engine.connect() as CON:
        #     # await create_rtds_btcusd_table(CON=CON)
        #     await ws_stream()
    else:
        logging.warning("DATABASE NOT BEING USED, NO DATA WILL BE RECORDED")
        await ws_stream()


if __name__ == '__main__':
    START_TIME = round(datetime.now().timestamp()*1000)

    logging.info(f'Log FilePath: {LOG_FILEPATH}')

    logging.basicConfig(
        force=True,
        filename=LOG_FILEPATH,
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        filemode='w'
    )
    logging.info(f"STARTED: {START_TIME}")
    
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logging.info("Stream stopped")