import os
from functools import partial
from typing import Callable, Literal, Optional, Tuple, Union

import flax
import jax
import msgpack
import safetensors.flax
from flax import struct
from flax.serialization import from_bytes, from_state_dict, to_bytes, to_state_dict
from flax.traverse_util import empty_node, flatten_dict, unflatten_dict
from jax import numpy as jnp
from tqdm.auto import tqdm

from fjformer.utils import get_logger

logger = get_logger(__name__)


def process_tensor(
	key,
	shard_fns,
	mismatch_allowed,
	manager,
	callback: Optional[Callable[[jax.Array, str], jax.Array]] = None,
):
	tensor = manager.get_tensor(key)
	istate = 0
	if shard_fns is not None:
		try:
			callable_func = shard_fns.get(key)
			if callable_func is None:
				if not mismatch_allowed:
					raise KeyError(
						f"Shard Function {key} is None and NoneType OBJ is not callable."
					)
				istate = 1
			tensor = callable_func(tensor)
		except KeyError as k_err:
			if mismatch_allowed:
				istate = 1
			else:
				raise KeyError(k_err) from None

	if callback is not None:
		tensor = callback(tensor, key)
	return key, tensor, istate


def is_flatten(pytree: Union[dict, struct.PyTreeNode]) -> bool:
	"""
	Check if the given PyTree is flattened.

	Args:
	    pytree: The PyTree to check.

	Returns:
	    True if the PyTree is flattened, False otherwise.
	"""
	return True if isinstance([k for k in pytree.keys()][0], tuple) else False


def get_dtype(tensor: jax.Array, dtype: Optional[Union[str, jnp.dtype]]) -> jax.Array:
	"""
	Get the tensor with the specified data type.

	Args:
	    tensor: The input tensor.
	    dtype: The desired data type.

	Returns:
	    The tensor with the specified data type.
	"""
	if dtype is None or dtype == "":
		return tensor
	if isinstance(dtype, str):
		dtype = {
			"bf16": jnp.bfloat16,
			"bfloat16": jnp.bfloat16,
			"fp16": jnp.float16,
			"float16": jnp.float16,
			"fp32": jnp.float32,
			"float32": jnp.float32,
			"fp64": jnp.float64,
			"float64": jnp.float64,
		}[dtype]
	float_dtypes = (jnp.bfloat16, jnp.float16, jnp.float32, jnp.float64)
	if getattr(tensor, "dtype", None) in float_dtypes:
		tensor = tensor.astype(dtype)
	return tensor


class CheckpointManager(object):
	"""
	A class to manage saving and loading checkpoints.

	Args:
	    checkpoint_dir: The directory to save checkpoints to.
	    enable: Whether to enable saving and loading checkpoints.
	    float_dtype: The floating-point data type to use for saving checkpoints.
	    save_optimizer_state: Whether to save the optimizer state in the checkpoint.
	    verbose: Whether to print verbose output.
	"""

	def __init__(
		self,
		checkpoint_dir: Union[str, os.PathLike],
		enable: bool = True,
		float_dtype: Union[str, jnp.dtype] = "bf16",
		save_optimizer_state: bool = True,
		verbose: bool = False,
	):
		self.float_dtype = float_dtype
		self.save_optimizer_state = save_optimizer_state
		self.checkpoint_dir = checkpoint_dir
		self.enable = enable
		self.verbose = verbose

	def save_checkpoint(
		self,
		state: struct.PyTreeNode,
		filename: Union[str, os.PathLike],
		gather_fns: Optional[dict[Callable]] = None,
		mismatch_allowed: bool = True,
	):
		"""
		Save a checkpoint to the given filename.

		Args:
		    state: The state dictionary to save.
		    filename: The filename to save the checkpoint to.
		    gather_fns: A dictionary of functions to gather the state before saving.
		    mismatch_allowed: Whether to allow mismatches between the state dictionary and gather functions.
		"""
		if self.enable:
			path = os.path.join(self.checkpoint_dir, filename)
		else:
			path = "/dev/null"
		self.save_state_to_file(
			state, path, gather_fns, self.float_dtype, mismatch_allowed=mismatch_allowed
		)

	@staticmethod
	def load_checkpoint_safe(
		path: Union[str, os.PathLike],
		target: Optional[struct.PyTreeNode] = None,
		shard_fns: Optional[dict[Callable]] = None,
		verbose: bool = False,
		mismatch_allowed: bool = True,
		callback: Optional[Callable[[jax.Array, str], jax.Array]] = None,
	) -> Tuple[Union[struct.PyTreeNode, dict], dict]:
		"""
		Load a checkpoint from the given path.

		Args:
			path: The path to the checkpoint file.
			target: The target PyTree to load the checkpoint into.
			shard_fns: A dictionary of functions to shard the state after loading.
			verbose: Whether to print verbose output.
			mismatch_allowed: Whether to allow mismatches between the state dictionary and shard functions.
		Returns:
			A tuple containing the loaded state dictionary and metadata.
		"""
		with safetensors.safe_open(path, framework="flax") as f:
			metadata = f.metadata()
			keys = list(f.keys())

			if shard_fns is not None:
				if not is_flatten(shard_fns):
					shard_fns = flatten_dict(shard_fns, sep=".")

			process_func = partial(
				process_tensor,
				shard_fns=shard_fns,
				mismatch_allowed=mismatch_allowed,
				manager=f,
				callback=callback,
			)
			results = [
				process_func(key)
				for key in tqdm(
					keys,
					desc="Loading",
					total=len(keys),
				)
			]

		state = {key: tensor for key, tensor, _ in results}
		shard_functions_mismatch = sum(mismatch for _, _, mismatch in results)

		if verbose and shard_functions_mismatch != 0:
			logger.info(f"Sharding mismatch: {shard_functions_mismatch}")

		if target is not None:
			flattened_target = flatten_dict(to_state_dict(target), keep_empty_nodes=True)
			for key, value in flattened_target.items():
				if key not in state and value == empty_node:
					state[key] = value

		state = unflatten_dict(state, sep=".")
		if target is None:
			return state, metadata

		return from_state_dict(target, state), metadata

	@staticmethod
	def save_checkpoint_safe(
		state: struct.PyTreeNode,
		path: Union[str, os.PathLike],
		gather_fns: Optional[dict[Callable]] = None,
		float_dtype: Optional[Union[str, jnp.dtype]] = None,
		verbose: bool = True,
		mismatch_allowed: bool = True,
		metadata: Optional[dict[str, str]] = None,
	):
		"""
		Save a checkpoint to the given path using SafeTensors.

		Args:
		    state: The state dictionary to save.
		    path: The path to the checkpoint file.
		    gather_fns: A dictionary of functions to gather the state before saving.
		    float_dtype: The floating-point data type to use for saving the checkpoint.
		    verbose: Whether to print verbose output.
		    mismatch_allowed: Whether to allow mismatches between the state dictionary and gather functions.
		    metadata: Additional metadata to store in the checkpoint.
		"""
		state = to_state_dict(state)
		gather_functions_mismatch = 0
		if is_flatten(state):
			state = unflatten_dict(state)

		if gather_fns is not None:
			# Example:
			# gather_fns = {"params": {"Dense_0": lambda x: x[0]}}

			if not is_flatten(gather_fns):
				gather_fns = flatten_dict(gather_fns)
			state = flatten_dict(state)
			pbar_gather = tqdm(
				list(state.keys()), desc="Gathering State", disable=not verbose
			)
			for key in pbar_gather:
				try:
					callable_func = gather_fns[key]
					if callable_func is None and not mismatch_allowed:
						raise KeyError(
							f"Gather Function {key} is None and NoneType OBJ is not callable."
						)
					if callable_func is None:
						gather_functions_mismatch += 1
					else:
						state[key] = callable_func(state[key])
				except KeyError as e:
					if mismatch_allowed:
						pbar_gather.set_postfix(gather_mismatch=gather_functions_mismatch)
					else:
						raise KeyError(e) from None
				pbar_gather.update(1)
		state = flax.traverse_util.flatten_dict(state, sep=".")
		for key in list(state.keys()):
			if not isinstance(state[key], jax.Array):
				state[key] = jnp.array(state[key])
			state[key] = get_dtype(state[key], float_dtype)

		safetensors.flax.save_file(tensors=state, filename=path, metadata=metadata)

	@staticmethod
	def save_state_to_file(
		state: struct.PyTreeNode,
		path: Union[str, os.PathLike],
		gather_fns: Optional[dict[Callable]] = None,
		float_dtype: Optional[Union[str, jnp.dtype]] = None,
		verbose: bool = False,
		mismatch_allowed: bool = True,
	):
		"""
		Save the state dictionary to a file.

		Args:
		    state: The state dictionary to save.
		    path: The path to the file to save the state dictionary to.
		    gather_fns: A dictionary of functions to gather the state before saving.
		    float_dtype: The floating-point data type to use for saving the state dictionary.
		    verbose: Whether to print verbose output.
		    mismatch_allowed: Whether to allow mismatches between the state dictionary and gather functions.
		"""
		state = to_state_dict(state)
		packer = msgpack.Packer()
		flatten_state = flatten_dict(state)
		if gather_fns is not None:
			gather_fns = flatten_dict(to_state_dict(gather_fns))
		pbar = tqdm(
			flatten_state.items(),
			disable=not verbose,
			desc="Saving State to File",
		)

		gather_functions_mismatch = 0

		with open(path, "wb") as stream:
			for key, value in pbar:
				if gather_fns is not None:
					try:
						callable_func = gather_fns[key]
						if callable_func is None and not mismatch_allowed:
							raise KeyError(
								f"Gather Function {key} is None and NoneType OBJ is not callable."
							)
						value = callable_func(value) if callable_func is not None else value
						if callable_func is None:
							gather_functions_mismatch += 1
					except KeyError as k_err:
						if mismatch_allowed:
							gather_functions_mismatch += 1
						else:
							raise KeyError(k_err) from None
				pbar.set_postfix(gather_functions_mismatch=gather_functions_mismatch)
				value = get_dtype(value, float_dtype)
				stream.write(packer.pack((key, to_bytes(value))))

	def save_pickle(self, obj: object, filename: Union[str, os.PathLike]):
		"""
		Save an object to a pickle file.

		Args:
		    obj: The object to save.
		    filename: The filename to save the object to.
		"""
		import pickle

		def save_pickle(obj_, path_):
			with open(path_, "wb") as stream:
				pickle.dump(obj_, stream)

		if self.enable:
			path = os.path.join(self.checkpoint_dir, filename)
		else:
			path = "/dev/null"
		save_pickle(obj, path)

	def save_all(
		self,
		state: struct.PyTreeNode,
		gather_fns: dict[Callable],
		metadata: Optional[dict] = None,
		dataset: Optional[object] = None,
		milestone: bool = False,
	):
		"""
		Save all components of a checkpoint.

		Args:
		    state: The state dictionary to save.
		    gather_fns: A dictionary of functions to gather the state before saving.
		    metadata: Additional metadata to save.
		    dataset: The dataset to save.
		    milestone: Whether to save a milestone checkpoint.
		"""
		step = int(jax.device_get(state.step))
		if self.save_optimizer_state:
			checkpoint_state = state
			checkpoint_name = "streaming_state"
			checkpoint_gather_fns = gather_fns
		else:
			checkpoint_state = state.params["params"]
			checkpoint_name = "streaming_params"
			checkpoint_gather_fns = gather_fns.params["params"]

		if milestone:
			# Save a milestone checkpoint that will not be overwritten
			self.save_pickle(metadata, f"metadata_{step}.pkl")
			self.save_pickle(dataset, f"dataset_{step}.pkl")
			self.save_checkpoint(
				checkpoint_state, f"{checkpoint_name}_{step}", checkpoint_gather_fns
			)
		else:
			# Save a normal checkpoint that can be overwritten
			self.save_pickle(metadata, "metadata.pkl")
			self.save_pickle(dataset, "dataset.pkl")
			self.save_checkpoint(
				checkpoint_state, f"{checkpoint_name}", checkpoint_gather_fns
			)

	@staticmethod
	def load_checkpoint(
		path: Union[str, os.PathLike],
		target: Optional[struct.PyTreeNode] = None,
		shard_fns: Optional[dict[Callable]] = None,
		remove_dict_prefix: Optional[Tuple[str, ...]] = None,
		verbose: bool = False,
		mismatch_allowed: bool = True,
		callback: Optional[Callable[[jax.Array, str], jax.Array]] = None,
	) -> Union[struct.PyTreeNode, dict]:
		"""
		Load a checkpoint from the given path.

		Args:
			path: The path to the checkpoint file.
			target: The target PyTree to load the checkpoint into.
			shard_fns: A dictionary of functions to shard the state after loading.
			remove_dict_prefix: A tuple of strings representing the prefix to remove from the state dictionary keys.
			verbose: Whether to print verbose output.
			mismatch_allowed: Whether to allow mismatches between the state dictionary and shard functions.
		Returns:
			The loaded state dictionary.
		"""
		if shard_fns is not None:
			shard_fns = flatten_dict(to_state_dict(shard_fns))
		if remove_dict_prefix is not None:
			remove_dict_prefix = tuple(remove_dict_prefix)
		flatten_state = {}

		shard_functions_mismatch = 0
		with open(path, "rb") as fin:
			unpacker = msgpack.Unpacker(fin, read_size=83886080, max_buffer_size=0)
			pbar = tqdm(unpacker, disable=not verbose, desc="Loading Checkpoints From File")
			for key, value in pbar:
				key = tuple(key)
				if remove_dict_prefix is not None:
					if key[: len(remove_dict_prefix)] == remove_dict_prefix:
						key = key[len(remove_dict_prefix) :]
					else:
						continue

				tensor = from_bytes(None, value)
				if shard_fns is not None:
					try:
						callable_func = shard_fns.get(key, None)
						if callable_func is None and not mismatch_allowed:
							raise KeyError(
								f"Shard Function {key} is None and NoneType OBJ is not callable."
							)
						tensor = callable_func(tensor) if callable_func is not None else tensor
						if callable_func is None:
							shard_functions_mismatch += 1
					except KeyError as k_err:
						if mismatch_allowed:
							shard_functions_mismatch += 1
						else:
							raise KeyError(k_err) from None
				flatten_state[key] = tensor
				tensor = callback(tensor, key)
				pbar.set_postfix(shard_functions_mismatch=shard_functions_mismatch)
		if target is not None:
			flattened_target = flatten_dict(to_state_dict(target), keep_empty_nodes=True)
			for key, value in flattened_target.items():
				if key not in flatten_state and value == empty_node:
					flatten_state[key] = value

		state = unflatten_dict(flatten_state)
		if target is None:
			return state

		return from_state_dict(target, state)

	@staticmethod
	def load_flax_checkpoint(
		path: Union[str, os.PathLike],
		target: Optional[struct.PyTreeNode] = None,
		shard_fns: Optional[dict[Callable]] = None,
	) -> Union[struct.PyTreeNode, dict]:
		"""
		Load a standard flax checkpoint that's not saved with the
		msgpack streaming format.

		Args:
		    path: The path to the checkpoint file.
		    target: The target PyTree to load the checkpoint into.
		    shard_fns: A dictionary of functions to shard the state after loading.

		Returns:
		    The loaded state dictionary.
		"""
		with open(path, "rb") as fin:
			encoded_bytes = fin.read()

		state_dict = flax.serialization.msgpack_restore(encoded_bytes)
		if shard_fns is not None:
			shard_fns = to_state_dict(shard_fns)
			state_dict = jax.tree_util.tree_map(lambda fn, x: fn(x), shard_fns, state_dict)

		if target is None:
			return state_dict
		return from_state_dict(target, state_dict)

	@classmethod
	def load_state_checkpoint(
		cls,
		load_type: Literal["state", "state_params", "params", "flax_params"],
		load_path: Union[str, os.PathLike],
		state_target: Optional[struct.PyTreeNode] = None,
		state_shard_fns: Optional[dict[Callable]] = None,
		disallow_state: bool = False,
		mismatch_allowed: bool = True,
	) -> Tuple[Optional[struct.PyTreeNode], Optional[flax.core.frozen_dict.FrozenDict]]:
		"""
		Load a state checkpoint from the given path.

		Args:
		    load_type: The type of checkpoint to load.
		    load_path: The path to the checkpoint file.
		    state_target: The target PyTree to load the state into.
		    state_shard_fns: A dictionary of functions to shard the state after loading.
		    disallow_state: Whether to disallow loading the full state.
		    mismatch_allowed: Whether to allow mismatches between the state dictionary and shard functions.

		Returns:
		    A tuple containing the loaded state and parameters.
		"""
		if state_target is not None:
			params_target = state_target.params["params"]
		else:
			params_target = None

		if state_shard_fns is not None:
			params_shard_fns = state_shard_fns.params["params"]
		else:
			params_shard_fns = None

		if disallow_state:
			assert load_type != "state", "Loading full state is not allowed!"
		state = None
		restored_params = None
		if load_type == "state":
			state = cls.load_checkpoint(
				path=load_path,
				target=state_target,
				shard_fns=state_shard_fns,
				mismatch_allowed=mismatch_allowed,
			)
		elif load_type == "state_params":
			restored_params = cls.load_checkpoint(
				path=load_path,
				target=params_target,
				shard_fns=params_shard_fns,
				remove_dict_prefix=("params", "params"),
				mismatch_allowed=mismatch_allowed,
			)
			restored_params = flax.core.frozen_dict.freeze({"params": restored_params})
		elif load_type == "params":
			restored_params = cls.load_checkpoint(
				path=load_path,
				target=params_target,
				shard_fns=params_shard_fns,
				mismatch_allowed=mismatch_allowed,
			)
			restored_params = flax.core.frozen_dict.freeze({"params": restored_params})
		elif load_type == "flax_params":
			restored_params = cls.load_flax_checkpoint(
				path=load_path,
				target=params_target,
				shard_fns=params_shard_fns,
			)
			restored_params = flax.core.frozen_dict.freeze({"params": restored_params})
		else:
			raise ValueError(f"Invalid load_from type: {load_type}")

		return state, restored_params
