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


### 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 | None = None
VAL_KEY = None

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

### CONSTANTS ###
WS_SYMBOL: str = 'ETH-USD'
FUNDING_RATE_INTERVAL_MIN = 60

### Globals ###
WSS_URL = f"wss://api.starknet.extended.exchange/stream.extended.exchange/v1/funding/{WS_SYMBOL}"


# HIST_TRADES = np.empty((0, 3))
# HIST_TRADES_LOOKBACK_SEC = 6

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


# ### Database Funcs ###
# async def create_rtds_btcusd_table(
#         CON: AsyncContextManager,
#         engine: str = 'mysql', # mysql | duckdb
#     ) -> None:
#     if CON is None:
#         logging.info("NO DB CONNECTION, SKIPPING Create Statements")
#     else:
#         if engine == 'mysql':
#             logging.info('Creating Table if Does Not Exist: binance_btcusd_trades')
#             await CON.execute(text("""
#             CREATE TABLE IF NOT EXISTS binance_btcusd_trades (
#                 timestamp_arrival BIGINT,
#                 timestamp_msg BIGINT,
#                 timestamp_value BIGINT,
#                 value DOUBLE,
#                 qty DOUBLE
#             );
#             """))
#             await CON.commit()
#         else:
#             raise ValueError('Only MySQL engine is implemented')

# async def insert_rtds_btcusd_table(
#         timestamp_arrival: int,
#         timestamp_msg: int,
#         timestamp_value: int,
#         value: float,
#         qty: float,
#         CON: AsyncContextManager,
#         engine: str = 'mysql', # mysql | duckdb
#     ) -> None:
#     params={
#         'timestamp_arrival': timestamp_arrival,
#         'timestamp_msg': timestamp_msg,
#         'timestamp_value': timestamp_value,
#         'value': value,
#         'qty': qty,
#     }
#     if CON is None:
#         logging.info("NO DB CONNECTION, SKIPPING Insert Statements")
#     else:
#         if engine == 'mysql':
#             await CON.execute(text("""
#             INSERT INTO binance_btcusd_trades 
#             (
#                 timestamp_arrival,
#                 timestamp_msg,
#                 timestamp_value,
#                 value,
#                 qty
#             )
#             VALUES 
#                 (
#                 :timestamp_arrival,
#                 :timestamp_msg,
#                 :timestamp_value,
#                 :value,
#                 :qty
#                 )
#             """),
#             parameters=params
#             )
#             await CON.commit()
#         else:
#             raise ValueError('Only MySQL engine is implemented')

### Websocket ###
async def ws_stream():
    async for websocket in websockets.connect(WSS_URL):
        logging.info(f"Connected to {WSS_URL}")
        try:
            async for message in websocket:
                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
                            VAL_KEY_OBJ = json.dumps({
                                '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,
                            })
                            VAL_KEY.set(VK_FUND_RATE, 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:
        VAL_KEY = None
        logging.warning("VALKEY NOT BEING USED, NO DATA WILL BE PUBLISHED")

    if USE_DB:
        engine = create_async_engine('mysql+asyncmy://root:pwd@localhost/polymarket')
        async with engine.connect() as CON:
            # await create_rtds_btcusd_table(CON=CON)
            await ws_stream()
    else:
        CON = None
        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")