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
import modules.utils as utils

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

### Database ###
USE_DB: bool = False
VK_CHANNEL = 'fund_rate_aster_all'

CON: AsyncContextManager
VAL_KEY: valkey.Valkey

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

### Globals ###
WSS_URL: str = "wss://fstream.asterdex.com/ws/!markPrice@arr"


### Websocket ###
async def ws_stream():
    async for websocket in websockets.connect(WSS_URL):
        logging.info(msg=f"Connected to {WSS_URL}")
        try:
            async for message in websocket:
                if isinstance(message, str):
                    try: 
                        data: dict = json.loads(message)
                        if data:
                            VAL_KEY.set(name=VK_CHANNEL, value=json.dumps(obj=data))
                            # print(f'VK_SAVED: {len(data)}')
                            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(msg=f'Connection closed: {e}')
            logging.error(msg=traceback.format_exc())
            utils.send_tg_alert(msg=f'WS: {VK_CHANNEL} - Failure: {e}')
        except Exception as e:
            logging.error(msg=f'Connection closed: {e}')
            logging.error(msg=traceback.format_exc())
            utils.send_tg_alert(msg=f'WS: {VK_CHANNEL} - Failure: {e}')

### Startup ###
async def main():
    global VAL_KEY
    global CON
    
    VAL_KEY = valkey.Valkey(host='localhost', port=6379, db=0)
    
    if USE_DB:
        raise NotImplementedError('DB not implemented for this ws.')
        # engine = create_async_engine('mysql+asyncmy://root:pwd@localhost/fund_rate')
        # async with engine.connect() as CON:
            # await ws_stream()
    else:
        logging.warning("DATABASE NOT BEING USED, NO DATA WILL BE RECORDED")
        await ws_stream()


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

    logging.info(msg=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(msg=f"STARTED: {START_TIME}")
    
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logging.info(msg="Stream stopped")