# MapReduce for python/repy!
#
#
# TODO list:
#
#

include sockettimeout.repy

# !! currently only works for a single mapper !!
# i'll worry about splitting the data file up at a later time
def initalize_replicas(data):
    
    # generate strings:
    # generate primary location string
    prim_str = format_message(mycontext['my_addr'])

    # generate replica list to send
    replica_str = pad_int_to_str(len(mycontext['scoreboard'].keys()), 4)
    replica_str += mycontext['replica_str']

#    print "replica_str:", replica_str

    replicas = mycontext['scoreboard'].keys()
    replicas.sort()

    # split up the data
    data_lines = data.split("\n")
    chunk_size = len(data_lines) / len(replicas)
    if len(data_lines) % len(replicas) != 0:
        chunk_size + 1

    data_dist = []
    for i in xrange(0, len(data_lines), chunk_size):
        data_dist.append(data_lines[i:i+chunk_size])

    split_data = []
    for single_data in data_dist:
        split_data.append("\n".join(single_data))

   
#     for i in range(len(split_data)):
#         print str(i) + "->", split_data[i]
            

    # generate data strings to send - size limited to 1e6-1 bytes
    for i in range(len(split_data)):
        split_data[i] = format_message(split_data[i])

    
    # prepare socket array
    mycontext['peer_sockets'] = [""] * len(mycontext['scoreboard'].keys())

    # send data to each replica
    for i in range(len(replicas)):
        replica = replicas[i]

        addr_parts = replica.partition(":")
        replica_ip = addr_parts[0]
        replica_port = int(addr_parts[2])

        print "opening connection to", replica
        socketobj = timeout_openconn(replica_ip, replica_port, timeout=10)
        socketobj.send(prim_str)
        socketobj.send(replica_str)
        socketobj.send(split_data[i])

        socketobj.settimeout(5)

        mycontext['scoreboard'][replica]["socket"] = socketobj

        print "receiving", socketobj.recv(1)


# only worried about active replicas now, add support for inactive peers later 
def construct_scoreboard(replica_list):
    mycontext['scoreboard'] = {}
    
    for replica in replica_list:
        active_entry = {}
        active_entry["socket"] = ""
        active_entry["active"] = 1
        active_entry["state"] = "Initialized"
        active_entry["heartbeat_failed"] = 0
        mycontext['scoreboard'][replica] = active_entry


def keep_scoreboard():
    
    active_replicas = mycontext['scoreboard'].keys()
    active_replicas.sort()

    mycontext['final_output'] = {}
    finished_replicas = 0

    scoreboard = mycontext['scoreboard']

    print "active_replicas:", active_replicas

    while finished_replicas < len(active_replicas):

        for replica in active_replicas:

            print
            print "retrying, only got %d of %d data" % (finished_replicas, len(active_replicas))

            print "scoreboard:", scoreboard
            # if this replica says it's done, skip it
            if scoreboard[replica]['state'] == "Done":
                continue

            print "current replica:", replica
            send_message(scoreboard[replica]['socket'],"heartbeat")
            
            cur_state = ""
            try:
                print "recving:"
                cur_state = recv_message(scoreboard[replica]['socket'],
                                         timeout=3)
            except SocketTimeoutError:
                print "timed out!"
                scoreboard[replica]['heartbeat_failed'] += 1
                if scoreboard[replica]['heartbeat_failed'] >= 3:
                    pass # add activating an inactive node here!
            except Exception, e:
                if "refused" not in str(e):
                    raise
            else:
                print "recving state (%s for %s).." % (cur_state, replica)
                scoreboard[replica]['state'] = cur_state
                if scoreboard[replica]['heartbeat_failed'] != 0:
                    scoreboard[replica]["heartbeat_failed"] -= 1

                # data is available for pickup!
                if cur_state == "ReducerDone":
                    red_data = recv_message_dict(scoreboard[replica]['socket'],
                                                 timeout=5)
                    mycontext['final_output'].update(red_data)
                    scoreboard[replica]['state'] = "Done"
                    finished_replicas += 1
                    print "received mapred data:", mycontext['final_output']
                    continue

            # give the replicas a bit of a break!
            sleep(5)
        
            
    # if we've reached this point, we've received all data
    buf = ""
    for key, value in mycontext['final_output'].iteritems():
        buf += key + "\t" + value + "\n"
                    
    fileobj = open("final_data.dat", "w")
    fileobj.write(buf)
    fileobj.close()
    print "finished data acquisition, job finished"


def send_message(socketobj, data):
    data = str(len(data)) + "*" + data
    socketobj.send(data)


def format_message(data):
    return str(len(data)) + "*" + data


def recv_message_dict(socketobj, initialread=2, timeout=None):
    serialized_dict = recv_message(socketobj, initialread, timeout)

    data_dict = {} 
    
    cur_key = ""
    for line in serialized_dict.split("\n"):
        if cur_key == "":
            cur_key = line
        elif line == "":
            cur_key = ""
        else:
            data_dict[cur_key] = line

    return data_dict


def recv_message(socketobj, initialread=2, timeout=None):
    buf = ""
    found = False

    # if timeout, we have a timeout_socket object; try recving, but can throw
    # a SocketTimeoutError here.
    if timeout:
        socketobj.settimeout(timeout)
        buf += socketobj.recv(1)
    
    while not found:
        buf += socketobj.recv(initialread)
        if "*" in buf:
            found = True

    parts = buf.partition("*")
    data_len = int(parts[0])
    
    data = parts[2]
    data += socketobj.recv(data_len - len(parts[2]))      
    return data



def pad_int_to_str(i, len):
    if i / (10**len) > 0:
        raise Exception("integer overflow; consider revising protocol") 
    return ('%0' + str(len) + "d") % i




def get_results(ip, port, sockobj, thiscommhandle, listencommhandle):
    
    # parse actual data, write to file
    buf = recv_message(sockobj)
    result_file = open("results.dat", "w")
    result_file.write(buf)
    result_file.close()


def print_usage():
    print """Usage: mappri.repy [options] <datafile> <replica-ip>:<replica-port>...
    -n               Do not parse the data file with \\t splitting key and value
    <datafile>       The data file to use in the map-reduce pipeline
    <replica>:<port> A list of the replica addresses to use for the job
"""
    exit()


if callfunc == 'initialize':
    mycontext['replica_list'] = {}

    if len(callargs) == 0: 
        print "* No arguments provided."
        print_usage()
    
    argpointer = 0
    port = 12346
    try:
        port = int(callargs[0])
        ip = getmyip()
        argpointer += 1
        
        if len(callargs) == 1:
            print "* No arguments provided"
            print_usage()

    except ValueError:
        ip = '127.0.0.1'

    mycontext['my_addr'] = ip + ":" + str(port)
#    print "my_addr:", mycontext['my_addr']

    if callargs[argpointer] == '-n':
        # do something special here, it's a no-op atm
        pass

    try:
        data_file = open(callargs[argpointer], "r")
        data = data_file.read()
        data_file.close()
        argpointer += 1
    except IOError:
        print "* Error reading file (does it exist?)"
        print_usage()

    input_replicas = callargs[argpointer:]
    if len(input_replicas) == 0:
        print "* List of replicas not found."
        print_usage()

    replica_list = []
    for replica in input_replicas:
        address = replica.partition(':')
        if not address[2]:
            print "* Port not found!"
            print_usage()
        replica_list.append(replica)

    construct_scoreboard(replica_list)

    # ensure that everything is sorted in the same way
    mycontext['replica_str'] = ""
    replica_list.sort()
    for replica in replica_list:
        mycontext['replica_str'] += format_message(replica)

    initalize_replicas(data)

    # once the job has started, we need to listen for reporting back 
    # by the replicas (right now, just handle the scoreboard!)
    
    keep_scoreboard()
