"""
<Description>
  Encapsilates the sending structure of a connection.

<Usage>
  TcpClient() - constructor, takes function to send through

  Handshakers:
    client.connect(...) - send SYN
    client.accept(...) - send SYNACK
    client.send(...) - send message
    client.close(...) - send FIN

  Callback:
    client.process(packet) - account for ACK

<Example>
  # Send out a message.
  include client.repy
 
  def callback(ip, port, mess, ch):
    packet = unpack(mess)
    client.process(packet)

  client = TcpClient()
  client.send("hello world", retries, timeout):
"""

include congestion_controller.repy
include constants.repy
include errors.repy
include packet.repy
include sleep_lock.repy
include utilities.repy

class TcpClient:
  """
  <Purpose>
    A connection has two personalities: 
      an active-sender-client and 
      a passive-receiver-server.

    This class is the active-sender-client.
  """

  def __init__(self, send_function):
    """
    <Purpose>
      Initializes varibles.
    
    <Arguments>
      send_function - a function to send through packets
    
    <Exceptions>
      None
    
    <Side Effects>
      Call getlock().
      Initializes a SleepLock
      Sets recv_window_size to default.
      Makes a random sequence number to start with.

    <Returns>
      New client object.
    """
    self.send_function = send_function
        
    self.packet_cache = {} # holds currently sent, unacked packets

    # kept track of + updated by Server
    self.recv_window_size = AD_WINDOW_SIZE # size of window

    # Congestion window -- slow start
    self.cong_controller = CongestionController()

    # Sliding window accounting
    self.send_window_size = None # size of window
    self.last_ack_recv = None # recv back by SYNACK, Beginning of window 
    self.last_seq_sent = mod_rand()  # End of window

    self.lock = getlock()
    self.sleep_lock = SleepLock()


  def send(self, data, destport, srcport, retries, timeout):
    """
    <Purpose>
      Actually sends data!  The most complex piece and
      whole point of all this stupid code.

      Does sliding window algorithm.  Uses slow start.

    <Arguments>
      data - string of data
      srcport - the source port to put into the packet headers
      destport - the destination port to put into the packet headers
      retries - send this many times after initial try
      timeout - wait time for an ACK, then retry
    
    <Exceptions>
      TimeoutError if not enough ACK's come back.
    
    <Side Effects>
      Sends UDP messages.
      Uses self.lock, sleep_lock
     
    <Returns>
      The bytes sent.
    """
    bytes_sent = 0  # total bytes sent

    # Split up into bytes
    data = list(data)
 
    # Send them -- Sliding Window Algorithm
    while True:
      self.lock.acquire()

      # the last ack changed and we have no more data
      # nor packets cached
      if not data and not self.packet_cache:
	self.lock.release()
	break;
      else: # reset variables
        last_ack = self.last_ack_recv # save to see if changes
        times_sent = 0
      self.lock.release()

      # while we haven't receieved a newer ack
      while True:
        self.sleep_lock.acquire()
        self.lock.acquire()
	# success
	if self.last_ack_recv != last_ack:
          self.sleep_lock.release()
	  self.lock.release()
	  break;

        # Pass True if unacked packets, False otherwise
        self.cong_controller.update(self.packet_cache)

        # give up
        if times_sent > retries:
          self.sleep_lock.release()
	  self.lock.release()
          raise TimeoutError("tried %d times" % times_sent)

        # Send or Re-send all those in cache
        # These must be modulo sorted because
        # Say we're never sending the first
        # one because the window size is too small.
        seqs = mod_sort(self.packet_cache.keys(), self.last_ack_recv, self.last_seq_sent)
        free_window = min(self.send_window_size, self.cong_controller.window_size)
        for seq_num in seqs:
          if len(self.packet_cache[seq_num].payload) <= free_window:
            free_window -= len(self.packet_cache[seq_num].payload)
            bytes_sent += len(self.packet_cache[seq_num].payload)
	    self.send_function(self.packet_cache[seq_num].pack())
          else: # can't send anymore
            free_window = 0
            break

        # make and send new packets
        while data and free_window:
            # can we send the whole window?
            payload_size = min(free_window, MSS)

            # make the packet
            next_seq_num = mod_add(self.last_seq_sent, 1)
            packet = TcpPacket(srcport, destport, next_seq_num, 0, "PSH", self.recv_window_size, ''.join(data[0:payload_size]))
            self.packet_cache[packet.seq_num] = packet
            self.last_seq_sent = mod_add(packet.seq_num, len(packet.payload) - 1)

            # send it as well
            bytes_sent += len(packet.payload)
            self.send_function(packet.pack())

	    # go around again?
            del data[0:payload_size]
	    free_window -= payload_size

        # increment tries
        times_sent += 1
        # take a deep breath
        self.sleep_lock.sleep(timeout)
        self.lock.release()

    return bytes_sent


  ##################################
  # Handshakers
  ###################################

  def accept(self, ack_num, destip, destport, \
                srcip, srcport, retries, timeout):
    """
    <Purpose>
      Used for SYNACK to send it, stops with Timeout or
      the window_size is set by ACK

    <Exceptions>
      raise UnknownStateError("whoa shouldn't have variables yet")
      raise TimeoutError("tried %d times" % times_sent)

    <Returns>
      Bytes sent.
    """
    self.lock.acquire()
    if not self.last_ack_recv is None:
      self.lock.release()
      raise UnknownStateError("whoa shouldn't have variables yet")
   
    synack = TcpPacket(srcport, destport, self.last_seq_sent, ack_num, "SYNACK", self.recv_window_size, "synacking...")
    start_ack = None 

    self.lock.release()
    # allow an interrupt here if you want.

    return self._send_packet(synack, start_ack, retries, timeout)

  def connect(self, destip, destport, \
                srcip, srcport, retries, timeout):
    """
    <Purpose>
      Used for SYN to send it, stops with Timeout or
      the window_size is set by SYNACK

    <Exceptions>
      raise UnknownStateError("whoa shouldn't have variables yet")
      raise TimeoutError("tried %d times" % times_sent)

    <Returns>
      Bytes sent.
    """
    
    self.lock.acquire()
    if not self.last_ack_recv is None:
      self.lock.release()
      raise UnknownStateError("whoa shouldn't have variables yet")
   
    syn = TcpPacket(srcport, destport, self.last_seq_sent, 0, "SYN", self.recv_window_size, "syncing...")
    start_ack = None 

    self.lock.release()
    # allow an interrupt here if you want.

    return self._send_packet(syn, start_ack, retries, timeout)


  def close(self, destip, destport, srcip, srcport, retries, timeout):
    """
    <Purpose>
      Used for FIN to send it, stops with Timeout or ACK of FIN.

    <Exceptions>
      raise TimeoutError("tried %d times" % times_sent)

    <Returns>
      Bytes sent.
    """
    self.lock.acquire()
    if self.last_ack_recv is None:
      self.lock.release()
      raise NotConnectedError("must be connected to close")
   
    self.last_seq_sent = mod_add(self.last_seq_sent, 1)
    fin = TcpPacket(srcport, destport, self.last_seq_sent, 0, "FIN", self.recv_window_size, "fining...")
    start_ack = self.last_seq_sent

    self.lock.release()
    # allow an interrupt here if you want.

    return self._send_packet(fin, start_ack, retries, timeout)


  ###########################
  # Callback
  #############################

  def process(self, packet):
    """
    <Purpose>
      Handles ACK's (and therefore SYNACK's)
      Sets up variables and stops sleeping. 
   
    <Arguments>
      packet
        the ACK or SYNACK
    
    <Exceptions>
      raise NotConnectedError("no packets sent to ack")
    
    <Side Effects>
      Changes window size and seq nums and last ack recv
     
    <Returns>
      True if the packet is a SYNACK and valid.
    """
    self.lock.acquire()
    valid = False
    if self.last_seq_sent == None:
      self.lock.release()
      raise NotConnectedError("no packets sent to ack")
    
    # handshake
    elif packet.is_synack():
      valid = self.process_synack(packet)

    # special case for payloadack of synack
    elif packet.is_payload():
      self.process_psh(packet)

    # usually
    elif packet.is_ack():
      valid = self.process_ack(packet)
    else:
      pass

    self.lock.release()
    return valid

  ##########
  # Getters/Setters
  ##########

  def get_seq_num(self):
    """
    <Purpose>
      atomic getter for Server to know seq num
    """
    self.lock.acquire()
    seq_num = mod_add(self.last_seq_sent, 1)
    self.lock.release() 
    return seq_num

  def set_recv_window(self, size):
    """
    <Purpose>
      atomic setter for Server to specify recv window
    """
    self.lock.acquire()
    self.recv_window_size = size
    self.lock.release() 

  ##########
  # Private
  ##########
  def process_synack(self, packet):
    """
    <Purpose>
      Accounts a synack.
      Verifies if valid ACK num of SYN sent.

    <Returns>
      True if the SYNACK is valid, false otherwise.
    """
    is_valid = False
    if self.last_ack_recv is None and packet.ack_num == mod_add(self.last_seq_sent,1):  # new synack
      self.last_ack_recv = packet.ack_num
      self.sleep_lock.wakeup()
      is_valid = True 

    elif not self.last_ack_recv is None and packet.ack_num == mod_add(self.last_seq_sent,1):  # old synack
      if self.last_ack_recv != mod_add(self.last_seq_sent, 1): # never will happen
	raise UnknownStateError("cannot be true that old synack is not current ack")
      is_valid = True 

    elif packet.ack_num != mod_add(self.last_seq_sent, 1):  # bad synack
      pass
    else: # no other cases
      pass

    if is_valid:
      # reset window
      self.send_window_size = packet.window 

    return is_valid


  def process_psh(self, packet):
    """
    <Purpose>
      Accounts to stop resending the SYNACK and
      Setup state for sending messages.
    
    <Exceptions>
#      raise AlreadyConnectedError("how could I pshack when already established")
    
    <Side Effects>
      Changes window size and last ack recv
      Wakes up sleeper if needed.
     
    <Returns>
      Nothing.
    """
#    if not self.send_window_size is None \
#	or not self.last_ack_recv is None:
#      raise AlreadyConnectedError("how could I pshack when already established")
	
    self.send_window_size = packet.window
    
    # assume they rcvd the right seq num
    self.last_ack_recv = mod_add(self.last_seq_sent, 1)

    self.sleep_lock.wakeup()  # stop resending SYNACK


  def process_ack(self, packet):
    """
    <Purpose>
      Handles ACK's from receiver.

      Throws an error if the last seq number is undefined
      or if the packet is not an ack.
      Accounts to stop resending the acked packet.
    
    <Arguments>
      packet
        the ACK
    
    <Exceptions>
      raise UnknownStateError("expected an ack")
    
    <Side Effects>
      Changes window size and seq nums and last ack recv
     
    <Returns>
      True if the ACK is new, False otherwise.
    """
    if not packet.is_ack():
      raise UnknownStateError("expected an ack")

    # one past the last seq
    lateist_ack = mod_add(self.last_seq_sent, 1)

    # ack of synack
    if self.last_ack_recv is None \
	 and lateist_ack == packet.ack_num: # good ack of syn
      self.last_ack_recv = packet.ack_num # == last_seq_sent
      self.send_window_size = packet.window 
      self.sleep_lock.wakeup()
      return True
    elif self.last_ack_recv is None: # bad ack of syn
      return False 
    elif self.last_ack_recv == lateist_ack: # already all acked
      return False
    elif self.last_ack_recv == packet.ack_num:
      return False # already ACK'd duplicate
    else:  # keep going
      pass

    earliest_ack = mod_add(self.last_ack_recv, 1)
    # if a newer ack that falls in sent
    if mod_between(earliest_ack, packet.ack_num, lateist_ack):
      last_seq = mod_sub(packet.ack_num, 1)
      # -- Cumulative ACK Processing
      for seq_num in self.packet_cache.keys():  
        # Delete all packets below this recvd ack_num.
        # We never received their acks but the sender must have
        # received the payloads of all seq_num before this ack by
        # the definition of our protocol.
        if mod_between(self.last_ack_recv, seq_num, last_seq):
          del self.packet_cache[seq_num]

	  # account for ACK
	  self.cong_controller.incr()

      # all done up to this ack
      self.last_ack_recv = packet.ack_num

      # reset window
      self.send_window_size = packet.window 

      # all in window are done so stop waiting
      if not self.packet_cache:
        self.sleep_lock.wakeup()

      return True

    else: # old ack 
      return False #ignore it


  def _send_packet(self, packet, start_ack, retries, timeout):
    """ 
    <Purpose>
      Send a FIN, SYN, SYNACK, etc. packet again and again until ACK'd.

    <Arguments>
      start_ack - how we know we've been ACK'd.  May be == None.
    """

    times_sent = 0
    bytes_sent = 0
    while True:
      self.sleep_lock.acquire()
     
      self.lock.acquire()
      if (not self.last_ack_recv is None and start_ack is None) or \
	  (not self.last_ack_recv is None and self.last_ack_recv == mod_add(start_ack, 1)):
        self.lock.release()
        self.sleep_lock.release()
        break
      
      # give up
      if times_sent > retries:
        self.sleep_lock.release()
        raise TimeoutError("tried %d times" % times_sent)

      bytes_sent += self.send_function(packet.pack())

      times_sent += 1
      self.sleep_lock.sleep(timeout)
      self.lock.release()

    return bytes_sent


