import yaml
import pickle
import pprint
from neo_api_client import NeoAPI
import sys

#sample prog for live streaming of bank nifty index with reuse of session.
# 1. install above packages using pip install 
# 2. store credentials in credKotak.yml
# 3. For first time during day you would receive TOTP prompt, enter totp 
# 4. Now you can do ctrl C and start again , streaming will start without asking any totp.

KOTAKHOME='./'
FILEPICKLE = KOTAKHOME + 'picklekotak.pkl'

def storePickle( obj ):
    #create a pickle file
    picklefile = open( FILEPICKLE , 'wb')
    #pickle and write it to file
    pickle.dump(obj.configuration , picklefile)
    #close the file
    picklefile.close()
    print ( "Wrote pickle file" )

def readPickle():
    #read the pickle file
    obj = None
    try:
        picklefile = open( FILEPICKLE , 'rb')
        #unpickle the dataframe
        obj = pickle.load( picklefile )
        picklefile.close()
        return  obj
    except:
        print ( "Error reading pickle file" )
        return None

def cb_message( x ):
    print( "On Msg" )
    pprint( x )

def cb_error( x ):
    print( "Error" )
    pprint( x )

def cb_open( x ):
    print( "on open" )
    pprint( x )

def cb_close( x ):
    print( "on close" )
    pprint( x )


def getApi():
    clientconf = readPickle()
    with open( KOTAKHOME  + 'credKotak.txt') as f:
        cred = yaml.load(f, Loader=yaml.FullLoader)
        print("Loading KOTAK credentials...")
    if clientconf:
        client = NeoAPI(consumer_key= cred['consumerKey'], consumer_secret=cred['consumerSecret'], environment='prod', on_message=cb_message, on_error=cb_error, on_close=cb_close, on_open= cb_open)
        client.configuration = clientconf
        print ("Returning a set session " )
        #print_all( client )
        return client

    else:
        try:
            client = NeoAPI(consumer_key= cred['consumerKey'], consumer_secret=cred['consumerSecret'], environment='prod', on_message=cb_message, on_error=cb_error, on_close=cb_close, on_open= cb_open)
            client.login(mobilenumber=cred['mobile'], password=cred['password'])
            totp = input("Enter TOTP: ")
            totp =  totp 
            client.session_2fa(OTP=totp)
            cred['factor2'] = totp
            storePickle( client )
            return client
        except Exception as e:
            print( "Error while auth" , str(e) )
            os._exit( 0 )
            return None


if '__main__' == __name__:

   #make tokens for live subscription, multiple can be added to list
   banknifty= 26009 
   tokens=[]
   segment = 'nse_fo'
   token= {"instrument_token": str( banknifty ), "exchange_segment": segment }
   tokens.append( token )

   client = getApi()
   client.subscribe( tokens )