
# Imports
import stouputils as stp
from beet import Context, LootTable
from beet.core.utils import JsonDict

from ....core.__memory__ import Mem
from ....core.cls.ingredients import Ingr
from ....core.cls.item import Item
from ....core.utils.io import write_function
from ....core.utils.loot_table import result_count_to_suffix


# Main entry point
@stp.measure_time(message="Execution time of 'stewbeet.plugins.datapack.loot_tables'")
def beet_default(ctx: Context):
	""" Main entry point for the loot tables plugin.
	This plugin sets up loot tables for items in the definitions and external items.

	Args:
		ctx (Context): The beet context.
	"""
	if Mem.ctx is None: # pyright: ignore[reportUnnecessaryComparison]
		Mem.ctx = ctx

	# Get data from memory
	assert ctx.meta["stewbeet"].get("source_lore", None) is not None, "Source lore is not set in the context metadata."
	assert ctx.project_id, "Project ID is not set. Please set it in the project configuration."
	ns: str = ctx.project_id

	# Creative loot table (sort of give all loot table)
	creative_loot_table: JsonDict = {"pools": []}

	# For each item in the definitions, create a loot table
	for item in Mem.definitions.keys():
		obj = Item.from_id(item)
		loot_table: JsonDict = {
			"pools": [{
				"rolls": 1,
				"entries": [{
					"type": "minecraft:item",
					"name": obj.base_item
				}]
			}]
		}

		# Set components
		set_components: JsonDict = {
			"function": "minecraft:set_components",
			"components": {}
		}
		for k, v in obj.components.items():
			if k.startswith("!"):
				set_components["components"][f"!minecraft:{k[1:]}"] = {}
			else:
				set_components["components"][f"minecraft:{k}"] = v

		# Add functions
		loot_table["pools"][0]["entries"][0]["functions"] = [set_components]

		# Create loot table with beet
		ctx.data[ns].loot_tables[f"i/{item}"] = LootTable(stp.json_dump(loot_table, max_level = 10))

		# Add the pool to the creative loot table
		creative_loot_table["pools"].append({"rolls": 1, "entries":[{"type":"minecraft:loot_table","value":f"{ns}:i/{item}"}] })

	# Same for external items
	for item in Mem.external_definitions.keys():
		obj = Item.from_id(item)
		ext_ns, item_name = item.split(":")
		loot_table: JsonDict = {
			"pools": [{
				"rolls": 1,
				"entries": [{
					"type": "minecraft:item",
					"name": obj.base_item
				}]
			}]
		}
		set_components: JsonDict = {
			"function": "minecraft:set_components",
			"components": {}
		}
		for k, v in obj.components.items():
			if k.startswith("!"):
				set_components["components"][f"!minecraft:{k[1:]}"] = {}
			else:
				set_components["components"][f"minecraft:{k}"] = v
		loot_table["pools"][0]["entries"][0]["functions"] = [set_components]

		# Create external loot table with beet
		ctx.data[ns].loot_tables[f"external/{ext_ns}/{item_name}"] = LootTable(stp.json_dump(loot_table, max_level=10))

	# Loot tables for items with crafting recipes
	for item in Mem.definitions.keys():
		obj = Item.from_id(item)
		if obj.recipes:
			results: list[int | JsonDict] = []
			for recipe in obj.recipes:
				# Verify the result is for this item
				if recipe.get("result"):
					result_id = Ingr(recipe["result"]).to_id(add_namespace=False)
					if result_id != item:
						continue

				# Verify count different than 1
				count: int | JsonDict = recipe.get("result_count", 1)
				if count != 1:
					results.append(count)

			# For each result count, create a loot table for it
			for result_count in results:
				loot_table: JsonDict = {
					"pools": [{
						"rolls": 1,
						"entries": [{
							"type": "minecraft:loot_table",
							"value": f"{ns}:i/{item}",
							"functions": [{
								"function": "minecraft:set_count",
								"count": result_count
							}]
						}]
					}]
				}
				loot_table_path = f"i/{item}{result_count_to_suffix(result_count)}"
				ctx.data[ns].loot_tables[loot_table_path] = LootTable(stp.json_dump(loot_table, max_level=10))

	# Second loot table for the manual (if present)
	if "manual" in Mem.definitions:
		loot_table: JsonDict = {
			"pools": [{
				"rolls": 1,
				"entries": [{
					"type": "minecraft:loot_table",
					"value": f"{ns}:i/manual"
				}]
			}]
		}
		ctx.data[ns].loot_tables[f"i/{ns}_manual"] = LootTable(stp.json_dump(loot_table, max_level=10))

	# Write the creative loot table
	if creative_loot_table["pools"]:
		ctx.data[ns].loot_tables["creative_loot_table"] = LootTable(stp.json_dump(creative_loot_table, max_level=2))

	# Make a give all command that gives chests with all the items
	CHEST_SIZE: int = 27
	total_chests: int = (len(Mem.definitions) + CHEST_SIZE - 1) // CHEST_SIZE

	# Get source lore from context metadata
	source_lore: JsonDict = ctx.meta["stewbeet"]["source_lore"]
	lore = stp.json_dump(source_lore, max_level=0).strip()

	chests: list[str] = []
	definitions_copy: list[tuple[str, Item]] = [(item, Item.from_id(item)) for item in Mem.definitions.keys()]
	for i in range(total_chests):
		chest_contents: list[str] = []

		# For each slot of the chest, append an item and remove it from the copy
		for j in range(CHEST_SIZE):
			if not definitions_copy:
				break
			item, obj = definitions_copy.pop(0)
			json_content = stp.json_dump(obj.components, max_level=0).strip()
			chest_contents.append(f'{{slot:{j},item:{{count:1,id:"{obj.base_item}",components:{json_content}}}}}')

		joined_content = ",".join(chest_contents)
		chests.append(f'give @s chest[container=[{joined_content}],custom_name={{"text":"Chest [{i+1}/{total_chests}]","color":"yellow"}},lore=[{lore}]]')

	# Write the give all function
	write_function(f"{ns}:_give_all", "\n" + "\n\n".join(chests) + "\n\n")

