import hashlib
import operator
from functools import cached_property
from more_itertools import ilen
from itertools import takewhile
from dataclasses import dataclass

from typing import TYPE_CHECKING, Tuple, Optional
if TYPE_CHECKING: # Self doesn't exist <3.11
	from typing import Self

from cbrrr import encode_dag_cbor, decode_dag_cbor, CID


@dataclass(frozen=True) # frozen == immutable == win
class MSTNode:
	"""
	k/v pairs are interleaved between subtrees like so: ::

		keys:          (0,    1,    2,    3)
		vals:          (0,    1,    2,    3)
		subtrees:   (0,    1,    2,    3,    4)

	If a method is implemented in this class, it's because it's a function/property
	of a single node, as opposed to a whole tree
	"""
	keys: Tuple[str] # collection/rkey
	vals: Tuple[CID] # record CIDs
	subtrees: Tuple[Optional[CID]] # a None value represents an empty subtree


	# NB: __init__ is auto-generated by dataclass decorator

	# these checks could be elided for "trusted" trees
	def __post_init__(self) -> None:
		# TODO: maybe check that they're tuples here?
		# implicitly, the length of self.subtrees must be at least 1
		if len(self.subtrees) != len(self.keys) + 1:
			raise ValueError("Invalid subtree count")
		if len(self.keys) != len(self.vals):
			raise ValueError("Mismatched keys/vals lengths")
		key_heights = [MSTNode.key_height(k) for k in self.keys]
		if any(k!=key_heights[0] for k in key_heights):
			raise ValueError("Inconsistent key heights")

	@classmethod
	def empty_root(cls) -> "Self":
		return cls(
			subtrees=(None,),
			keys=(),
			vals=()
		)

	# this should maybe not be implemented here?
	@staticmethod
	def key_height(key: str) -> int:
		digest = int.from_bytes(hashlib.sha256(key.encode()).digest(), "big")
		leading_zeroes = 256 - digest.bit_length()
		return leading_zeroes // 2

	# since we're immutable, this can be cached
	@cached_property
	def cid(self) -> CID:
		return CID.cidv1_dag_cbor_sha256_32_from(self.serialised)

	# likewise
	@cached_property
	def serialised(self) -> bytes:
		e = []
		prev_key = b""
		for subtree, key_str, value in zip(self.subtrees[1:], self.keys, self.vals):
			key_bytes = key_str.encode()
			shared_prefix_len = ilen(takewhile(bool, map(operator.eq, prev_key, key_bytes))) # I love functional programming
			e.append({
				"k": key_bytes[shared_prefix_len:],
				"p": shared_prefix_len,
				"t": subtree,
				"v": value,
			})
			prev_key = key_bytes
		return encode_dag_cbor({
			"e": e,
			"l": self.subtrees[0]
		})

	@classmethod
	def deserialise(cls, data: bytes) -> "Self":
		cbor = decode_dag_cbor(data)
		if len(cbor) != 2: # e, l
			raise ValueError("malformed MST node")
		subtrees = [cbor["l"]]
		keys = []
		vals = []
		prev_key = b""
		for e in cbor["e"]: # TODO: make extra sure that these checks are watertight wrt non-canonical representations
			if len(e) != 4: # k, p, t, v
				raise ValueError("malformed MST node")
			prefix_len: int = e["p"]
			suffix: bytes = e["k"]
			if prefix_len > len(prev_key):
				raise ValueError("invalid MST key prefix len")
			if prev_key[prefix_len:prefix_len+1] == suffix[:1]:
				raise ValueError("non-optimal MST key prefix len")
			this_key = prev_key[:prefix_len] + suffix
			if this_key <= prev_key:
				raise ValueError("invalid MST key sort order")
			keys.append(this_key.decode())
			vals.append(e["v"])
			subtrees.append(e["t"])
			prev_key = this_key

		return cls(
			subtrees=tuple(subtrees),
			keys=tuple(keys),
			vals=tuple(vals)
		)

	def is_empty(self) -> bool:
		return self.subtrees == (None,)

	def _to_optional(self) -> Optional[CID]:
		"""
		returns None if the node is empty
		"""
		if self.is_empty():
			return None
		return self.cid


	@cached_property
	def maybe_height(self) -> Optional[int]:
		# if there are keys at this level, query one directly
		if self.keys:
			return self.key_height(self.keys[0])

		# we're an empty tree
		if self.subtrees[0] is None:
			return 0

		# this should only happen for non-root nodes with no keys (aka an empty intermediate node)
		return None
		# NOTE: a Node class cannot see what's below it. You'll need to track
		# state externally (like NodeWrangler does) if you want to find out

	def definitely_height(self) -> int:
		if self.maybe_height is None:
			raise ValueError("indeterminate node height")
		return self.maybe_height

	def gte_index(self, key: str) -> int:
		"""
		find the index of the first key greater than or equal to the specified key
		if all keys are smaller, it returns len(keys)
		"""
		i = 0 # this loop could be a binary search but not worth it for small fanouts
		while i < len(self.keys) and key > self.keys[i]:
			i += 1
		return i


"""
if __name__ == "__main__":
	from .blockstore import MemoryBlockStore, OverlayBlockStore
	from .blockstore.car_reader import ReadOnlyCARBlockStore

	if 0:
		import sys
		sys.setrecursionlimit(999999999)
		f = open("/home/david/programming/python/bskyclient/retr0id.car", "rb")
		bs = OverlayBlockStore(MemoryBlockStore(), ReadOnlyCARBlockStore(f))
		commit_obj = dag_cbor.decode(bs.get_block(bytes(bs.lower.car_roots[0])))
		mst_root: CID = commit_obj["data"]
		ns = NodeStore(bs)
		wrangler = NodeWrangler(ns)
		#print(wrangler)
		#enumerate_mst(ns, mst_root)
		enumerate_mst_range(ns, mst_root, "app.bsky.feed.generator/", "app.bsky.feed.generator/\xff")

		root2 = wrangler.del_record(mst_root, "app.bsky.feed.generator/alttext")
		root2 = wrangler.del_record(root2, "app.bsky.feed.like/3kas3fyvkti22")
		root2 = wrangler.put_record(root2, "app.bsky.feed.like/3kc3brpic2z2p", hash_to_cid(b"blah"))

		c, d = mst_diff(ns, mst_root, root2)
		print("CREATED:")
		for x in c:
			print("created", x.encode("base32"))
		print("DELETED:")
		for x in d:
			print("deleted", x.encode("base32"))

		for op in record_diff(ns, c, d):
			print(op)
		
		e, f = very_slow_mst_diff(ns, mst_root, root2)
		assert(e == c)
		assert(f == d)
	else:
		bs = MemoryBlockStore()
		ns = NodeStore(bs)
		wrangler = NodeWrangler(ns)
		root = ns.get_node(None).cid
		print(ns.pretty(root))
		root = wrangler.put_record(root, "hello", hash_to_cid(b"blah"))
		print(ns.pretty(root))
		root = wrangler.put_record(root, "foo", hash_to_cid(b"bar"))
		print(ns.pretty(root))
		root_a = root
		root = wrangler.put_record(root, "bar", hash_to_cid(b"bat"))
		root = wrangler.put_record(root, "xyzz", hash_to_cid(b"bat"))
		root = wrangler.del_record(root, "foo")
		print("=============")
		print(ns.pretty(root_a))
		print("=============")
		print(ns.pretty(root))
		#exit()
		print("=============")
		enumerate_mst(ns, root)
		c, d = mst_diff(ns, root_a, root)
		print("CREATED:")
		for x in c:
			print("created", x.encode("base32"))
		print("DELETED:")
		for x in d:
			print("deleted", x.encode("base32"))
		
		e, f = very_slow_mst_diff(ns, root_a, root)
		assert(e == c)
		assert(f == d)
		
		exit()
		root = wrangler.delete(root, "foo")
		root = wrangler.delete(root, "hello")
		print(ns.pretty(root))
		root = wrangler.delete(root, "bar")
		print(ns.pretty(root))
		root = wrangler.delete(root, "bar")
		print(ns.pretty(root))
"""
