import subprocess
import sys
import ast
import importlib.util
import pathlib
import contextlib
import os
import torch

# Check for wrapt
wrapt = None
try:
	wrapt = __import__("wrapt")
except:
	subprocess.Popen([sys.executable, "-m", "pip", "install", "wrapt"]).wait()
	wrapt = __import__("wrapt")

BLACKLIST = [
	"_way_in",
	"_way_out",
	"_junc_in",
	"_junc_out",
	"..."
]

HEAD_LOG = "[ComfyUI-0️⃣ 2️⃣ 4️⃣ 6️⃣ ] "

######################################################################################
######################################## HACK ########################################
######################################################################################

class TautologyStr(str):
	def __ne__(self, other):
		return False

class ByPassTypeTuple(tuple):
	def __getitem__(self, index):
		if index > 0:
			index = 0
		item = super().__getitem__(index)
		if isinstance(item, str):
			return TautologyStr(item)
		return item

class TautologyRest(dict):
	def __iter__(self):
		yield False
		while True:
			yield True

	def __getitem__(self, index):
		if index == 0:
			return False
		return True

class TautologyAll(dict):
	def __iter__(self):
		while True:
			yield True

	def __getitem__(self, index):
		return True

######################################################################################
######################################## UTIL ########################################
######################################################################################

@contextlib.contextmanager
def temp_dir(new_dir):
	old_dir = os.getcwd()
	os.chdir(new_dir)
	try:
		yield
	finally:
		os.chdir(old_dir)

def import_module(module_name, file_path, package_name = None):
	try:
		if module_name in sys.modules:
			return sys.modules[module_name]
		else:
			spec = importlib.util.spec_from_file_location(module_name, pathlib.Path(file_path))
			module = importlib.util.module_from_spec(spec)
			if package_name is not None:
				module.__package__ = package_name
			sys.modules[module_name] = module
			spec.loader.exec_module(module)
			return module
	except (FileNotFoundError, ModuleNotFoundError, OSError):
		return None

class RevisionDict(dict):
	def __init__(self, *args, **kwargs):
		self.update(*args, **kwargs)

	def path_count(self, path):
		count = 0
		for key in self.keys():
			if key[0:len(path)] == path:
				count += 1
		return count
	
	def path_exists(self, path):
		for key in self.keys():
			if key[0:len(path)] == path:
				return True
		return False
	
	def path_iter(self, path):
		for key in self.keys():
			if key[0:len(path)] == path:
				yield key

	def path_keys(self, path):
		res = []
		for key in self.path_iter(path):
			res.append(key[len(path):])
		return res

class FlatIter:
	def __init__(self, data):
		self.data = data
		self.keys_list = list(data.keys())
		self.key_indices = {key: 0 for key in self.keys_list}
		self.index = 0
		self.elem_total = sum(len(v) for v in data.values())

	def __iter__(self):
		# Reset the indices
		self.key_indices = {key: 0 for key in self.keys_list}
		self.index = 0
		
		while self.index < self.elem_total:
			next_key = None
			next_key_index = None

			# Instead of finding the smallest value, find the next index to be incremented
			for key in self.keys_list:
				if self.key_indices[key] < len(self.data[key]):
					if next_key is None or self.key_indices[key] < next_key_index:
						next_key = key
						next_key_index = self.key_indices[key]

			if next_key is None:
				break

			# Yield the next item from the current key
			current_index = self.key_indices[next_key]
			yield next_key, self.data[next_key][current_index]

			# Increment the index for the next key and the total index
			self.key_indices[next_key] += 1
			self.index += 1

	def __getitem__(self, index):
		# Iterate to get to the item at the specified index
		for i, item in enumerate(self):
			if i == index:
				return item
		raise IndexError("Index out of range")

def flat_zip(names, flat_iter: FlatIter):
	keys_list = list(flat_iter.data.keys())
	name_idx = 0

	for key, value in flat_iter:
		key_idx = keys_list.index(key)
		if key_idx != name_idx:
			name_idx = key_idx
		current_name = names[name_idx]
		yield current_name, (key, value)

def hijack(scope, name, param_func, res_func = None, call_func = None, out_scope = None):
	old_func = getattr(scope, name)
	def new_func(*args, **kwargs):
		new_args, new_kwargs = param_func(*args, **kwargs)
		final_args = tuple(new_args[i] if i < len(new_args) else args[i] for i in range(len(args)))
		for key in new_kwargs:
			kwargs[key] = new_kwargs[key]
		res_value = None
		if call_func is None:
			res_value = old_func(*final_args, **kwargs)
		else:
			res_value = call_func(old_func, *final_args, **kwargs)
		return res_value if res_func is None else res_func(res_value, *final_args, **kwargs)
	setattr(scope if out_scope is None else out_scope, name, new_func)

WRAP_DATA = "_self_0246_data_dict"

class Wrapper(wrapt.ObjectProxy):
	def __init__(self, wrapped):
		super().__init__(wrapped)
		setattr(self, WRAP_DATA, {})

def check_update(data):
	if isinstance(data, list):
		data = data[0]
	if isinstance(data, str):
		data = ast.literal_eval(data) # [TODO] Maybe properly handle this in the future
	return data["update"]

def at_idx(head, rest, index):
	full_length = len(rest) + 1
	norm_index = index % full_length
	if norm_index == 0:
		return head
	return rest[norm_index - 1]

def beautify_structure(data, indent=0, mode=0, stop=False):
	"""
	Recursively format the structure of a dictionary, list, or tuple without formatting the actual data.
	Now it format the structure for each element within lists and tuples.
	"""
	# Determine the indentation string
	indent_str = '  ' * indent
	res_str = ""

	if isinstance(data, dict):
		res_str += f"{indent_str}Dict of {len(data)}:\n"
		for key, value in data.items():
			res_str += f"{indent_str}  Key ({type(key).__name__}): {key}\n"
			res_str += beautify_structure(value, indent + 1, mode)
	elif isinstance(data, list):
		res_str += f"{indent_str}List of {len(data)}:\n"
		for item in data:
			res_str += beautify_structure(item, indent + 1, mode)
	elif isinstance(data, tuple):
		res_str += f"{indent_str}Tuple of {len(data)}:\n"
		for item in data:
			res_str += beautify_structure(item, indent + 1, mode)
	elif isinstance(data, set):
		res_str += f"{indent_str}Set of {len(data)}:\n"
		for item in data:
			res_str += beautify_structure(item, indent + 1, mode)
	else:
		res_str += f"{indent_str}Type: {type(data).__name__}\n"
		# Attempt to loop through each attribute
		match mode:
			case 1:
				try:
					if isinstance(data, str) or isinstance(data, torch.Tensor):
						raise TypeError("stub")
					iterator = iter(data) # Try to get an iterator from data.
					if hasattr(data, '__len__'):
						res_str += f"{indent_str}Iterable of {len(data)}:\n"
					else:
						res_str += f"{indent_str}Iterable:\n"
					for item in iterator:
						res_str += beautify_structure(item, indent + 1, mode)
				except TypeError:
					# If data is not iterable, just print its repr.
					res_str += f"{indent_str}  Data: {repr(data)}\n"
			case 2:
				try:
					if isinstance(data, (int, float, str, bool, type(None))):
						res_str += f"{indent_str}  Data: {repr(data)}\n"
					else:
						for attr in dir(data):
							attr_value = getattr(data, attr)
							if attr.startswith('_') or callable(attr_value):
								continue

							if isinstance(attr_value, (int, float, str, bool, type(None))):
								res_str += f"{indent_str}  Data: {repr(data)}\n"
							elif isinstance(attr_value, (dict, list, tuple, set)):
								res_str += f"{indent_str}  Attribute: {attr}\n"
								res_str += beautify_structure(attr_value, indent + 2, mode, True)
							elif stop:
								res_str += f"{indent_str}  {attr} ({type(attr_value).__name__}): {attr_value}\n"
							else:
								res_str += f"{indent_str}  Attribute: {attr}\n"
								res_str += beautify_structure(attr_value, indent + 2, mode, True)
				except Exception as e:
					pass
			case _:
				pass

	return res_str

######################################################################################
######################################## LANG ########################################
######################################################################################

# Special thanks to ChatGPT for this
HIGHWAY_OPS = {'>': 'set', '<': 'get', '!': 'eat'}

def parse_query(input, ops):
	# States
	inside_backticks = False
	operation = None
	name = []
	result = {value: [] for value in ops.values()}  # Initialize result with values from ops
	errors = []
	order = []  # To keep track of the order of operations
	i = 0  # Current index in the input string

	# Helper function to add a command to the result and order
	def add_command(op, name_str):
		if op and name_str:  # Only add if both operation and name are present
			result[op].append(name_str)
			order.append((op, name_str))

	# Iterate over the input string
	while i < len(input):
		char = input[i]

		# Handle operation characters
		if char in ops:
			if inside_backticks:
				name.append(char)
			elif operation is not None:
				errors.append(f"Error at char {i+1}: Multiple operation symbols")
				break
			else:
				operation = ops[char]

		# Handle backticks
		elif char == '`':
			inside_backticks = not inside_backticks
			if not inside_backticks and operation:  # Closing backtick
				if name:
					add_command(operation, ''.join(name))
				name = []
				operation = None

		# Handle valid name characters
		elif char.isalnum() or char in ['-', '_'] or (inside_backticks and char):
			if operation is not None or inside_backticks:
				name.append(char)
			else:
				errors.append(f"Error at char {i+1}: Operation symbol expected before name")
				break

		# Handle spaces
		elif char.isspace():
			if inside_backticks:
				name.append(char)
			elif name:
				add_command(operation, ''.join(name))
				name = []
				operation = None

		# Handle semicolons
		elif char == ';':
			if inside_backticks:
				name.append(char)
			elif name:  # Semicolon outside of backticks ends the current command
				add_command(operation, ''.join(name))
				name = []
				operation = None

		# Handle any other character that is not whitespace (error case)
		elif not char.isspace():
			errors.append(f"Error at char {i+1}: Invalid character '{char}'")
			break

		i += 1  # Move to the next character

	# Check if we're in a valid state after parsing all characters
	if inside_backticks:
		errors.append("Error: Unclosed backticks")

	# If there's an unfinished operation at the end
	if operation and name:
		add_command(operation, ''.join(name))

	# If there's an operation symbol but no name and no other errors
	if operation and not name and not errors:
		errors.append(f"Error at char {i + 1}: Operation '{operation}' without a name")

	# Return the result, any errors, and the order of operations
	return (result, order, errors)

def highway_check(result, errors):
	# Check if duplicate names exist within results
	exists = set()
	for name in result['set']:
		if name in exists:
			errors.append(f"Error: Duplicate input name '{name}'")
		else:
			exists.add(name)

	exists = set()
	for name in result['get'] + result['eat']:
		if name in exists:
			errors.append(f"Error: Duplicate output name '{name}'")
		else:
			exists.add(name)

def check_used(_way_in, elem):
	if elem[1] in _way_in["used"]:
		raise Exception(f"Output \"{elem[1]}\" is already used.")

######################################################################################

def parse_offset(input):
	# Split the string by semicolons
	segments = input.split(';')
	
	# Initialize an empty list to store the parsed data
	parsed_data = []
	
	# Iterate over each segment
	for segment in segments:
		# Trim whitespace and check if the segment is not empty
		segment = segment.strip()
		if segment:
			# Split the segment by comma
			parts = segment.split(',')
			
			# Check if there are exactly two parts after splitting by comma
			if len(parts) != 2:
				return (None, f"Segment '{segment}' is invalid: expected a pair separated by a single comma.")
			
			# Trim whitespace from the string part
			string_part = parts[0].strip()
			
			# Ensure the string part is not empty
			if not string_part:
				return (None, f"Segment '{segment}' is invalid: string part is empty.")
			
			# Concatenate the number part to remove spaces and keep the operator
			number_part = ''.join(parts[1].split())
			try:
				# Check for multiple operators or incorrect placement
				if number_part.count('+') + number_part.count('-') > 1 or not number_part.lstrip('+-').isdigit():
					return (None, f"Segment '{segment}' is invalid: number part has multiple operators or incorrect placement.")

				# Check if the substring is an integer
				int(number_part.lstrip('+-'))
			except ValueError:
				return (None, f"Segment '{segment}' is invalid: number part is not an integer.")
			
			# Add the tuple to the list
			parsed_data.append((string_part, number_part))
	
	return (parsed_data, None)

######################################################################################

"""
# Assume we are dealing with KSampler
[0, 1, 0, 1, 0, 1] = model;
	# Input must be a batch list
	# Unnamed, which is depends on order of input pin position (currently 0)
	[0(2), 1(3)] = model; # [0, 0, 1, 1, 1]
	[0(2), 1(3), @] = model; # [0, 0, 1, 1, 1, 0, 0, 1, 1, 1, 0, etc.]
	[0(@),] = model; # [0, 0, 0, 0, etc.]
		[0(@)][1(@)] = model; # Basically always pin to arr[0][1]
	[-1] = model; # Can have negative index for last position, otherwise behave same as positive index
pos = `positive`; # Input and Output can be escaped just in case
`neg` = 2; # negative, since negative is on input pin 2 count from 0
	`neg` = !8; # negative, but use internal python func param index instead
	`neg` = !-3; # Index can be negative number
	`neg` = !negative; # Use internal python func param name
latent[0] = latent_image;
	# Same syntax within [] like model
	# Same requirement of input must be a batch list
	# Named
(1.0) = denoise; # Basic support for int, string, float and combo
(`normal`) = scheduler; # Same thing here
sampler_name; # sampler_name = sampler_name
!sampler_name; # !sampler_name = !sampler_name
6 = cfg;
	6[0] = cfg; # Optional [] with same syntax as model

[] = steps;
	# Basically act like steps = steps
	# Differ from "<num> = <name>" as this rely on input in position (currently 8)

# If pin not provided, the value will default to whatever INPUT_TYPES gives (for steps, it will be 20)
# Also support basic single line comment

@@;
	# If there only exist this within the entire code, then Highway will map param to pin name one by one
!@@;
	# Same thing here but highway key map to func param

"""

"""
# Full syntax
[0, 1, 0, 1, 0, 1] = model; # Basic inline comment
# More inline comment
[0, 1, 0, 1, 0, 1,] = model;
	# Allows indent
	[0(2), 1(3)] = model;
	[0(2), 1(3), @] = model;
	[0(@)] = model;
	[0(@),] = model;
	[-1] = model;
pos = "positive";
"neg" = 2;
neg = !8;
"neg" = !-3;
"neg" = !negative;
latent[0] = latent_image;
"latent"[0(2), 1(3)] = latent_image;
(1.0) = denoise;
("normal") = scheduler;
sampler_name;
!sampler_name = sampler;
6 = cfg;
!6 = cfg;
6[0] = cfg;
[] = steps;
"""

"""
# Execute specific node
@`class_type` {
	<param-code>
}

# Execute from group
@[`group_name`] {
	<param-code>
}

# Execute from .json graph
@(`path.json`) {
	<param-code>
}

# Execute graph from other tabs
@{`tab_name`} {
	<param-code>
}

# Output a Highway
@... ${...}
@... $`Highway`{...}

# Output a Junction
@... $${...}
@... $`Junction`{...}

# Output a Junction Batch
@... $$${...}
@... $`JunctionBatch`{...}
"""

def trim_space(line):
	while len(line) > 0 and line[0].isspace():
		line = line[1:]
	
	return line

def parse_comment(input, state):
	if input[state["idx"]] != '#':
		return True
	
	state["idx"] += 1
	temp = ""

	while state["idx"] < state["len"]:
		token = input[state["idx"]]
		if token == '\n':
			temp = trim_space(temp)
			state["res"].append({
				"type": "#",
				"value": temp
			})
			state["idx"] += 1
			return
		temp += token
		state["idx"] += 1

def parse_space(input, state):
	if not input[state["idx"]].isspace():
		return True

	while state["idx"] < state["len"]:
		if not input[state["idx"]].isspace():
			break
		state["idx"] += 1

def parse_name(input, state):
	if not input[state["idx"]].isalpha():
		return True

	temp = ""
	while state["idx"] < state["len"]:
		token = input[state["idx"]]
		if token.isalnum() or token == '_':
			temp += token
		else:
			break
		state["idx"] += 1

	state["res"].append({
		"type": "name",
		"value": temp
	})

def parse_string(input, state):
	if input[state["idx"]] != '`':
		return True

	state["idx"] += 1

	temp = ""
	while state["idx"] < state["len"]:
		token = input[state["idx"]]
		if token == '`':
			state["res"].append({
				"type": "name",
				"value": temp
			})
			state["idx"] += 1
			return
		temp += token
		state["idx"] += 1

	state["err"] = f"Error at char {state['idx']}: Unterminated \"`\""

def parse_number(input, state):
	if not (input[state["idx"]].isdigit() or input[state["idx"]] == '-'):
		return True

	temp = ""
	if input[state["idx"]] == '-':
		temp += '-'
		state["idx"] += 1

	while state["idx"] < state["len"]:
		token = input[state["idx"]]
		if token.isdigit() or token == '.':
			temp += token
		else:
			break
		state["idx"] += 1

	try:
		temp_num = None
		try:
			if temp.startswith("0x") or temp.startswith("0X"):
				temp_num = int(temp, 16)
			elif temp.startswith("0o") or temp.startswith("0O"):
				temp_num = int(temp, 8)
			elif temp.startswith("0b") or temp.startswith("0B"):
				temp_num = int(temp, 2)
			else:
				temp_num = int(temp)
		except ValueError:
				temp_num = float(temp) # Propagate
		state["res"].append({
			"type": "number",
			"value": temp_num
		})
	except ValueError:
		state["err"] = f"Error at char {state['idx']}: Invalid number char \"{temp}\""

def parse_access(input, state):
	if input[state["idx"]] != '[':
		return True

	state["idx"] += 1

	res = []
	state["res"].append(res)
	state["stk"].append(res)
	state["res"] = res

	res.append({
		"type": "[",
	})

	while state["idx"] < state["len"]:
		token: str = input[state["idx"]]
		if token == ']':
			res.append({
				"type": "]",
			})
			state["idx"] += 1
			state["stk"].pop()
			state["res"] = state["stk"][-1]
			return
		elif token == '@':
			res.append({
				"type": "@",
			})
			state["idx"] += 1
		elif token == ',':
			res.append({
				"type": ",",
			})
			state["idx"] += 1
		elif (
			parse_number(input, state) and
			parse_extra(input, state) and
			parse_comment(input, state)
		):
			state["idx"] += 1

		if state["err"] is not None:
			return
	
	state["err"] = f"Error at char {state['idx']}: Unterminated \"[\""

def parse_extra(input, state):
	if input[state["idx"]] != '(':
		return True

	state["idx"] += 1

	res = []
	state["res"].append(res)
	state["stk"].append(res)
	state["res"] = res

	res.append({
		"type": "(",
	})

	while state["idx"] < state["len"]:
		token: str = input[state["idx"]]
		if token == ')':
			res.append({
				"type": ")",
			})
			state["idx"] += 1
			state["stk"].pop()
			state["res"] = state["stk"][-1]
			return
		elif token == '@':
			res.append({
				"type": "@",
			})
			state["idx"] += 1
		elif (
			parse_number(input, state) and
			parse_string(input, state) and
			parse_comment(input, state)
		):
			state["idx"] += 1

		if state["err"] is not None:
			return
		
	state["err"] = f"Error at char {state['idx']}: Unterminated \"(\""

# LL(1) parser is enough
# Modular design to easy manage
def parse_lang(input: str):
	ast = []

	state = {
		"idx": 0,
		"len": len(input),
		"err": None
	}
	state["res"] = []
	state["stk"] = [state["res"]]

	while state["idx"] < state["len"]:
		token = input[state["idx"]]
		if token == '!':
			state["res"].append({
				"type": "!",
			})
			state["idx"] += 1
		elif token == '=':
			state["res"].append({
				"type": "=",
			})
			state["idx"] += 1
		elif token == ';':
			ast.append(state["res"])
			state["res"] = []
			state["stk"] = [state["res"]]
			state["idx"] += 1
		elif (
			parse_space(input, state) and
			parse_name(input, state) and
			parse_string(input, state) and
			parse_access(input, state) and
			parse_extra(input, state) and
			parse_number(input, state) and
			parse_comment(input, state)
		):
			state["idx"] += 1

		if state["err"] is not None:
			return (None, state["err"])

	return (ast, None)

######################################################################################

# Trash that may be used later, don't mind me :)
# https://pastebin.com/raw/Z3Y9HimQ
# https://pastebin.com/raw/5kn6KKbT: The OG RevisionDict, if we needs it ever again

# Put a COND_EXEC to remove the curr_id from the EXEC_QUEUE
# def cond(exec_arr: list):
# 	for i in range(len(exec_arr)):
# 		if exec_arr[i][1] == curr_id:
# 			del exec_arr[i]
# 			return True
# 	return False
# COND_EXEC.append(cond)

# else:
# 	COND_EXEC.append(_event[0]["cond"](_prompt[0], _workflow[0], exec))

# def cond(_prompt, _workflow, id_arr):
# 	def inner_cond(exec_arr):
# 		temp = Count.COUNT_DB[_id]
# 		Count.COUNT_DB[_id] += 1
# 		if temp < int(_event) - 1:
# 			for curr_id in id_arr:
# 				trace_node(_prompt, curr_id, _workflow)
# 				exec_arr.append((0, curr_id))
# 			return False
# 		return True
	
# 	return inner_cond