from panflute import Doc, Element, Div, Span, Header

from typing import Union, Callable
from types import ModuleType
import os
import warnings

from .command import Command

CommandCallable = Callable[[Command, 'Context'], list[Element]] # This is here because of a wild circular import dependency between many functions and classes

# This class is used to keep state while transforming the document using
# transform.py. For the context to be available to the html and TeX generators,
# individual keys must be manually assigned to the individual elements. This is
# done in transform.py.
#
# The context is also aware of its parent contexts and relevant data (such as
# metadata and commands) can be read from the closest parent context. Writing
# only happens to the current one.
#
# This class is basically an extension to panflute's doc, this is why metadata
# is read directly from it.
def default_number_generator(e: Header, context: 'Context') -> str:
	l = e.level
	context.section_counters[l-1] += 1
	for i in range(l, len(context.section_counters)):
		context.section_counters[i] = 0
	return ".".join(map(str, context.section_counters[:l]))


class Context:
	parent: Union["Context", None]
	_commands: dict[str, Union[CommandCallable, None]]
	doc: Doc
	trusted: bool
	path: str
	dir: str
	filename: str

	section_counters: list[int]
	number_generator: Callable[[Header, 'Context'], str]

	def __init__(self, doc: Doc, path: str, parent: Union['Context', None]=None, trusted: bool=True):
		self.parent = parent
		self._commands = {}
		self.doc = doc
		self.trusted = trusted
		self.path = path
		self.dir = os.path.dirname(path) if os.path.dirname(path) != "" else "."
		self.filename = os.path.basename(path)
		if self.get_metadata("flags", immediate=True) is None:
			self.set_metadata("flags", {})
		self.number_generator = default_number_generator
		self.section_counters = [0 for i in range(6)]

	def get_command(self, command: str) -> Union[CommandCallable, None]:
		if command in self._commands:
			return self._commands[command]
		elif self.parent:
			return self.parent.get_command(command)
		else:
			return None

	def set_command(self, command: str, val: CommandCallable):
		self._commands[command] = val

	def unset_command(self, command: str):
		del self._commands[command]

	def add_commands_from_module(self, module: Union[dict[str, CommandCallable], ModuleType], module_name: str=""):
		if isinstance(module, ModuleType):
			module = module.__dict__
		prefix = module_name+"." if module_name else ""
		for name, func in module.items():
			if isinstance(func, Callable):
				self.set_command(prefix+name, func)

	def is_flag_set(self, flag: str):
		if self.get_metadata("flags."+flag) is not None:
			if self.get_metadata("flags."+flag):
				return True
			else:
				return False
		elif self.parent:
			return self.parent.is_flag_set(flag)
		else:
			return False

	def set_flag(self, flag: str, val: bool):
		self.set_metadata("flags."+flag, val)

	def unset_flag(self, flag: str):
		self.unset_metadata("flags."+flag)

	def get_metadata(self, key: str, simple: bool=True, immediate: bool=False):
		value = self.doc.get_metadata(key, None, simple)
		if value is not None:
			return value
		elif self.parent and not immediate:
			return self.parent.get_metadata(key)
		else:
			return None

	def set_metadata(self, key: str, value):
		meta = self.doc.metadata
		keys = key.split(".")
		for k in keys[:-1]:
			meta = meta[k]
		meta[keys[-1]] = value

	def unset_metadata(self, key: str):
		meta = self.doc.metadata
		keys = key.split(".")
		for k in keys[:-1]:
			meta = meta[k]
		del meta.content[key[-1]] # A hack because MetaMap doesn't have a __delitem__
	
	def import_metadata(self, data, key: str=""):
		if isinstance(data, dict) and isinstance(self.get_metadata(key), dict):
			for subkey, value in enumerate(data):
				self.import_metadata(value, key+"."+subkey if key != "" else subkey)
		else:
			self.set_metadata(key, data)



# This is a custom element which creates \begingroup \endgroup groups in TeX
# and also causes KaTeX math blocks to be isolated in a similar way.
#
# Whenever a new context is created, its content should be eclosed in a group and vice-versa.
class Group(Element):
	metadata: dict
	context: Context
	def __init__(self, *args, context:Context, metadata={}, **kwargs):
		self.metadata = metadata # This is only here for backwards compatibility with old html.py, tex.py and transform.py. FIXME: Remove this when the time comes.
		self.context = context
		super().__init__(*args, **kwargs)

class BlockGroup(Group, Div):
	pass

class InlineGroup(Group, Span):
	pass