
'''
- header at top of file : docstring
- header for each function (you already have the descriptions) : docstring
- more space between functions, 3-4 space lines
'''

include sockettimeout.repy

include mapper.repy
include reducer.repy
include hasher.repy


def get_data(ip, port, socketobj, thiscommhandle, listencommhandle):
    """ Listens for connections on a well-defined port, imports data files """

    # get a list of all of our neighbors!
    mycontext['primary'] = recv_message(socketobj)
    print "Primary init thread: got primary loc: ", mycontext['primary']

    # we need to know how many peers we have..
    mycontext['num_peers'] = int(socketobj.recv(4))
    print "Primary init thread: got num_peers: ", mycontext['num_peers']

    mycontext['peers'] = []
    for i in range(mycontext['num_peers']):
        mycontext['peers'].append(recv_message(socketobj))

    # parse and save data file:
    buf = recv_message(socketobj)
    print "Primary init thread: got file data"
    dataobj = open("map_data.dat", "w")
    dataobj.write(buf)
    dataobj.close()

    # make sure each replica has the same order!
    # this line isn't needed --> mycontext['peers'].sort()
    print "Primary init thread: got my peers: ", mycontext['peers']

    # which peer am I?  save this for future reference..
    address_str = mycontext['myip'] + ":" + str(mycontext['myport']-1)
    mycontext['my_peer_index'] = mycontext['peers'].index(address_str)        
    print "Primary init thread: my_peer_index is", mycontext['my_peer_index']

    # start the peer_sockets variable
    mycontext['peer_sockets'] = [""] * mycontext['num_peers']
        
    # save the primary socket for heartbeats
    mycontext['primary_socket'] = socketobj

    # try ACKing
    socketobj.send("i")

    # destroy the listen socket as we're done initializing
    mycontext['state'] = 'Initialized'
    stopcomm(listencommhandle)


    
# respond to any queries from the primary
def heartbeat_response():
    msg = ""
    while True:
        try:
            msg = recv_message(mycontext['primary_socket'], timeout=5)
        except SocketTimeoutError:
            print "Primary control thread: timed out waiting for heartbeat"
        else:
            if "control" in msg:
                pass  # change peers in here
            elif msg == "heartbeat":
                state = mycontext['state']
                print "Primary control thread: sent state %s to primary" % state
                send_message(mycontext['primary_socket'], state)

                if state == "ReducerDone":
                    print "Primary control thread: sending datadict:", mycontext['reduce_result']
                    send_message_dict(mycontext['primary_socket'], 
                                      mycontext['reduce_result'])
                    mycontext['state'] = "Done"
                    break



def init_replica_sockets():
    # we only want to make one socket for each link, so let's be smart about it
    # and only make new connections to equal or higher indexed peers..

    print "entered init_replica_sockets"

    # put a fake socket into mycontext['peer_sockets'] for self
    mycontext['peer_sockets'][mycontext['my_peer_index']] = "self"

    # if we're the last peer, just skip this.  we'll listen for all sockets.
    if mycontext['my_peer_index'] + 1 >= len(mycontext['peers']):
        print "Main thread: quit early, I'm last in list."
        return

    print "Main thread: continuing init_replica_sockets()"
    # + 1 is added to skip creating self-socket
    peer_subset = mycontext['peers'][mycontext['my_peer_index'] + 1:]
    for peer in peer_subset:
        addr_parts = peer.partition(":")
        peer_index = mycontext['peers'].index(peer)

        print "Main thread: attempting to connect to peer %d, (addr: %s:%d)" % (peer_index, addr_parts[0], int(addr_parts[2])+1)
        socketobj = timeout_openconn(addr_parts[0], int(addr_parts[2])+1,
                                     mycontext['myip'], 12347)
        socketobj.send("hello")
        socketobj.recv(1)

        mycontext['connection_lock'].acquire()
        mycontext['peer_sockets'][peer_index] = socketobj
        mycontext['peer_conn_left'] -= 1
        mycontext['connection_lock'].release()

        print "succesfully communicated with", peer



# listens for incoming tcp connections, establishes a port
def recv_init_replica_conn(ip, port, sockobj, thiscommhandle, listencommhandle):
    data = sockobj.recv(5)
    
    print "Peer init thread: receiving peer_init message from %s, %s" % (ip, port)

    if not data == "hello":
        print "Received an unintelligible message from a peer %s:%d, %s" % (ip, port, data)
        exitall()
    else:
        sockobj.send("i")

        # save socket obj for later
        incoming_index = mycontext['peers'].index(ip+":12344")

        mycontext['connection_lock'].acquire()
        mycontext['peer_sockets'][incoming_index] = sockobj

        # decrement our connection counter
        mycontext['peer_conn_left'] -= 1

        print "Peer init thread: adding new socket to index " + str(incoming_index)
        print "Peer init thread: peer_sockets:", mycontext['peer_sockets']

        mycontext['connection_lock'].release()

        print "Peer init thread: successfully received init message from", ip+":"+str(port)


# listens for incoming map data (this is run asynchronously)
def recv_peer_map_data():
    peer_socket_recvd = []    

    print "Peer datarecv thread: peer_sockets:", mycontext['peer_sockets']

    while len(peer_socket_recvd) != mycontext['num_peers']:

        # cycle through all the sockets, trying to receive
        for socket in mycontext['peer_sockets']:

            print "Peer datarecv thread: trying to receive on socket", socket
            print "Peer datarecv thread: received from", peer_socket_recvd
            print "Wanted %d receptions, only got %d" % (len(mycontext['peer_sockets']), len(peer_socket_recvd))

            # pass over the socket if we've received data already
            if socket in peer_socket_recvd:
                continue

            # pass over our own socket
            if socket == "self":
                continue
            
            # add socket if we've added data to ourself in partition
            if socket == "selfrecvd":
                peer_socket_recvd.append(socket)
                continue

            # try to receive
            recv_data = {}
            try:
                print "Peer datarecv thread: recving on socket", socket
                recv_data = recv_message_dict(socket, timeout=5)
            except SocketTimeoutError:
                print "Peer datarecv thread: timed out on socket", socket, "... continuing"
                continue
            else:
                mycontext['connection_lock'].acquire()
                mycontext['reduce_data'].update(recv_data)
                print "Peer datarecv thread: received data"
                peer_socket_recvd.append(socket)
                mycontext['connection_lock'].release()
            
    print "Peer datarecv thread: finished recving data"
    mycontext['state'] = "ReducerWaiting"



# Assumptions to make this simpler:
# - all this data fits in memory (<2 GB) in the variable map_result
# - data is stored in the files/string as "(key)(\t)(value)"  
def do_map():
    mycontext['state'] = "Mapping"

    data = open("map_data.dat", "r")

    map_result = []
    for line in data:
        line_parts = line.partition('\t')
        # I assume that results are returned in the form "<key>\t<value>"
        # map.mapper takes key, value as two separate arguments
        map_result.extend(map_func(line_parts[0], line_parts[2]))
#        print "map_result" , map_result
    map_result.sort()
    return map_result
    

    
# the user must define their own partition function (hash_func)
def partition(map_result):
    mycontext['state'] = "Partitioning"

    key_value = {}

    # generate a key-value dict
    # ! could use simple dict([]) here, but duplicate keys get overwritten
    for kv_pair in map_result:
        if kv_pair[0] in key_value:
            key_value[kv_pair[0]].append(kv_pair[1])
        else:
            key_value[kv_pair[0]] = [kv_pair[1]]

    print "map_result" , map_result
    print
    print "kv_pair", kv_pair
    print
    print "key_value", key_value
    print

    # try to partition objects bashed on their hash
    partition_hash = {}
    for key, values in key_value.iteritems():
        hashcode = hash_func(key)
        
        # the following is always key in index 0, values following
        # this is so we don't lose track of the key if the hash_func is not
        # totally unique/reversible
        values.insert(0,key)
        if hashcode in partition_hash:
            partition_hash[hashcode].append(values)
        else:
            partition_hash[hashcode] = [values]

    # we should have a grouping of elements by hash now
    print "partition hash:", partition_hash
    print

    # attempt to partition elements based on the number of replicas,
    #   first, get all the hash values and sort them
    hash_list = partition_hash.keys()
    hash_list.sort()

    # initialize the peer_data list
    peer_data = []
    for index in range(mycontext['num_peers']):
        peer_data.append({})
    

    # go through each hash, and assign it to a replica
    for hashcode in hash_list:
        cur_replica = hashcode % mycontext['num_peers']
        
        for kv_pairs in partition_hash[hashcode]:
            print "kv_pairs", kv_pairs
            kv = {kv_pairs[0]: kv_pairs[1:]}
            peer_data[cur_replica].update(kv)


    # print out the peer data, just for my reference
    index = 0
    print "peer_data:"
    for peer in peer_data:
        print index, "->", peer
        index += 1

    # send data to replicas, each replica should have a dictionary of kv pairs  
    for peer_index in range(mycontext['num_peers']):
        if peer_index == mycontext['my_peer_index']:
            print "Main thread: sending data to self"
            mycontext['connection_lock'].acquire()
            mycontext['reduce_data'].update(peer_data[peer_index])
            mycontext['peer_sockets'][mycontext['my_peer_index']] = "selfrecvd"
            mycontext['connection_lock'].release()
        else:
            print "Main thread: sending data to peer %d..." % peer_index
            mycontext['connection_lock'].acquire()
            socketobj = mycontext['peer_sockets'][peer_index]
            send_message_dict(socketobj, peer_data[peer_index])
            mycontext['connection_lock'].release()

    
def do_reduce():
    #data = open("reduce_data.dat", "r")
    data = mycontext['reduce_data']

    print "reduce_data", mycontext['reduce_data']
        
    reduce_result = {}
    results_list = []


    for key, values in data.iteritems():
        print "reducing..."
        returned_dict = reduce_func(key, values)
        results_list.append(returned_dict)

    for single_dict in results_list:
        for key, value in single_dict.iteritems():
            if key in reduce_result:
                reduce_result[key] += " " + value
            else:
                reduce_result[key] = value
                          
    print "Reducing: reduce_result:", reduce_result

    mycontext['reduce_result'] = reduce_result


# TODO...
def report_results(map_results):
    pass
    

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



def send_message_dict(socketobj, data_dict):
    buf = ""
    for key,values in data_dict.iteritems():
        buf += str(key) + "\n"
        if isinstance(values, list):
            for value in values:
                buf += str(value) + "\n"
        else:
            buf += str(values) + "\n"
        
        # after last value, add another return to close key
        buf += "\n"

    send_message(socketobj, buf)



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
            data_dict[cur_key] = []
        elif line == "":
            cur_key = ""
        else:
            data_dict[cur_key].append(line)

    return data_dict


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

#    print "recv_message: trying to recv from " + socketobj.remote_address

    # if timeout, we have a timeout_socket object; try recving, but can throw
    # a SocketTimeoutError here.
    if timeout:
        print "timeout was set to", timeout, "- entered timeout recv section"
        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



if callfunc == 'initialize':
    mycontext['num_mappers'] = 1
    mycontext['num_reducers'] = 1
    mycontext['state'] = 'Ready'
    
    if len(callargs) > 1:
        raise Exception("too many args")
    elif len(callargs) == 1:
        port = int(callargs[0])
        ip = getmyip()
    else:
        port = 12345
        ip = '127.0.0.1'

    # save my ip/port for future reference
    mycontext['myip'] = ip
    mycontext['myport'] = port
    
    # wait for primary initialization data
    print "Main thread: waiting for connection on ", ip, ":",  port
    listencommhandle = timeout_waitforconn(ip, port-1, get_data)
    
    # block until we've been initialized
    while mycontext['state'] == 'Ready':
        sleep(.1)

    # maintain connection with primary, keep scoreboard updated
    # this is the main channel of communication with the primary
    settimer(0, heartbeat_response, ())
    
    # try to open connections to all hosts
    mycontext['connection_lock'] = getlock()
    mycontext['peer_conn_left'] = mycontext['num_peers']

    # listen for connections from peers
    listen_replica_init = timeout_waitforconn(ip, port, recv_init_replica_conn)
    print "Main thread: listening on %s:%d for peer connections" % (ip, port)

    # sleep to slow the socket thing down, we're working over WAN so it's tricky
    default_timeout = 4 * mycontext['num_peers']
    lag = 4 * mycontext['my_peer_index']
    sleeptime = default_timeout - lag

    print "Main thread: sleeping for %d seconds to allow listeners to catch up" % sleeptime
    sleep(sleeptime)

    # open connections to peers
    init_replica_sockets()

    while mycontext['peer_conn_left'] > 1:
        print "Main thread: sleeping for peer_conn_left (%d)" % mycontext['peer_conn_left']
        sleep(0.1)

    # destroy the listener
    print "Main thread: destroyed peer-peer initialization listener"
    stopcomm(listen_replica_init)
    mycontext['state'] = 'Connected'

    # start listening for map data from peers
    # !! CAN'T DO IT LIKE THIS: use the sockettimeout thing?
    # listencommhandle = waitforconn(ip, port, recv_peer_map_data)
    
    # this is the better way to do it... sockettimeout.repy
    mycontext['reduce_data'] = {}
    recv_handle = settimer(0, recv_peer_map_data, ())

    # start mapping, synchronous call
    map_result = do_map()

    # send map results to all reducers, split as necessary
    partition(map_result)
    
    # block until all partition data is acquired
    while mycontext['state'] != 'ReducerWaiting':
        print "Main thread: sleeping for state ReducerWaiting (state: %s)" % mycontext['state']
        sleep(0.5)


    # start reducing, synchronous call 
    do_reduce()
    mycontext['state'] = "ReducerDone"

    # wait for a primary heartbeat to come in, then return data, change state
    # and terminate
    while mycontext['state'] != "Done":
        print "Main thread: sleeping for state Done (state: %s)" % mycontext['state']
        sleep(0.5)




    
