import logging as Log
import grpc,argparse,time,sys,threading

import mllp_grpc.asr_pb2_grpc as asr_pb2_grpc
import mllp_grpc.asr_pb2 as asr_pb2
import mllp_grpc.asr_common_pb2 as asr_common_pb2
import google.protobuf.empty_pb2 as empty_pb2

from pydub import AudioSegment
from math import ceil

def get_data(args,sys_id):
    Log.info("Loading '%s'"%args.wav)
    sound= AudioSegment.from_wav(args.wav)
    
    if (sound.frame_rate != 16000):
        sound= sound.set_frame_rate(16000)
        
    nbuffers= ceil(len(sound)/args.bsize)
    Log.info(('c: %d r: %d rate: %d len: %d nbuf:'+
              ' %d bsize: %d')%(sound.channels,
                                sound.sample_width,
                                sound.frame_rate,
                                len(sound),
                                nbuffers,
                                args.bsize))
    sended= 0
    Log.info('Sending data')
    yield asr_pb2.DecodeRequest(system_id=sys_id)
    for n in range(0,nbuffers):
        buf= sound[n*args.bsize:(n+1)*args.bsize]
        yield asr_pb2.DecodeRequest(data=buf.raw_data)
        sended+= args.bsize
        if sended>=1000:
            sended-= 1000
            time.sleep(1)
    Log.info('Transmission completed')
        
def parse_cmdline():
    p= argparse.ArgumentParser(description='ASR GRPC client')
    p.add_argument('host',type=str,help='server host')
    p.add_argument('port',type=int,help='server port')
    p.add_argument('wav',type=str,help='input 1 channel wav at 16KHz')
    p.add_argument('--bsize',type=int,default=250,
                   help='buffer size')
    args= p.parse_args()
    return args

def run(args):
    addr= '%s:%d'%(args.host,args.port)
    with grpc.insecure_channel(addr) as channel:
        stub= asr_pb2_grpc.ASRStub(channel)
        for info in stub.GetSystemsInfo(empty_pb2.Empty()):
            print(info)
        print_host= True
        for o in stub.Decode(get_data(args,1)):
            if print_host:
                print_host= False
                print(o.host_info)
            if o.status.code!=asr_common_pb2.DecodeResponse.Status.Code.OK:
                print(o)
            elif o.hyp_novar!='':
                sys.stdout.write(o.hyp_novar+' ')
            sys.stdout.flush()
        sys.stdout.write('\n')
        sys.stdout.flush()
        
if __name__=='__main__':
    Log.basicConfig(level=Log.INFO)
    args= parse_cmdline()
    run(args)
