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

#begin include timeout_socket.repy
"""
<Description>
  A socket that guarentees the receipt of a message.   Raises TimeoutError if it does not
  receive any message before a given timeout.
  If actually receives the message, returns the message and continues.

<Usage>
  Text-replacable for Repy Sockets:
    timeout_openconn(desthost, destport, localip=None, localport=None, timeout = 5)
    timeout_waitforconn(localip, localport, function)

  Object:
    sockobj.send(data)
    sockobj.recv(bytes)
    sockobj.close()

<Date>
  Sun Mar  1 10:27:35 PST 2009

<Example>
  # hello world
  include sockettimer.repy

  def callback(ip, port, timeout_sockobj, commhandle, listenhandle):
    hw_message = timeout_sockobj.recv(1047)

    # cleanup
    stopcomm(commhandle)
    stopcomm(listenhandle)
    timeout_sockobj.close()

    print hw_message # => "hello world!"
  
  def server():
    sockobj = timeout_waitforconn(getmyip(), 12345, callback)

  def client():
    sockobj = timeout_openconn(getmyip(), 12345)
    sockobj.send("hello world!")

  def main():
    server()
    client()
    exitall()

  if callfunc == 'initialize':
    main() 
"""

class SocketTimeoutError(Exception):
  """The socket timed out before receiving a response"""

def timeout_openconn(desthost, destport, localip=None, localport=None, timeout = 5):
  """
  <Purpose> 
    Wrapper for Repy like socket interface

  <Args>
    Same as Repy openconn

  <Exception>
    Timeout exception if the dest address doesnt respond.

  <Returns>
    socket obj on success
  """

  tsock = TimeoutSocket()
  tsock.settimeout(timeout)
  if localip and localport:
    tsock.bind((localip, localport))
  tsock.connect((desthost, destport))
  return tsock

def timeout_waitforconn(localip, localport, function):
  """
  <Purpose> 
    Wrapper for Repy like socket interface

  <Args>
    Same as Repy waitforconn

  <Side Effects>
    Sets up event listener which calls function on messages.

  <Returns>
    Handle to listener.
  """

  tsock = TimeoutSocket()
  tsock.bind((localip, localport))
  tsock.setcallback(function)
  return tsock.listen()

class TimeoutSocket:
  """
  <Purpose>
    Provide an socket object like the Repy usual one.

  <Side Effects>
    Uses a getlock() to watch for a timeout
    Uses waitforconn and openconn to simulate socket
  """

  ################
  # Constructors
  ################

  def __init__(self):
    """ Constructor for socket """
#    self.lock = getlock() # general lock BUG: Do we need to lock everything?
    self.timeout_lock = getlock() # special lock for Timeout condition
    self.timeout = 5 # seconds to wait

    # user vars   
    self.local_address = None # ip, port
    self.remote_address = None # ip, port
    self.callback = None # the user's function to call

    # repy socket vars
    self.sockobj = None #  the Repy socket
    self.commhandle = None # the current comm
    self.listencommhandle = None # the listener comm

  ################
  # Mutator methods
  #################

  def settimeout(self, value):
    """ Setter for timeout"""
    self.timeout = value

  def setcallback(self, function):
    """ Setter for callback function"""
    self.callback = function

  ####################
  # Public Methods
  ####################

  def bind(self, local_address = None):
    """
    <Purpose>
      Set local address

    <Args>
      Tuple of (ip, port) local.
    """
    self.local_address = local_address

  def listen(self):
    """
    <Purpose>
      Listen for peer
    
    <Side Effects>
      Calls Repy waitforconn()
    """
    return self._waitforconn()

  def connect(self, remote_address):
    """
    <Purpose>
      Connect to peer.

    <Args>
      Tuple of (ip, port) remote.
   
    <Side Effects>
      Calls Repy openconn.
    """
    self.remote_address = remote_address
    self._openconn()

  def recv(self, maxLen): # timeout as optional arg ???
    """
    <Purpose>
      If it fails to finish within the timeout, I close the socket and raise a
      TimeoutError exception. I.e. if there's no message, we call it an error
      and raise it.
      
    <Arguments>
      maxLen - bytes to recv

    <Exception>
      Raises TimeoutError exception if the recv times out
      without receiving a message.

    <Side Effects>
      Closes the connection if times out.

    <Returns>
      The message.
    """
    return self._recv_or_close(maxLen)

  def send(self, data):
    """
    <Purpose>
      Just like normal Repy socket.  Sends messages.
      
    <Arguments>
      data - the string message

    <Exception>
      Same as Repy socket.
 
    <Returns>
      The bytes sent.
    """
    return self._send(data)

  def close(self):
    self.local_address = None # ip, port
    self.remote_address = None # ip, port
    self.callback = None # the user's function to call

    self.sockobj.close()
    self.sockobj = None #  the Repy socket
    stopcomm(self.commhandle)
    self.commhandle = None # the current comm
    stopcomm(self.listencommhandle)
    self.listencommhandle = None # the listener comm


  ########################
  # Private
  #########################

  def _openconn(self):
    """Handle current state variables and call Repy openconn."""

    destip, destport = self.remote_address
    if self.local_address:
      srcip, srcport = self.local_address
      self.sockobj = openconn(destip, destport, srcip, srcport, self.timeout)
    else:
      self.sockobj = openconn(destip, destport)

  def _waitforconn(self):
    """Setup way between Repy waitforconn event"""
    localip, localport = self.local_address
    self.listencommhandle = waitforconn(localip, localport, self._callback)
    return self.listencommhandle

  def _callback(self, ip, port, sockobj, ch, lh):
    """Pass on through to user callback"""
    self.sockobj = sockobj
    self.listencommhandle = lh # same as the 1st from wait for comm, right?
    self.commhandle = ch # should we care?
    
    print "sockettimeout|remote_address:", self.remote_address
    if not self.remote_address:
      self.remote_address = (ip, port)
    else: 
      raise Exception("what! peer does not match?")

    self.callback(ip, port, self, ch, lh)

  def _send(self, data):
    """Send data"""
    return self.sockobj.send(data)

  def _recv(self, maxLen):
    """Recv data of length maxLen"""
    return self.sockobj.recv(maxLen)

  def _recv_or_close(self, amount):
    """Raise the Timeout Error if no receipt.  Keep track by timeout_lock."""
    timerhandle = settimer(self.timeout, self._clobbersocket, ())
    try:
      retdata = self._recv(amount)
    except Exception, e:
      # if it's not the timeout, reraise...
      if self.timeout_lock.acquire(False):
        raise
      raise SocketTimeoutError
    
    # I acquired the lock, I should stop the timer because I succeeded...
    if self.timeout_lock.acquire(False):
      # even if this isn't in time, the lock prevents a race condition 
      # this is merely an optimization to prevent the timer from ever firing...
      canceltimer(timerhandle)
      self.timeout_lock.release() # Alper's bug 3/10/09
      return retdata
    else:
      raise SocketTimeoutError

  def _clobbersocket(self):

    # alpers - don't close the socket if we timeout!  It might be just an error
    return

    """If I can acquire the lock without blocking, then close the socket to abort"""
    if self.timeout_lock.acquire(False):
      self.close()


############################
# Deprecated functions
##############################

# private function...
def sockettimeout_clobbersocket(sockobj,mylock):
  # if I can acquire the lock without blocking, then close the socket to abort
  if mylock.acquire(False):
    sockobj.close()

# if it fails to finish within the timeout, I close the socket and raise a
# SocketTimeout exception...
def sockettimeout_recv_or_close(sockobj, amount, timeout):
  # A lock I'll use for this attempt
  mylock = getlock()
  timerhandle = settimer(timeout,clobbersocket, (sockobj, mylock))
  try:
    retdata = sockobj.recv(amount)
  except Exception, e:
    # if it's not the timeout, reraise...
    if mylock.acquire(False):
      raise
    raise SocketTimeout
    
  # I acquired the lock, I should stop the timer because I succeeded...
  if mylock.acquire(False):
    # even if this isn't in time, the lock prevents a race condition 
    # this is merely an optimization to prevent the timer from ever firing...
    canceltimer(timerhandle)
    return retdata
  else:
    raise SocketTimeout


#end include timeout_socket.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()
