#!/usr/bin/env python

import sys, socket, select, time, json, random, math, collections

# Amount of time to wait to send a heartbeat
HEARTBEAT_TIME = .1

# Timeout for an append entry request to be re-sent
APPEND_ENTRY_TIMEOUT = .1

# Respond to a client's message
def handle_client_message(msgtype, mid, client_address, msg):
	global my_id
	global sock
	global state
	global current_leader
	global command_queue
	global current_term
	global log

	# If we are not the leader, send a redirect message
	if not (state == "leader"):
		#print "SENDING REDIRECT " + current_leader
		msg = {'src': my_id, 'dst': client_address, 'leader': current_leader, 'type': 'redirect', 'MID': mid}
		sock.send(json.dumps(msg))
	# If we are the leader, deal with the user's message
	elif state == "leader":
		#print "Add entry to queue"

		# Add the entry to the queue
		entry = {'client': client_address, 'type': msgtype, 'key': msg['key'], 'term': current_term, 'MID': mid, 'leader': current_leader}

		# Add value to entry if it is in the message
		if 'value' in msg:
			entry['value'] = msg['value']

		command_queue.append(entry)

		# See if the leader needs to be told
		# to begin the commit process
		if len(command_queue) == 1:
			handle_queue_command()

# Handle a general command
def handle_queue_command():
	global command_queue
	global log
	global num_replica_logs

	# Make sure there are items in the queue
	if len(command_queue) >= 1:
		# First append it to our own log
		# each entry has a term and command
		# keep track of what command it is and the
		# key and value associated with it
		# as well as the client who sent it
		entry = command_queue[0]
		
		# Make sure that this message has not been
		# dealt with yet by checking the message id
		for l in log:
			# If we have seen it, remove it from queue
			# and go to next command
			if 'MID' in l and l['MID'] == entry['MID']:
				#print "ALREADY HAVE THIS IN LOG"
				del command_queue[0]	
				handle_queue_command()
				return

		#print "ADD ENTRY TO LEADER LOG"
		log.append(entry)

		num_replica_logs = 1
	
		send_append_entries_rpc(len(log), 'FFFF')

# The leader initializes the next index
# for each follower to be its own last
# log index + 1
def initialize_follower_index():
	global follower_next_index
	global log
	global replica_ids

	# Get the index of of the last entry in the log
	index = len(log)

	for r in replica_ids:
		follower_next_index[r] = index + 1
		
# The leader must keep track of a timeout
# for each of the append entries it sends out
# to the followers, do this by keeping track
# of the time of when the message was sent out
# will be initialized to -1
def initialize_append_entry_timeout():
	global replica_ids
	global append_entries_timeout_tracker
	
	for r in replica_ids:
		append_entries_timeout_tracker[r] = -1
	 
# Get a random number for the election
# timeout
def get_election_timeout():
	i = random.randint(250, 500)
	
	# want the number to be in ms so
	# divide by 100
	return i / 100

# Check to see if the election timeout
# has occurred
def check_election_timeout(current_time):
	global last_received_rpc
	global election_timeout

	# if a timeout has occurred, begin an election
	if current_time - last_received_rpc >= election_timeout:
		begin_election()

# Check to see if the rpc timeout has occurred
def check_send_heartbeat(current_time):
	global last_sent_rpc
	global log

	if current_time - last_sent_rpc >= HEARTBEAT_TIME:
		send_append_entries_rpc(-1, 'FFFF')

# Check to see if we need to resend an append entry
# to one of the followers
def check_append_entry_timeout(current_time):
	global replica_ids
	global append_entries_timeout_tracker
	global follower_next_index
	
	# If a timeout occurs, resend
	for r in replica_ids:
		if (not append_entries_timeout_tracker[r] == -1) and current_time - append_entries_timeout_tracker[r] >= APPEND_ENTRY_TIMEOUT:
			#print "RESENDING APPEND ENTRY DUE TO TIMEOUT to follower " + str(r) + " entry number " + str(follower_next_index[r])
			send_append_entries_rpc(follower_next_index[r], r)	
	
# Begin an election
def begin_election():
	global my_id
	global state
	global current_term
	global num_votes
	global voted_for
	global current_leader
	global last_received_rpc

	#print "Election started by " + my_id

	# increment term by 1
	current_term = current_term + 1

	# transition to candidate state
	state = "candidate"

	# vote for yourself
	voted_for = my_id

	# increment number of votes
	num_votes = 1

	# set current leader to ourselves
	current_leader = my_id

	# reset for election timeout
	last_received_rpc = time.time()

	# issue out request rpcs to everyone else
	send_request_vote_rpcs('FFFF')

# Send out request vote rpcs
def send_request_vote_rpcs(destination):
	global sock
	global my_id
	global current_leader
	global current_term
	global last_sent_rpc
	global log 

	#print "SENDING VOTE REQUEST"
	last_sent_rpc = time.time()
	
	#send the log length so that we can compare them to the replica's 
	log_length = len(log)
 
	msg = {'src': my_id, 'dst': destination, 'leader': my_id, 'type': 'request_vote', 'term': current_term, 'log_length': log_length}
	sock.send(json.dumps(msg))

# Send out append entries rpc
# Takes in the the entries index
# and the follower(s) to send it to
def send_append_entries_rpc(entry_index, destination):
	global my_id
	global current_term
	global sock
	global log
	global commit_index
	global last_sent_rpc
	global append_entries_timeout_tracker

	# Get the approriate entry/entries based on the entry index
	# if it is -1, it means it is a heartbeat so send
	# nothing, otherwise get the entry at log[entry_index - 1]
	# need to subtract 1 when getting the entry out of the log
	# since the list starts at index 0
	if entry_index == -1:
		entry = ''
	else:
		# add all of the entries up to the last one in
		# the leader's log
		temp = entry_index - 1
		entry = []
		while temp < len(log) and (temp - entry_index - 1) < 50:
			entry.append(log[temp])
			temp = temp + 1
	
	# Get the last index and term
	last_index = entry_index - 1
	if (last_index > 0):
		last_term = log[last_index]['term']
	else:
		last_term = 0
	
	# setup the message
	msg = {'src': my_id, 'dst': destination, 'leader': my_id, 'type': 'append_entries', 'term': current_term, 'entries': entry, 'commit_index': commit_index, 'last_index': last_index, 'last_term': last_term}

	# update the append_entries_timeout list
	if not entry == '' and not destination == 'FFFF':
		append_entries_timeout_tracker[destination] = time.time()
	elif not entry == '':
		for r in replica_ids:
			append_entries_timeout_tracker[r] = time.time()

	# update
	last_sent_rpc = time.time()

	# send the message
	sock.send(json.dumps(msg))

# If a message of type request_vote is
# received, handle it here - a replica
# gets one of these messages if another
# server is trying to become a leader or
# if itself is trying to become a leader and 
# is receiving responses
def handle_request_vote_messages(msg):
	global current_term
	global current_leader
	global sock
	global num_votes
	global majority
	global voted_for
	global state
	global my_id
	global log
	global command_queue
	global commit_index
	global commit_queue

	# if this is a response to our election
	if 'vote_granted' in msg and state == "candidate":
		#print "RECEIVED A VOTE FROM " + msg['src']

		# check the answer, only do something if yes
		if msg['vote_granted'] == 'yes' and msg['term'] == current_term and msg['leader'] == my_id:
			num_votes = num_votes + 1			
			
			# if we have enough votes to become leader
			# send out append entry messages to everyone
			# also setup follower next index dictionary
			if num_votes >=  majority:
				#print "BECOMING LEADER " + current_leader
				state = "leader"
				initialize_follower_index()
				initialize_append_entry_timeout()
				send_append_entries_rpc(-1, 'FFFF')
				command_queue = []
				
				# Update the commit index to length of log
				# and keep track of past log entries to commit later
				while commit_index < len(log):
					commit_queue.append(log[commit_index - 1])
					commit_index = commit_index + 1
					
				handle_queue_command()
											
	# if this is a message from a candidate
	else:
		# if the candidate's term number is higher, than the candidate's log is more up-to-date than us -> vote_granted = yes
		# if the terms are equal, then check the length of the log. 
			# if the candidate's log is longer than ours, then vote_granted = yes
			# if the length of our log is longer, deny the vote 
		if msg['term'] > current_term:
			#print "sending vote yes"
			current_term = msg['term']
			current_leader = msg['leader']
			voted_for = msg['src']
			response = {'src': my_id, 'dst': msg['src'], 'leader': msg['leader'], 'type':'request_vote', 'vote_granted':'yes','term': msg['term']}
			sock.send(json.dumps(response))
		if msg['term'] == current_term and voted_for == None:
			if msg['log_length'] > len(log): 
				current_leader = msg['leader']
				voted_for = msg['src']
				response = {'src': my_id, 'dst': msg['src'], 'leader': msg['leader'], 'type':'request_vote', 'vote_granted':'yes','term': msg['term']}
                        	sock.send(json.dumps(response)) 
			
	
def handle_append_entries_messages(msg):
	global current_term
	global current_leader
	global state
	global my_id
	global num_replica_logs
	global commit_index
	global log
	global follower_next_index

	# If this messsage has a higher term than ours
	# or the same term but a different leader
	# then make sure our information matches with theirs
	# because highest term wins and also only a replica
	# who has gotten the majority of votes and is the real
	# leader can send append entries rpcs
	if msg['term'] > current_term or (msg['term'] == current_term and not msg['leader'] == current_leader):
		current_leader = msg['leader']
		current_term = msg['term']
		
		if not current_leader == my_id:
			state = 'follower'
		else:
			state = 'leader'
		
	# ** Now that we have the correct term and leader sorted
	# ** out we may begin dealing with the information from
	# ** the rpc 
	
	# If it is a message from the leader
	if 'entries' in msg and msg['term'] == current_term and msg['src'] == current_leader:
		# Get the leader's last commit
		leader_commit = msg['commit_index']
		
		# Get our last index and term
		last_index = len(log)
		
		if (len(log) > 0):
			last_term = log[-1]['term']
		else:
			last_term = 0			
		
		# if it is a heartbeat from the leader check if
		# we need to commit anything to state machine 
		if msg['entries'] == '': 	
			check_to_commit(leader_commit)

		# if entries is not empty, check if our last term/last 
		# index match up with the leader's
		elif (msg['last_index'] == last_index and msg['last_term'] == last_term):
			# if we need to add something to the log
			#print "APPEND ENTRY TO FOLLOWER LOG " + str(my_id) + " " + str(msg['last_index'])
					
			# append entry to log
			log = log + msg['entries']
			last_index = len(log)
					
			# send back message of success and what our new last_index is
			response = {'src': my_id, 'dst': msg['src'], 'leader': current_leader, 'type': 'append_entries', 'term': current_term, 'status': 'success', 'last_index': last_index}
			sock.send(json.dumps(response))
			
			# check if we need to commit anything
			# to state machine
			check_to_commit(leader_commit)
			
		# If they do not match up with our log, see if it was a
		# resend, if it is then send success otherwise send failure
		# that includes our last term and index
		else:
			msg_last_index = msg['last_index']
			follower_last_index = len(log)
			
			# first go through the entries the leader sent and
			# see if we already have them b/c this is a resent message
			# check up to our last entry, if more to add do so and send success
			for m in msg['entries']:
				# if we already have this entry, ignore
				if m in log and msg_last_index == log.index(m):
					msg_last_index = msg_last_index + 1
				# if we do not have this entry and it matches up
				# with the indices of our log, append it
				elif not m in log and msg_last_index == follower_last_index:
					log.append(m)
					msg_last_index = msg_last_index + 1
					follower_last_index = follower_last_index + 1
				# if not in log and the last indices do not match up
				elif not m in log and not msg_last_index == follower_last_index:
					# If our last index is greater than or equal to 
					# the last index sent from the leader, remove the
					# unmatching entry/entries
					while len(log) > msg_last_index and len(log) > 0:
						#print "REMOVING ENTRIES FROM FOLLOWER LOG " + str(len(log)) + " " + my_id + " " + str(msg_last_index)
						del log[-1]
				
					response = {'src': my_id, 'dst': msg['src'], 'leader': current_leader, 'type': 'append_entries', 'term': current_term, 'status': 'consistency_failure', 'last_index': last_index, 'last_term':last_term}
					sock.send(json.dumps(response))
					
					return;
					
			#print 'REACHED END OF FOR LOOP'
			# If we reach here, then send success
			last_index = len(log)
			response = {'src': my_id, 'dst': msg['src'], 'leader': current_leader, 'type': 'append_entries', 'term': current_term, 'status': 'success', 'last_index': last_index, 'last_term':last_term}
			sock.send(json.dumps(response))

	# if we are the leader and this is a
	# response to an append_entries we sent
	elif msg['term'] == current_term and state == "leader" and msg['leader'] == my_id:
		# Update the timeout tracker for append entry
		append_entries_timeout_tracker[msg['src']] = -1

		# if the entry we are dealing with now
		# was successfully added to the replica's log
		if msg['status'] == 'success':
			#print "entry successfully added to replica " + str(msg['src']) + " " + str(msg['last_index'])

			follower_last_index = msg['last_index']
			
			# update the next index of the follower
			follower = msg['src']
			follower_next_index[follower] = follower_last_index + 1

			# Check to see if we need to send more append entries
			# messages to the follower to catch them up to our log
			#print str(follower_next_index[follower]) + " " + str(commit_index)
			if follower_next_index[follower] < commit_index:
				#print str(follower_next_index[follower]) + " " + str(commit_index)
				#print "NEED TO SEND MORE TO FOLLOWER"
				send_append_entries_rpc(follower_next_index[follower], follower)
			
			# check if we have already committed this on our state
			# machine or not, if we have then ignore it
			if follower_next_index[follower] - 1 == commit_index + 1:
				# add one to the counter and see if the
				# majority has added it to their logs
				num_replica_logs = num_replica_logs + 1

				# if they have, apply the command to
				# state machine, respond to client
				# and make sure that all other replicas
				# are up to date
				if num_replica_logs >= majority:
					#print "ENTRY ADDED TO MAJORITY OF FOLLOWER LOGS"
					apply_command()
		elif msg['status'] == 'consistency_failure':
			# Get the follower's last index and term
			follower_last_index = msg['last_index']
			follower_last_term = msg['last_term']

			# Get what we have stored for the follower's
			# next index to see if we need to decrement
			follower = msg['src']
			follower_next_entry_index = follower_next_index[follower]

			# Get what we have stored as the follower's last index
			# and last term
			leader_version_follower_last_index = follower_next_entry_index - 1
			
			if leader_version_follower_last_index - 1 <= 0:
				leader_version_follower_last_term = 0
			else:
				leader_version_follower_last_term = log[leader_version_follower_last_index - 1]['term']

			# If what we have as the follower's last term and index is
			# correct, then do not decrement, otherwise decrement the
			# follower's next index by 1
			if not leader_version_follower_last_term == follower_last_term and not leader_version_follower_last_index == follower_last_index:
				follower_next_index[follower] = follower_next_index[follower] - 1
								
			# Send an append entry with the appropriate entry or entries attached
			#print "SENDING APPEND ENTRY TO RESOLVE FAILURE " +  str(follower) + " " + str(follower_next_index[follower])
			send_append_entries_rpc(follower_next_index[follower], follower)
		
# If we are a follower, check to see if we
# should commit anything to our state machine
def check_to_commit(leader_commit_index):
	global log
	global commit_index
	global current_term
	global current_leader
	
	while leader_commit_index > commit_index:
		if len(log) > commit_index and log[commit_index]['term'] == current_term and log[commit_index]['leader'] == current_leader:
			#print "COMMITTING TO FOLLOWER STATE MACHINE"
			apply_command()
		else:
			break;
		
# Apply the command to a replica's state machine
# machine and respond to the client
def apply_command():
	global commit_index			
	global log
	global command_queue
	global num_replica_logs
	global state
	global commit_queue

	#print "APPLYING COMMAND TO STATE MACHINE"
	
	while len(commit_queue) > 0:
		#print "applying old commit"
		command = commit_queue[0]

		if command['type'] == 'get':
			handle_get_command(command)
		else:
			handle_put_command(command)

		del commit_queue[0]

	# Get the correct log entry
	command = log[commit_index]
	
	# Increment the commit index of leader
	commit_index = commit_index + 1
	
	# Check the type of command and handle it accordingly
	if command['type'] == 'get':
		handle_get_command(command)
	elif command['type'] == 'put':
		handle_put_command(command)

	global state_machine
	global my_id
	#print "replica state_machine " + str(my_id) + " " + str(len(state_machine))
	
	# Only need to do the following steps
	# if we are the leader
	if state == "leader":
		# Remove this from the queue
		del command_queue[0]	

		num_replica_logs = 0

		# Handle the next element in the queue
		handle_queue_command()

# Handle when a client sends a
# get request
def handle_get_command(command):
	global state_machine
	global my_id
	global current_leader
	global sock
	global state

	# For gets, only need to do anything if
	# we are the leader
	if state == "leader":
		key = command['key']
		
		# if the key exists, send okay message to client
		if key in state_machine:
			#print "GET SUCCESS"
	
			value = state_machine[key]
			response = {'src': my_id, 'dst': command['client'], 'leader': current_leader, 'type': 'ok', 'MID': command['MID'], 'value': value}
		# otherwise send a message with a blank value
		else:
			response = {'src': my_id, 'dst': command['client'], 'leader': current_leader, 'type': 'ok', 'MID': command['MID'], 'value': ''}
			
		sock.send(json.dumps(response))

# Handle when a client sends a put request
def handle_put_command(command):
	global my_id
	global current_leader
	global state_machine
	global sock
	global state

	key = command['key']
	value = command['value']

	if state == "leader":
		# add to state machine and send success
		if not key == None and not value == None:
			#print "PUT SUCCESS"
			state_machine[key] = value	
			response = {'src': my_id, 'dst': command['client'], 'leader': current_leader, 'type': 'ok', 'MID': command['MID']}
		else:
			response = {'src': my_id, 'dst': command['client'], 'leader': current_leader, 'type': 'fail', 'MID': command['MID']} 
	
		sock.send(json.dumps(response))
	elif state == "follower":
		if not key == None and not value == None:
			state_machine[key] = value
			
# Your ID number
my_id = sys.argv[1]

# The ID numbers of all the other replicas
replica_ids = sys.argv[2:]

# The number that is considered the majority
majority = math.ceil((len(replica_ids) + 1) / 2.0)

# The state the replica is in
# this can be leader, follower or candidate
# all start as followers
state = "follower"

# Each replica has an election timeout
# which will be a random number between
# 150 and 300 ms
election_timeout = get_election_timeout()

# Keep track of when the last RPC from a
# leader was received, initialize to
# the current time
last_received_rpc = time.time()

# Keep track of the last heartbeat we sent
# this is only relevant when we are the leader
# initialize to -1, will be reset if/when we
# become the leader
last_sent_rpc = -1

# Keep track of whether or not the leader
# should resend an append entries to a follower
# based on a timeout
append_entries_timeout_tracker = {}

# Keep track of the replica's term
# all will start at 0
current_term = 0

# Candidate who received this replica's
# vote in the current term
voted_for = None

# If we are in the candidate state, we want
# to keep track of how many votes we have gotten
num_votes = 0

# Current leader, when replicas startup
# leader is unknown so set to 'FFFF'
current_leader = 'FFFF'

# Keep track of a state machine
state_machine = {}

# Keep track of a log, each entry has a 
# term number and the command
log = []

# Keep track of a queue of client commands
command_queue = []

# Keep track of the number of successful
# logs of the current message
num_replica_logs = 0

# When we are the leader, keep track of
# the highest known commit index to let
# replicas know when it is okay to 
# commit a command to their state machines
commit_index = 0

# Keep track of entries in the log that
# were not committed yet and need to be
# when a leader changes
commit_queue = []

# Keep track of a next index for each
# of the other replicas - used when we are
# the leader
follower_next_index = {}

# Connect to the network. All messages to/from other replicas and clients will
# occur over this socket
sock = socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET)
sock.connect(my_id)

last = time.time()

while True:
	ready = select.select([sock], [], [], 0.1)[0]
	
	if sock in ready:
		msg_raw = sock.recv(32768)
		
		if len(msg_raw) == 0: continue
		msg = json.loads(msg_raw)
		
		# If it if a client request, handle it
		if msg['type'] in ['get', 'put']:
        		handle_client_message(msg['type'], msg['MID'], msg['src'], msg)

		# Handle a request_vote rpc
		elif msg['type'] == 'request_vote':
			handle_request_vote_messages(msg)
			last_received_rpc = time.time()

		# Handle an append_entries rpc
		elif msg['type'] == 'append_entries':
			handle_append_entries_messages(msg)
			last_received_rpc = time.time()
		
	# Do periodic checks
	clock = time.time()

	if clock - last >= .1:
		last = time.time()
		
		# Check if election timeout has occurred
		if not (state == "leader"):
			check_election_timeout(clock)
		# Check if we need to send a heartbeat
		# and if any append entry requests have
		# timed out and should be re-sent
		else:
			check_send_heartbeat(clock)
			check_append_entry_timeout(clock)





