import asyncio
import json
import logging
import os
import time
import traceback
from dataclasses import asdict
from datetime import datetime
from typing import AsyncContextManager
import modules.structs as structs
import pandas as pd
import requests
import valkey
from dotenv import load_dotenv
import modules.manual_leverage as leverage
import modules.aster_auth as aster_auth
import modules.utils as utils

### MANUAL LEVERAGE DATA ###
df_leverage_by_exch = pd.DataFrame(data=leverage.LEVERAGE_BY_EXCH)

### Database ###
# CON: AsyncContextManager | None = None
VAL_KEY: valkey.Valkey

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

### CONSTANTS ###
LOOP_SLEEP_SEC: int = 5
REFRESH_MKT_INFO_EVERY_SEC: int = 90
REFRESH_MKT_VOLUME_EVERY_SEC: int = 30

MINUTES_LOOKBACK: int = 60

### GLOBALS ###
Mkt_Info_Last_Refresh_TS_ms: int = 0 
Mkt_Volume_Last_Refresh_TS_ms: int = 0

### TODO: score by volume, how long since last trade?, volatility, volume by time of day (active or dormant period?), funding rate consistency (% one side last 24hrs and from active close to active open periods). trade cost estimate?, max tradeable notional.
### TODO: figure out what is max percent of volume i can trade - TCA kinda? what is ideal slice size?
### TODO: Redesign so Algo allocates across the best markets with a waterfall method until at target collateral usage. order waterfall by score above^^
### TODO: NG display grid of markets sorted by above score. top left is control panel, top right is graph (goes to mkt you click on from table) (maybe tabs for different graph views/groups, e.g. PnL total or all mkts percent to liquidate, pov by market etc.) middle bottom is markets table (tabs for open orders, open positions, pnl)

### Funcs - Load Data ###
async def get_extended_markets_info() -> pd.DataFrame:
    r: dict = json.loads(s=requests.get(url='https://api.starknet.extended.exchange/api/v1/info/markets').text)
    
    df: pd.DataFrame = pd.DataFrame(data=r['data'])
    df['funding_rate'] = df['marketStats'].apply(lambda x: x.get('fundingRate',{}))
    df['funding_rate_ts'] = df['marketStats'].apply(lambda x: x.get('nextFundingRate',{}))
    df['daily_volume'] = df['marketStats'].apply(lambda x: x.get('dailyVolume',{})).astype(float)
    df['min_order_size'] = df['tradingConfig'].apply(lambda x: x.get('minOrderSize',{}))
    df['min_price'] = df['tradingConfig'].apply(lambda x: x.get('minPriceChange',{}))
    df['min_notional'] = 0
    df['min_lot_size'] = df['tradingConfig'].apply(lambda x: x.get('minOrderSizeChange',{}))
    df['max_leverage'] = df['tradingConfig'].apply(lambda x: x.get('maxLeverage',{}))

    print('Extend markets info refreshed successfully')

    return df

async def get_aster_exch_info() -> pd.DataFrame:
    ### ASTER EXCHANGE INFO ###
    fut_acct_exchangeInfo: dict = {
        "url": "/fapi/v3/exchangeInfo",
        "method": "GET",
        "params": {}
    }
    r: dict = await aster_auth.post_authenticated_url(fut_acct_exchangeInfo)  # ty:ignore[invalid-assignment]
    df = pd.DataFrame(r['symbols'])
    df['min_order_size'] = df['filters'].apply(lambda x: [f for f in x if f.get('filterType', None) == 'LOT_SIZE'][0]['minQty'] )
    df['min_price'] = df['filters'].apply(lambda x: [f for f in x if f.get('filterType', None) == 'PRICE_FILTER'][0]['minPrice'] )
    df['min_notional'] = df['filters'].apply(lambda x: [f for f in x if f.get('filterType', None) == 'MIN_NOTIONAL'][0]['notional'] )
    df['min_lot_size'] = df['filters'].apply(lambda x: [f for f in x if f.get('filterType', None) == 'LOT_SIZE'][0]['stepSize'] )

    fut_acct_ticker_stats: dict = {
        "url": "/fapi/v3/ticker/24hr",
        "method": "GET",
        "params": {}
    }
    r: dict = await aster_auth.post_authenticated_url(fut_acct_ticker_stats)  # ty:ignore[invalid-assignment]
    df_stats = pd.DataFrame(r)
    df_stats['last_trade_ts_ast'] = df_stats['closeTime']

    df = df.merge(df_stats[['symbol','quoteVolume','last_trade_ts_ast']].rename({'quoteVolume':'daily_volume'}, axis=1), on='symbol', how='left')
    df['daily_volume'] = df['daily_volume'].astype(float)


    print('Aster markets info refreshed successfully')

    return df

def load_aster_current_fr(df_aster_exch_info: pd.DataFrame) -> pd.DataFrame:
    vk_get: str = VAL_KEY.get(name='fund_rate_aster_all')  # ty:ignore[invalid-assignment]
    if not vk_get:
        raise ValueError(f'fund_rate_aster_all is empty: {vk_get}')
    df = pd.DataFrame(data=json.loads(vk_get))
    
    df: pd.DataFrame = df[['s','E','r','T']].rename({'s':'symbol','E':'funding_rate_updated_ts_ms','r':'funding_rate','T':'next_funding_ts'}, axis=1)
    df['funding_rate_updated_dt'] = pd.to_datetime(df['funding_rate_updated_ts_ms'], unit='ms')
    df['funding_rate'] = df['funding_rate'].astype(float)
    df['time_delta_to_next_funding'] = pd.to_datetime(df['next_funding_ts'], unit='ms') - pd.Timestamp.now()
    df = df.merge(df_aster_exch_info[['symbol','daily_volume','min_order_size','min_price','min_lot_size','min_notional', 'last_trade_ts_ast']], on='symbol', how='left')
    
    return df

def load_extend_current_fr(df_mkt_stats: pd.DataFrame) -> pd.DataFrame:
    df = pd.DataFrame(data=json.loads(s=VAL_KEY.get(name='fund_rate_extended_all')))  # ty:ignore[invalid-argument-type]
    
    df: pd.DataFrame = df[['symbol','funding_rate_updated_ts_ms','funding_rate']]
    df['funding_rate_updated_dt'] = pd.to_datetime(df['funding_rate_updated_ts_ms'], unit='ms')
    df['funding_rate'] = df['funding_rate'].astype(float)

    df = df.merge(df_mkt_stats[['name','assetName','status','funding_rate_ts','min_order_size','min_price','min_lot_size','min_notional','daily_volume']].rename({'name':'symbol','funding_rate_ts':'next_funding_ts'}, axis=1), on='symbol', how='left')
    df: pd.DataFrame = df.loc[df['status']=='ACTIVE',:]
    df['USDT_Symbol'] = df['assetName'] + 'USDT'
    
    df['time_delta_to_next_funding'] = pd.to_datetime(arg=df['next_funding_ts'], unit='ms') - pd.Timestamp.now()

    return df

async def get_candles(symbol: str, limit: int = MINUTES_LOOKBACK) -> pd.DataFrame:
    ### Candles for Midpoint Dispersion ###
    # Aster
    symbol_ast = utils.symbol_to_aster_fmt(symbol)
    aster_candles = {
        "url": "/fapi/v3/klines",
        "method": "GET",
        "params": {
            'symbol': symbol_ast,
            'interval': '1m',
            'limit':str(limit)
        }
    }
    j = await aster_auth.post_authenticated_url(aster_candles)
    df_candles_aster = pd.DataFrame(j, columns=['open_ts','open_px','high_px','low_px','close_px','volume','close_ts','quote_asset_volume','count_trades','taker_buy_base_asset_volume','taker_buy_quote_asset_volume','_drop'])
    df_candles_aster = df_candles_aster[['open_px', 'low_px', 'high_px', 'close_px', 'volume', 'open_ts']]
    df_candles_aster[['open_px', 'low_px', 'high_px', 'close_px', 'volume']] = df_candles_aster[['open_px', 'low_px', 'high_px', 'close_px', 'volume']].astype(float)

    df_candles_aster['med_px'] = ( df_candles_aster['high_px'] + df_candles_aster['low_px'] ) / 2
    df_candles_aster['typical_px'] = ( df_candles_aster['open_px'] + df_candles_aster['high_px'] + df_candles_aster['low_px'] + df_candles_aster['close_px'] ) / 4

    # Extend
    symbol_ext = utils.symbol_to_extend_fmt(symbol)
    ext_params = {
        'interval':'1m',
        'limit': limit,
    }
    r = json.loads(requests.get(f'https://api.starknet.extended.exchange/api/v1/info/candles/{symbol_ext}/trades', params=ext_params).text)
    df_candles_extended = pd.DataFrame(r['data'])
    df_candles_extended = df_candles_extended.rename({'o':'open_px','l':'low_px','h':'high_px','c':'close_px','v':'volume','T':'open_ts'}, axis=1)
    df_candles_extended[['open_px', 'low_px', 'high_px', 'close_px', 'volume']] = df_candles_extended[['open_px', 'low_px', 'high_px', 'close_px', 'volume']].astype(float)
    df_candles_extended['med_px'] = ( df_candles_extended['high_px'] + df_candles_extended['low_px'] ) / 2
    df_candles_extended['typical_px'] = ( df_candles_extended['open_px'] + df_candles_extended['high_px'] + df_candles_extended['low_px'] + df_candles_extended['close_px'] ) / 4

    df_candles_comb = df_candles_aster.merge(df_candles_extended, on='open_ts', how='inner', suffixes=('_ast','_ext'))
    df_candles_comb['open_dt'] = pd.to_datetime(df_candles_comb['open_ts'], unit='ms')
    df_candles_comb['med_ratio_aster_over_extend'] = ( df_candles_comb['med_px_ast'] / df_candles_comb['med_px_ext'] ) - 1

    return df_candles_comb


async def loop() -> None:
    global Mkt_Info_Last_Refresh_TS_ms
    try:
        while True:
            ts_arrival = round(datetime.now().timestamp() * 1000)
            if ( ts_arrival - Mkt_Info_Last_Refresh_TS_ms ) > ( REFRESH_MKT_INFO_EVERY_SEC * 1000 ):
                df_extend_mkt_stats = await get_extended_markets_info()
                df_aster_exch_info = await get_aster_exch_info()
                Mkt_Info_Last_Refresh_TS_ms = round(datetime.now().timestamp() * 1000)

            df_aster_fr = load_aster_current_fr(df_aster_exch_info=df_aster_exch_info)
            df_extend_fr = load_extend_current_fr(df_mkt_stats=df_extend_mkt_stats)
            
            df_comb_fr = df_extend_fr.merge(df_aster_fr, left_on='USDT_Symbol', right_on='symbol', how='inner', suffixes=('_ext', '_ast'))
            df_comb_fr['next_funding_at_same_time'] = (abs(df_comb_fr['time_delta_to_next_funding_ext'].dt.total_seconds() - df_comb_fr['time_delta_to_next_funding_ast'].dt.total_seconds()) / 60) < 1
            df_comb_fr['net_funding_rate'] = (df_comb_fr[['funding_rate_ext', 'funding_rate_ast']].max(axis=1) - df_comb_fr[['funding_rate_ext', 'funding_rate_ast']].min(axis=1)).where(df_comb_fr['next_funding_at_same_time'], df_comb_fr['funding_rate_ext'])
            df_comb_fr['net_funding_rate_abs'] = df_comb_fr['net_funding_rate'].abs()
            
            ### NET MULT ###
            df_comb_fr = df_comb_fr.merge(right=df_leverage_by_exch.loc[df_leverage_by_exch['exchange']=='EXTEND'], left_on='assetName', right_on='lh_asset').merge(df_leverage_by_exch.loc[df_leverage_by_exch['exchange']=='ASTER'], left_on='assetName', right_on='lh_asset', suffixes=('_ext', '_ast'))
            df_comb_fr['net_mult'] = 1 / ( ( 0.5 / df_comb_fr['max_leverage_ext'] ) + ( 0.5 / df_comb_fr['max_leverage_ast'] ) )
            df_comb_fr['net_mult'] = df_comb_fr['net_mult'].round(2)
            df_comb_fr['net_mult_x_net_fr_abs'] = df_comb_fr['net_funding_rate_abs'] * df_comb_fr['net_mult']

            df_best_fr_rate = df_comb_fr[
                ['symbol_ext','symbol_ast','daily_volume_ext','daily_volume_ast','min_price_ext','min_price_ast','min_order_size_ext',
                'min_order_size_ast','min_lot_size_ext','min_lot_size_ast','min_notional_ext','min_notional_ast','funding_rate_ext',
                'funding_rate_ast','max_leverage_ext','max_leverage_ast','lh_asset_ext','lh_asset_ast','rh_asset_ext','rh_asset_ast',
                'net_mult_x_net_fr_abs','net_funding_rate_abs','net_funding_rate','next_funding_at_same_time','last_trade_ts_ast']
                ].sort_values(by='net_mult_x_net_fr_abs', ascending=False).reset_index(drop=True)

            last_trade_max_ts = []
            for index, row in df_best_fr_rate.iterrows():
                r = json.loads(requests.get(f'https://api.starknet.extended.exchange/api/v1/info/markets/{row['symbol_ext']}/trades').text)
                max_ts = max([t['T'] for t in r['data']])
                last_trade_max_ts.append({'symbol_ext':row['symbol_ext'],'last_trade_ts_ext': max_ts})
                time.sleep(0.01)

            df_best_fr_rate = df_best_fr_rate.merge(pd.DataFrame(last_trade_max_ts), on='symbol_ext', how='left')
            
            df_best_fr_rate['last_trade_ts_dt_ast'] = pd.to_datetime(df_best_fr_rate['last_trade_ts_ast'], unit='ms')
            df_best_fr_rate['last_trade_ts_dt_ext'] = pd.to_datetime(df_best_fr_rate['last_trade_ts_ext'], unit='ms')
            

            candles_ratios = []

            for index, row in df_best_fr_rate.iterrows():
                try:
                    df = await get_candles(symbol=row['symbol_ext'])
                except Exception as e:
                    logging.warning(f'BFR failed to get candles...sleeping and retrying: {e}')
                    time.sleep(5)
                    df = await get_candles(symbol=row['symbol_ext'])

                buy_ratio_ext = float(df['med_ratio_aster_over_extend'].median())
                buy_ratio_std = float(df['med_ratio_aster_over_extend'].std())
                candles_ratios.append({'symbol_ext':row['symbol_ext'], 'buy_ratio_std': buy_ratio_std, 'buy_ratio_ext':buy_ratio_ext,'buy_ratio_ast':buy_ratio_ext*-1})

            df_best_fr_rate = df_best_fr_rate.merge(pd.DataFrame(candles_ratios), on='symbol_ext', how='left')
            
            ### Set Unfiltered Master Data ###
            master_data = df_best_fr_rate[
                ['symbol_ast','max_leverage_ast','lh_asset_ast','rh_asset_ast','funding_rate_ast','min_price_ast','min_order_size_ast','min_lot_size_ast','min_notional_ast','buy_ratio_ast',
                 'symbol_ext','max_leverage_ext','lh_asset_ext','rh_asset_ext','funding_rate_ext','min_price_ext','min_order_size_ext','min_lot_size_ext','min_notional_ext','buy_ratio_ext', 'buy_ratio_std','next_funding_at_same_time']
                ].to_json(orient='records')
            VAL_KEY.set(name='fr_engine_best_fund_rate_master', value=str(master_data))
            
            
            ### Filter BFR Data ###
            df_best_fr_rate = df_best_fr_rate.loc[( (datetime.now().timestamp()*1000 )-df_best_fr_rate['last_trade_ts_ast']) < (5*60*1000) ] # Last traded in 3min
            df_best_fr_rate = df_best_fr_rate.loc[( (datetime.now().timestamp()*1000 )-df_best_fr_rate['last_trade_ts_ext']) < (15*60*1000) ] # Last traded in 15min

            min_daily_volume = 100_000
            df_best_fr_rate = df_best_fr_rate.loc[ (df_best_fr_rate['daily_volume_ast']>=min_daily_volume) & (df_best_fr_rate['daily_volume_ext']>min_daily_volume) ,:].reset_index(drop=True)
            
            # print(df_best_fr_rate.columns)
            # print(df_best_fr_rate.iloc[0])

            if len(df_best_fr_rate) < 1:
                raise ValueError(f'NO BFR RATE: {df_best_fr_rate}')

            try:
                ASTER = structs.Perpetual_Exchange(
                    mult = int(df_best_fr_rate['max_leverage_ast'].iloc[0]),
                    lh_asset = df_best_fr_rate['lh_asset_ast'].iloc[0],
                    rh_asset = df_best_fr_rate['rh_asset_ast'].iloc[0],
                    symbol_asset_separator = '',
                    initial_funding_rate=float(df_best_fr_rate['funding_rate_ast'].iloc[0]),
                    fund_rate_at_same_time=bool(df_best_fr_rate['next_funding_at_same_time'].iloc[0]),
                    min_price=float(df_best_fr_rate['min_price_ast'].iloc[0]),
                    min_order_size=float(df_best_fr_rate['min_order_size_ast'].iloc[0]),
                    min_lot_size=float(df_best_fr_rate['min_lot_size_ast'].iloc[0]),
                    min_notional=float(df_best_fr_rate['min_notional_ast'].iloc[0]),
                    buy_ratio=float(df_best_fr_rate['buy_ratio_ast'].iloc[0]),
                    buy_ratio_std=float(df_best_fr_rate['buy_ratio_std'].iloc[0]),
                )          
                EXTEND = structs.Perpetual_Exchange(
                    mult = int(df_best_fr_rate['max_leverage_ext'].iloc[0]),
                    lh_asset = df_best_fr_rate['lh_asset_ext'].iloc[0],
                    rh_asset = df_best_fr_rate['rh_asset_ext'].iloc[0],
                    symbol_asset_separator = '-',
                    initial_funding_rate=float(df_best_fr_rate['funding_rate_ext'].iloc[0]),
                    fund_rate_at_same_time=bool(df_best_fr_rate['next_funding_at_same_time'].iloc[0]),
                    min_price=float(df_best_fr_rate['min_price_ext'].iloc[0]),
                    min_order_size=float(df_best_fr_rate['min_order_size_ext'].iloc[0]),
                    min_lot_size=float(df_best_fr_rate['min_lot_size_ext'].iloc[0]),
                    min_notional=float(df_best_fr_rate['min_notional_ext'].iloc[0]),
                    buy_ratio=float(df_best_fr_rate['buy_ratio_ext'].iloc[0]),
                    buy_ratio_std=float(df_best_fr_rate['buy_ratio_std'].iloc[0]),
                )
            except Exception as e:
                logging.critical(f'Failed to build ASTER/EXTEND objs err: {e}; df cols: {df_best_fr_rate.columns}')
                logging.error(traceback.format_exc())
                continue
            best_next_funding_pair: dict[str, dict] = {'ASTER': asdict(obj=ASTER), 'EXTEND': asdict(obj=EXTEND)}
            VAL_KEY.set(name='fr_engine_best_fund_rate_output', value=json.dumps(obj=best_next_funding_pair))
            

            print(df_best_fr_rate[['symbol_ext','max_leverage_ext','buy_ratio_ext','net_funding_rate','daily_volume_ast','buy_ratio_ast']].head(10))
            logging.info(f'BFR REFRESHED @ {datetime.now()}')
            time.sleep(LOOP_SLEEP_SEC)
            continue
    except valkey.exceptions.ConnectionError as e:
        logging.info(f"Could not connect to Valkey. Please check the publish server is up; {e}")
    except KeyboardInterrupt:
        logging.info('SHUTTING DOWN...')
    except Exception as e:
        logging.error(traceback.format_exc())
        logging.critical(f'*** CRASHED: {e}')


### STARTUP ###
async def main() -> None:
    global VAL_KEY
    # global CON

    VAL_KEY = valkey.Valkey(host='localhost', port=6379, db=0, decode_responses=True)
    # engine = create_async_engine('mysql+asyncmy://root:pwd@localhost/fund_rate')

    await loop()

if __name__ == '__main__':
    START_TIME = 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}")

    asyncio.run(main())