import asyncio
import json
import logging
import socket
import traceback
from datetime import datetime, timezone
from typing import AsyncContextManager
import math

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 sys
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
USE_VK: bool = True
VK_FUND_RATE = 'fund_rate_extended'

CON: AsyncContextManager
VAL_KEY: valkey.Valkey

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

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

### Globals ###
ALLOW_SYMBOL_CHG: bool = True
LOCAL_FUNDING_RATES = []

def time_round_down(dt, interval_mins=5) -> int: # returns timestamp in seconds
    interval_secs = interval_mins * 60
    seconds = dt.timestamp()
    rounded_seconds = math.floor(seconds / interval_secs) * interval_secs
    
    return rounded_seconds

### Websocket ###
async def ws_stream():
    global SYMBOL
    global LOCAL_FUNDING_RATES
    
    while True:
        # CHANGE_SYMBOL = False
        WSS_URL = "wss://api.starknet.extended.exchange/stream.extended.exchange/v1/funding/"
        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:
                        fr_algo_working_symbol = VAL_KEY.get(name='fr_algo_working_symbol')
                        if not fr_algo_working_symbol:
                            logging.critical(f'fr_algo_working_symbol is empty - killing: {fr_algo_working_symbol}')
                            sys.exit(1)
                        best_symbol_by_exchange: dict = json.loads(fr_algo_working_symbol)  # ty:ignore[invalid-argument-type]
                        best_symbol: str = best_symbol_by_exchange['EXTEND']['symbol']
                        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('data', None) is not None:
                                print(f'FR: {data}')
                                fr_next_update_ts = (time_round_down(dt=datetime.now(timezone.utc), interval_mins=60)+(60*60))*1000
                                fr_update = {
                                    'sequence_id': data['seq'],
                                    'timestamp_arrival': ts_arrival,
                                    'timestamp_msg': data['ts'],
                                    'symbol': data['data']['m'],
                                    'funding_rate': float(data['data']['f']),
                                    'funding_rate_updated_ts_ms': data['data']['T'],
                                    'next_funding_time_ts_ms': fr_next_update_ts,
                                }
                                if fr_update.get('symbol') == SYMBOL:
                                    VAL_KEY_OBJ = json.dumps(fr_update)
                                    VAL_KEY.set(VK_FUND_RATE, VAL_KEY_OBJ)
                                LOCAL_FUNDING_RATES = utils.upsert_list_of_dicts_by_id(LOCAL_FUNDING_RATES, fr_update, id='symbol', seq_check_field=None)
                                VAL_KEY.set('fund_rate_extended_all', json.dumps(LOCAL_FUNDING_RATES))
                                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")