import heapq
import time
import operator
import asyncio

from itertools import chain
from collections import OrderedDict
from liaa.utils import shared_prefix, bytes_to_bit_string


class KBucket:
	def __init__(self, range_lower, range_upper, ksize):
		self.range = (range_lower, range_upper)
		self.nodes = OrderedDict()
		self.replacement_nodes = OrderedDict()
		self.touch_last_updated()
		self.ksize = ksize

	def touch_last_updated(self):
		self.last_updated = time.monotonic()

	def get_nodes(self):
		return list(self.nodes.values())

	def split(self):
		midpoint = (self.range[0] + self.range[1]) / 2
		one = KBucket(self.range[0], midpoint, self.ksize)
		two = KBucket(midpoint + 1, self.range[1], self.ksize)
		nodes = chain(self.nodes.values(), self.replacement_nodes.values())
		for node in nodes:
			bucket = one if node.long_id <= midpoint else two
			bucket.add_node(node)

		return (one, two)

	def remove_node(self, node):
		if node.digest_id in self.replacement_nodes:
			del self.replacement_nodes[node.digest_id]

		if node.digest_id in self.nodes:
			del self.nodes[node.digest_id]

			if self.replacement_nodes:
				newnode_id, newnode = self.replacement_nodes.popitem()
				self.nodes[newnode_id] = newnode

	def has_in_range(self, node):
		return self.range[0] <= node.long_id <= self.range[1]

	def is_new_node(self, node):
		return node.digest_id not in self.nodes

	def add_node(self, node):
		"""
		Add a C{Node} to the C{KBucket}.  Return True if successful,
		False if the bucket is full.

		If the bucket is full, keep track of node in a replacement list,
		per section 4.1 of the paper.
		"""
		if node.digest_id in self.nodes:
			del self.nodes[node.digest_id]
			self.nodes[node.digest_id] = node
		elif len(self) < self.ksize:
			self.nodes[node.digest_id] = node
		else:
			if node.digest_id in self.replacement_nodes:
				del self.replacement_nodes[node.digest_id]
			self.replacement_nodes[node.digest_id] = node
			return False
		return True

	def depth(self):
		vals = self.nodes.values()
		sprefix = shared_prefix([bytes_to_bit_string(n.digest_id) for n in vals])
		return len(sprefix)

	def head(self):
		return list(self.nodes.values())[0]

	def __getitem__(self, node_id):
		return self.nodes.get(node_id, None)

	def __len__(self):
		return len(self.nodes)


class TableTraverser:
	def __init__(self, table, start_node):
		index = table.get_bucket_index_for(start_node)
		table.buckets[index].touch_last_updated()
		self.current_nodes = table.buckets[index].get_nodes()
		self.left_buckets = table.buckets[:index]
		self.right_buckets = table.buckets[(index + 1):]
		self.left = True

	def __iter__(self):
		return self

	def __next__(self):
		"""
		Pop an item from the left subtree, then right, then left, etc.
		"""
		if self.current_nodes:
			return self.current_nodes.pop()

		if self.left and self.left_buckets:
			self.current_nodes = self.left_buckets.pop().get_nodes()
			self.left = False
			return next(self)

		if self.right_buckets:
			self.current_nodes = self.right_buckets.pop(0).get_nodes()
			self.left = True
			return next(self)

		raise StopIteration


class RoutingTable:
	def __init__(self, protocol, ksize, node):
		"""
		@param node: The node that represents this server.  It won't
		be added to the routing table, but will be needed later to
		determine which buckets to split or not.
		"""
		self.node = node
		self.protocol = protocol
		self.ksize = ksize
		self.flush()

	def flush(self):
		self.buckets = [KBucket(0, 2 ** 160, self.ksize)]

	def split_bucket(self, index):
		one, two = self.buckets[index].split()
		self.buckets[index] = one
		self.buckets.insert(index + 1, two)

	def lonely_buckets(self):
		"""
		Get all of the buckets that haven't been updated in over
		an hour.
		"""
		hrago = time.monotonic() - 3600
		return [b for b in self.buckets if b.last_updated < hrago]

	def remove_contact(self, node):
		index = self.get_bucket_index_for(node)
		self.buckets[index].remove_node(node)

	def is_new_node(self, node):
		index = self.get_bucket_index_for(node)
		return self.buckets[index].is_new_node(node)

	def add_bucket(self, bucket):
		self.buckets.append(bucket)

	def add_contact(self, node):
		index = self.get_bucket_index_for(node)
		bucket = self.buckets[index]

		# this will succeed unless the bucket is full
		if bucket.add_node(node):
			return

		# Per section 4.2 of paper, split if the bucket has the node
		# in its range or if the depth is not congruent to 0 mod 5
		if bucket.has_in_range(self.node) or bucket.depth() % 5:
			self.split_bucket(index)
			self.add_contact(node)
		else:
			asyncio.ensure_future(self.protocol.call_ping(bucket.head()))

	def get_bucket_index_for(self, node):
		"""
		Get the index of the bucket that the given node would fall into.
		"""
		index = None
		for i, bucket in enumerate(self.buckets):
			if node.long_id < bucket.range[1]:
				index = i
				break
		return index

	def get_bucket_for(self, node):
		idx = self.get_bucket_index_for(node)
		return self.buckets[idx]

	def find_neighbors(self, node, k=None, exclude=None):
		k = k or self.ksize
		nodes = []
		for neighbor in TableTraverser(self, node):
			notexcluded = exclude is None or not neighbor.is_same_node(exclude)
			if neighbor.digest_id != node.digest_id and notexcluded:
				heapq.heappush(nodes, (node.distance_to(neighbor), neighbor))
			if len(nodes) == k:
				break

		return list(map(operator.itemgetter(1), heapq.nsmallest(k, nodes)))
