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

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

### Database ###
# CON: AsyncContextManager | None = None
VAL_KEY: valkey.Valkey
VK_OUT: str = 'fr_engine_best_fund_rate_output'

### 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

### GLOBALS ###
Mkt_Info_Last_Refresh_TS_ms: int
Mkt_Volume_Last_Refresh_TS_ms: int

### Funcs - Load Data ###
def get_extended_markets_info() -> pd.DataFrame:
    global Mkt_Info_Last_Refresh_TS_ms

    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['min_order_size'] = df['tradingConfig'].apply(lambda x: x.get('minOrderSize',{}))
    df['min_price_change'] = df['tradingConfig'].apply(lambda x: x.get('minPriceChange',{}))
    df['max_leverage'] = df['tradingConfig'].apply(lambda x: x.get('maxLeverage',{}))
    
    Mkt_Info_Last_Refresh_TS_ms = round(datetime.now().timestamp() * 1000)

    print('Extend markets info refreshed successfully')

    return df

def load_aster_current_fr() -> pd.DataFrame:
    df = pd.DataFrame(data=json.loads(s=VAL_KEY.get(name='fund_rate_aster_all')))  # ty:ignore[invalid-argument-type]
    
    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()

    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: pd.DataFrame = df.merge(df_mkt_stats[['name','assetName','status', 'funding_rate_ts']].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 loop() -> None:
    global Mkt_Info_Last_Refresh_TS_ms
    
    df_extend_mkt_stats = get_extended_markets_info()
    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 = get_extended_markets_info()

            df_aster_fr = load_aster_current_fr()
            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: pd.DataFrame = df_comb_fr[['symbol_ext','symbol_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']].sort_values(by='net_mult_x_net_fr_abs', ascending=False).reset_index(drop=True)
            
            ASTER = structs.Perpetual_Exchange(
                mult = int(df_best_fr_rate['max_leverage_ast'][0]),
                lh_asset = df_best_fr_rate['lh_asset_ast'][0],
                rh_asset = df_best_fr_rate['rh_asset_ast'][0],
                symbol_asset_separator = '',
            )
            
            EXTEND = structs.Perpetual_Exchange(
                mult = int(df_best_fr_rate['max_leverage_ext'][0]),
                lh_asset = df_best_fr_rate['lh_asset_ext'][0],
                rh_asset = df_best_fr_rate['rh_asset_ext'][0],
                symbol_asset_separator = '-',
            )
            
            best_next_funding_pair: dict[str, dict] = {'ASTER': asdict(obj=ASTER), 'EXTEND': asdict(obj=EXTEND)}

            VAL_KEY.set(name=VK_OUT, value=json.dumps(obj=best_next_funding_pair))
            # print(best_next_funding_pair)
            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())