from comfy_api.latest import ComfyExtension, io


VALUE_PREFIX = "value"
INITIAL_VALUE_INPUTS = 10
MAX_VALUE_INPUTS = 100
MISSING = object()


class LazyFallbackSwitch(io.ComfyNode):
	@classmethod
	def define_schema(cls) -> io.Schema:
		value_type = io.MatchType.Template("value_type")
		value_inputs = [
			io.MatchType.Input(
				f"{VALUE_PREFIX}{index}",
				template=value_type,
				lazy=True,
				optional=True,
			)
			for index in range(INITIAL_VALUE_INPUTS)
		]

		return io.Schema(
			node_id="SparknightLazyFallbackSwitchStatic",
			display_name="Lazy Fallback Switch",
			category="Sparknight/Logic",
			description="Selects valueN by index and lazily uses default when that input is not connected.",
			search_aliases=["lazy index switch", "switch case", "default case", "fallback"],
			inputs=[
				io.Int.Input(
					"index",
					default=0,
					step=1,
					tooltip="Selects the matching valueN input. Any missing value uses default.",
				),
				io.MatchType.Input(
					"default",
					template=value_type,
					lazy=True,
					optional=True,
					tooltip="Used when the indexed value input is not connected. An unconnected default returns None.",
				),
				*value_inputs,
			],
			outputs=[
				io.MatchType.Output(
					template=value_type,
					display_name="output",
					tooltip="The selected indexed value, or default.",
				),
			],
		)

	@classmethod
	def check_lazy_status(cls, index: int, default=MISSING, **values) -> list[str]:
		value_name = f"{VALUE_PREFIX}{index}"
		if value_name in values:
			return [value_name] if values[value_name] is None else []

		return ["default"] if default is None else []

	@classmethod
	def execute(cls, index: int, default=MISSING, **values) -> io.NodeOutput:
		value_name = f"{VALUE_PREFIX}{index}"
		if value_name in values:
			return io.NodeOutput(values[value_name])

		return io.NodeOutput(None if default is MISSING else default)


class LazyFallbackSwitchAutogrow(io.ComfyNode):
	@classmethod
	def define_schema(cls) -> io.Schema:
		value_type = io.MatchType.Template("value_type")
		value_inputs = io.Autogrow.TemplatePrefix(
			input=io.MatchType.Input("value", template=value_type, lazy=True, optional=True),
			prefix=VALUE_PREFIX,
			min=INITIAL_VALUE_INPUTS - 1,
			max=MAX_VALUE_INPUTS,
		)

		return io.Schema(
			node_id="SparknightLazyFallbackSwitch",
			display_name="Lazy Fallback Switch (Autogrow)",
			category="Sparknight/Logic",
			description="Selects valueN by index and lazily uses default when that input is not connected.",
			search_aliases=["lazy index switch", "switch case", "default case", "fallback", "autogrow"],
			inputs=[
				io.Int.Input(
					"index",
					default=0,
					step=1,
					tooltip="Selects the matching valueN input. Any missing value uses default.",
				),
				io.MatchType.Input(
					"default",
					template=value_type,
					lazy=True,
					optional=True,
					tooltip="Used when the indexed value input is not connected. An unconnected default returns None.",
				),
				io.Autogrow.Input(
					"values",
					template=value_inputs,
					optional=True,
					tooltip="Connect only the indexed routes that need custom processing.",
				),
			],
			outputs=[
				io.MatchType.Output(
					template=value_type,
					display_name="output",
					tooltip="The selected indexed value, or default.",
				),
			],
		)

	@classmethod
	def check_lazy_status(cls, index: int, default=MISSING, values: io.Autogrow.Type = None) -> list[str]:
		value_name = f"{VALUE_PREFIX}{index}"
		if values is not None and value_name in values:
			value, input_name = values[value_name]
			return [input_name] if value is None else []

		return ["default"] if default is None else []

	@classmethod
	def execute(cls, index: int, default=MISSING, values: io.Autogrow.Type = None) -> io.NodeOutput:
		value_name = f"{VALUE_PREFIX}{index}"
		if values is not None and value_name in values:
			return io.NodeOutput(values[value_name])

		return io.NodeOutput(None if default is MISSING else default)


class LazyFallbackSwitchExtension(ComfyExtension):
	async def get_node_list(self) -> list[type[io.ComfyNode]]:
		return [LazyFallbackSwitch, LazyFallbackSwitchAutogrow]


async def comfy_entrypoint() -> LazyFallbackSwitchExtension:
	return LazyFallbackSwitchExtension()