# IMPORT
from PIL import Image
import os
os.chdir(os.path.dirname(os.path.abspath(__file__)))


# ====
# DIRS
# ====

trim_texture_file = "tooltrims_ultimate_key.png"
extra_texture_file = "tooltrims_extra_key.png"
item_atlas_file = "minecraft/atlases/items.json"

color_palettes_dir = "tooltrims/textures/trims/color_palettes/"
trim_bases_dir = "tooltrims/textures/item/trim_bases/"
trim_key_dir = "tooltrims/textures/trims/items/"
trident_entity_dir = "tooltrims/textures/item/trident/"
item_model_dir = "tooltrims/models/trims/"
item_definition_dir = "minecraft/items/"

# =========
# MAIN VARS
# =========
selected_to_run = {"generate_images": False, "set_atlas": False, "generate_models": False, "generate_definitions": True}

palette_names = ["key", "amethyst", "copper", "diamond", "emerald", "gold", "iron", "lapis", "netherite", "quartz", "redstone", "resin"]

tool = ["palettes","netherite_sword","netherite_pickaxe","netherite_axe","netherite_shovel","netherite_hoe","diamond_sword","diamond_pickaxe","diamond_axe","diamond_shovel","diamond_hoe","iron_sword","iron_pickaxe","iron_axe","iron_shovel","iron_hoe","golden_sword","golden_pickaxe","golden_axe","golden_shovel","golden_hoe","stone_sword","stone_pickaxe","stone_axe","stone_shovel","stone_hoe","wooden_sword","wooden_pickaxe","wooden_axe","wooden_shovel","wooden_hoe","bow","bow_pulling_0","bow_pulling_1","bow_pulling_2","crossbow_standby","crossbow_pulling_0","crossbow_pulling_1","crossbow_pulling_2","crossbow_arrow","crossbow_firework","mace","trident","copper_sword","copper_pickaxe","copper_axe","copper_shovel","copper_hoe","netherite_spear","diamond_spear","iron_spear","golden_spear","stone_spear","wooden_spear","copper_spear"]
tool_extra = ["trident_entity", "netherite_spear_in_hand", "diamond_spear_in_hand", "iron_spear_in_hand", "golden_spear_in_hand", "stone_spear_in_hand", "wooden_spear_in_hand", "copper_spear_in_hand"]

tool_models = tool + ["trident_entity", "trident_entity_throwing", "netherite_spear_in_hand", "diamond_spear_in_hand", "iron_spear_in_hand", "golden_spear_in_hand", "stone_spear_in_hand", "wooden_spear_in_hand", "copper_spear_in_hand"]
tool_model_definitions = ["netherite_sword","netherite_pickaxe","netherite_axe","netherite_shovel","netherite_hoe","diamond_sword","diamond_pickaxe","diamond_axe","diamond_shovel","diamond_hoe","iron_sword","iron_pickaxe","iron_axe","iron_shovel","iron_hoe","golden_sword","golden_pickaxe","golden_axe","golden_shovel","golden_hoe","stone_sword","stone_pickaxe","stone_axe","stone_shovel","stone_hoe","wooden_sword","wooden_pickaxe","wooden_axe","wooden_shovel","wooden_hoe","bow","crossbow","mace","trident","copper_sword","copper_pickaxe","copper_axe","copper_shovel","copper_hoe","netherite_spear","diamond_spear","iron_spear","golden_spear","stone_spear","wooden_spear","copper_spear"]

trims = ["untrimmed","linear","tracks","charge","frost"]

files_created = 0
files_skipped = 0
last_message = ""
last_message_sent = False

def file_created():
    global files_created
    global last_message
    global last_message_sent

    files_created +=1
    last_message_sent = False

    last_message = str(files_created) + " files created or modified"
    if ( (files_created) / 10 ).is_integer():
        print(last_message)
        last_message_sent = True


# ==========
# GEN IMAGES
# ==========

input_image = (Image.open(trim_texture_file)).convert("RGBA")
image_map = input_image.load()

width, height = input_image.size
grid_size = 16

palette_size = 8
palette_colors = []
count_amount = 0

# gen palettes

if selected_to_run["generate_images"]:
    for current_palette_number in range(len(palette_names)):

        # get colour values in a list
        current_list = []
        for current_color in range(0,palette_size):
            current_height = current_palette_number % grid_size
            current_width = (palette_size * (current_palette_number // grid_size) ) + current_color

            current_list.append(input_image.getpixel((current_width, current_height)))
        palette_colors.append(current_list)

        # cut palettes into images
        palette_image = input_image.crop((  (current_palette_number // grid_size) * palette_size,
                                            (current_palette_number) % grid_size,
                                            (current_palette_number // grid_size) * palette_size + palette_size, 
                                            (current_palette_number) % grid_size + 1
                                        ))
        
        # save file
        palette_image.save(color_palettes_dir + palette_names[current_palette_number] + ".png")
        file_created()
        
    # erase non-palette pixels from the input image
    for current_width in range(width):
        for current_height in range(height):
            r, g, b, p = input_image.getpixel((current_width, current_height))
        
            if not( (r, g, b, p) in palette_colors[0] ):
                image_map[current_width, current_height] = (0,0,0,0)

    # cut image for tooltrims:trims/items
    for j in range(1, int((height)/grid_size)):
        for i in range(1, int(width/grid_size)):
            box = (i*grid_size, j*grid_size, (i+1)*grid_size, (j+1)*grid_size)
            splited = input_image.crop(box)

            # save file
            splited.save(trim_key_dir + f"{tool[j]}_{trims[i]}.png", format="png")
            file_created()



    # erase non-palette pixels from the trident input image
    extra_input_image = (Image.open(extra_texture_file)).convert("RGBA")
    extra_image_map = extra_input_image.load()
    extra_grid_size = 32

    extra_image_width, extra_image_height = extra_input_image.size

    for current_width in range(extra_image_width):
        for current_height in range(extra_image_height):
            r, g, b, p = extra_input_image.getpixel((current_width, current_height))
        
            if not( (r, g, b, p) in palette_colors[0] ):
                extra_image_map[current_width, current_height] = (0,0,0,0)

    # cut image for the trident entity texture
    for j in range(1, int((extra_image_height)/extra_grid_size)):
        for i in range(1, int(extra_image_width/extra_grid_size)):
            box = (i*extra_grid_size, j*extra_grid_size, (i+1)*extra_grid_size, (j+1)*extra_grid_size)
            splited = extra_input_image.crop(box)
            
            # save file
            splited.save(trim_key_dir + f"{tool_extra[j]}_{trims[i]}.png", format="png")
            file_created()



    # generate a cutout image for the tools (if needed)
    tools_image = (Image.open(trim_texture_file)).convert("RGBA")
    tools_map = tools_image.load()

    for current_width in range(width):
        for current_height in range(height):
            r, g, b, p = tools_image.getpixel((current_width, current_height))
        
            if (r, g, b, p) in palette_colors[0]:
                tools_map[current_width, current_height] = (0,0,0,0)

    for j in range(1, int((height)/grid_size)):
        for i in range(1, int(width/grid_size)):
            box = (i*grid_size, j*grid_size, (i+1)*grid_size, (j+1)*grid_size)
            splited = tools_image.crop(box)

            # save file
            splited.save(trim_bases_dir + f"{tool[j]}_{trims[i]}.png", format="png")
            file_created()


# =================
# ADD KEYS TO ATLAS
# =================

if selected_to_run["set_atlas"]:
    key_list = os.listdir(trim_key_dir)
    key_list_text = ""

    for current_key in range(len(key_list)):
        if current_key == len(key_list) - 1:
            key_list_text += "        \"tooltrims:" + trim_key_dir[19::] + key_list[current_key][:-4:] + "\""
        else:
            key_list_text += "        \"tooltrims:" + trim_key_dir[19::] + key_list[current_key][:-4:] + "\",\n"


    palette_list = os.listdir(color_palettes_dir)
    palette_list.remove("key.png")
    palette_list_text = ""

    for current_palette in range(len(palette_list)):
        if current_palette == len(palette_list) - 1:
            palette_list_text += "        \"" + palette_list[current_palette][:-4:] + "\": \"tooltrims:" + color_palettes_dir[19::] + palette_list[current_palette][:-4:] + "\""
        else:
            palette_list_text += "        \"" + palette_list[current_palette][:-4:] + "\": \"tooltrims:" + color_palettes_dir[19::] + palette_list[current_palette][:-4:] + "\",\n"


    with open(item_atlas_file,"w+") as file:
        file.write(f"""{{
"sources": [
    {{
    "type": "minecraft:paletted_permutations",
    "palette_key": "tooltrims:trims/color_palettes/key",
    "permutations": {{
{palette_list_text}
    }},
    "textures": [
{key_list_text}
    ]
    }}
]
}}""")
        file_created()



# =========================
# GENERATE ITEM MODEL FILES
# =========================

if selected_to_run["generate_models"]:
    for current_tool in tool_models[1::]:
        for current_trim in trims[1::]:
            for current_palette in palette_names[1::]:
                
                with open(item_model_dir + f"{current_tool}_{current_trim}_{current_palette}.json","w+") as file:
                    
                    # trident exceptions
                    if current_tool == "trident_entity":
                        file.write(f"""{{
    "parent": "tooltrims:item/trident_entity",
    "textures": {{
        "layer0": "tooltrims:item/trim_bases/trident_entity",
        "layer1": "tooltrims:trims/items/trident_entity_{current_trim}_{current_palette}"
    }}
}}""")
                    elif current_tool == "trident_entity_throwing":
                        file.write(f"""{{
    "parent": "tooltrims:item/{current_tool}",
    "textures": {{
        "layer0": "tooltrims:item/trim_bases/trident_entity",
        "layer1": "tooltrims:trims/items/trident_entity_{current_trim}_{current_palette}"
    }}
}}""")
                    # crossbow exception
                    elif current_tool == "crossbow_standby":
                        file.write(f"""{{
    "parent": "minecraft:item/crossbow",
    "textures": {{
        "layer0": "minecraft:item/{current_tool}",
        "layer1": "tooltrims:trims/items/{current_tool}_{current_trim}_{current_palette}"
    }}
}}""")
                        
                    # the other tools
                    else:
                        file.write(f"""{{
    "parent": "minecraft:item/{current_tool}",
    "textures": {{
        "layer0": "minecraft:item/{current_tool}",
        "layer1": "tooltrims:trims/items/{current_tool}_{current_trim}_{current_palette}"
    }}
}}""")
                    
                    file_created()



# ==============================
# GENERATE ITEM DEFINATION FILES
# ==============================

if selected_to_run["generate_definitions"]:
    def main_data_condition_text(data_conditions, trim_conditions, extra_keys):
        root = f"""{{
    "model": {{
        "type": "minecraft:condition",
        "property": "minecraft:has_component",
        "component": "custom_model_data",
        "on_true": {{ {data_conditions}
        }},
        "on_false": {{ {trim_conditions}
        }}
    }}{extra_keys}
}}"""
        return root
                


    def trim_condition_text(current_tool, *vanilla_tool):
        if vanilla_tool == ():
            vanilla_tool = current_tool
            condition_text = f"""
"type": "minecraft:select",
"fallback": {{
    "type": "minecraft:model",
    "model": "minecraft:item/{vanilla_tool}"
}},
"property": "minecraft:component",
"component": "minecraft:trim",
"cases": [
"""
            
        elif len(vanilla_tool) == 1:
            condition_text = f"""
"type": "minecraft:select",
"fallback": {{
    "type": "minecraft:model",
    "model": "minecraft:item/{vanilla_tool[0]}"
}},
"property": "minecraft:component",
"component": "minecraft:trim",
"cases": [
"""

        elif vanilla_tool[1] == "trident_special":
            condition_text = f"""
"type": "minecraft:select",
"fallback": {{
    "type": "minecraft:special",
    "base": "minecraft:item/{vanilla_tool[0]}",
    "model": {{
        "type": "minecraft:trident"
    }},
    "transformation": {{
        "left_rotation": [
        0.0,
        0.0,
        0.0,
        1.0
        ],
        "right_rotation": [
        0.0,
        0.0,
        0.0,
        1.0
        ],
        "scale": [
        1.0,
        -1.0,
        -1.0
        ],
        "translation": [
        0.0,
        0.0,
        0.0
        ]
    }}
}},
"property": "minecraft:component",
"component": "minecraft:trim",
"cases": [
"""
        
        for current_trim in trims[1::]:
            for current_palette in palette_names[1::]:
                condition_text += f'    {{"model": {{"type": "model", "model": "tooltrims:trims/{current_tool}_{current_trim}_{current_palette}"}}, "when": {{"pattern":"tooltrims:{current_trim}", "material": "{current_palette}"}}}},\n'

        condition_text = condition_text[:-2:] + """
]
"""
        return condition_text

    def data_condition_text(current_tool, *vanilla_tool):
        if vanilla_tool == ():
            vanilla_tool = current_tool
            condition_text = f"""
"type": "minecraft:range_dispatch",
"fallback": {{
    "type": "minecraft:model",
    "model": "minecraft:item/{vanilla_tool}"
}},
"property": "minecraft:custom_model_data",
"entries": [
"""
            
        elif len(vanilla_tool) == 1:
            condition_text = f"""
"type": "minecraft:range_dispatch",
"fallback": {{
    "type": "minecraft:model",
    "model": "minecraft:item/{vanilla_tool[0]}"
}},
"property": "minecraft:custom_model_data",
"entries": [
"""

        elif vanilla_tool[1] == "trident_special":
            condition_text = f"""
"type": "minecraft:range_dispatch",
"fallback": {{
    "type": "minecraft:special",
    "base": "minecraft:item/{vanilla_tool[0]}",
    "model": {{
        "type": "minecraft:trident"
    }},
    "transformation": {{
        "left_rotation": [
        0.0,
        0.0,
        0.0,
        1.0
        ],
        "right_rotation": [
        0.0,
        0.0,
        0.0,
        1.0
        ],
        "scale": [
        1.0,
        -1.0,
        -1.0
        ],
        "translation": [
        0.0,
        0.0,
        0.0
        ]
    }}
}},
"property": "minecraft:custom_model_data",
"entries": [
"""

        # How custom model data is structured: 
        #   linear: 31100x (amethyst-redstone)
        #   tracks: 31101x (amethyst-redstone)
        #   charge: 31102x (amethyst-redstone)
        #   frost:  31103x (amethyst-redstone)
        #   resin:  31104x (linear-frost)

        custom_model_data = 311001.0
        resin_custom_model_data = 311041.0

        for current_trim in trims[1::]:
            for current_palette in palette_names[1::]:
                if current_palette == "resin":
                    condition_text += f'    {{"model": {{"type": "model", "model": "tooltrims:trims/{current_tool}_{current_trim}_{current_palette}"}}, "threshold": {resin_custom_model_data} }},\n'
                    resin_custom_model_data += 1
                else:
                    condition_text += f'    {{"model": {{"type": "model", "model": "tooltrims:trims/{current_tool}_{current_trim}_{current_palette}"}}, "threshold": {custom_model_data} }},\n'
                    if custom_model_data == 311040.0:
                        custom_model_data +=11
                    else:
                        custom_model_data +=1

        condition_text = condition_text[:-2:] + """
]
"""
        return condition_text
                


    for current_definition in tool_model_definitions:
        with open(item_definition_dir + f"{current_definition}.json","w+") as file:

            # bow exception
            if current_definition == "bow":
                file.write(
                    main_data_condition_text(
f"""
            "type": "minecraft:condition",
            "property": "minecraft:using_item",
            "on_false": {{ 
{data_condition_text(current_definition)}
            }},
            "on_true": {{
                "type": "minecraft:range_dispatch",
                "property": "minecraft:use_duration",
                "scale": 0.05,
                "entries": [
                    {{ 
                        "model": {{
{data_condition_text("bow_pulling_1")}
                        }},
                        "threshold": 0.65
                    }},
                    {{ 
                        "model": {{
{data_condition_text("bow_pulling_2")}
                        }},
                        "threshold": 0.9
                    }}
                ],
                "fallback": {{ 
{data_condition_text("bow_pulling_0")}
                }}
            }}""",

f"""
            "type": "minecraft:condition",
            "property": "minecraft:using_item",
            "on_false": {{ 
{trim_condition_text(current_definition)}
            }},
            "on_true": {{
                "type": "minecraft:range_dispatch",
                "property": "minecraft:use_duration",
                "scale": 0.05,
                "entries": [
                    {{ 
                        "model": {{
{trim_condition_text("bow_pulling_1")}
                        }},
                        "threshold": 0.65
                    }},
                    {{ 
                        "model": {{
{trim_condition_text("bow_pulling_2")}
                        }},
                        "threshold": 0.9
                    }}
                ],
                "fallback": {{ 
{trim_condition_text("bow_pulling_0")}
                }}
            }}""",
            
            ""
                    )
                )
            







            # crossbow exception
            elif current_definition == "crossbow":
                file.write(
                    main_data_condition_text(
f"""
            "type": "minecraft:condition",
            "on_false": {{
                "type": "minecraft:select",
                "cases": [
                    {{
                    "model": {{
{data_condition_text("crossbow_arrow")}
                    }},
                    "when": "arrow"
                    }},
                    {{
                    "model": {{
{data_condition_text("crossbow_firework")}
                    }},
                    "when": "rocket"
                    }}
                ],
                "fallback": {{
{data_condition_text("crossbow_standby", "crossbow")}
                }},
                "property": "minecraft:charge_type"
            }},
            "on_true": {{
                "type": "minecraft:range_dispatch",
                "entries": [
                    {{
                    "model": {{
{data_condition_text("crossbow_pulling_1")}
                    }},
                    "threshold": 0.58
                    }},
                    {{
                    "model": {{
{data_condition_text("crossbow_pulling_2")}
                    }},
                    "threshold": 1.0
                    }}
                ],
                "fallback": {{
{data_condition_text("crossbow_pulling_0")}
                }},
                "property": "minecraft:crossbow/pull"
            }},
            "property": "minecraft:using_item\"""",

f"""
            "type": "minecraft:condition",
            "on_false": {{
                "type": "minecraft:select",
                "cases": [
                    {{
                    "model": {{
{trim_condition_text("crossbow_arrow")}
                    }},
                    "when": "arrow"
                    }},
                    {{
                    "model": {{
{trim_condition_text("crossbow_firework")}
                    }},
                    "when": "rocket"
                    }}
                ],
                "fallback": {{
{trim_condition_text("crossbow_standby", "crossbow")}
                }},
                "property": "minecraft:charge_type"
            }},
            "on_true": {{
                "type": "minecraft:range_dispatch",
                "entries": [
                    {{
                    "model": {{
{trim_condition_text("crossbow_pulling_1")}
                    }},
                    "threshold": 0.58
                    }},
                    {{
                    "model": {{
{trim_condition_text("crossbow_pulling_2")}
                    }},
                    "threshold": 1.0
                    }}
                ],
                "fallback": {{
{trim_condition_text("crossbow_pulling_0")}
                }},
                "property": "minecraft:crossbow/pull"
            }},
            "property": "minecraft:using_item\"""",
            
            ""
                    )
                )








            # trident exception
            elif current_definition == "trident":
                file.write(
                    main_data_condition_text(
f"""
            "type": "minecraft:select",
            "cases": [
            {{
                "model": {{
{data_condition_text("trident")}
                }},
                "when": [ "gui", "ground", "fixed" ]
            }}
            ],
            "fallback": {{
            "type": "minecraft:condition",
            "on_false": {{
{data_condition_text("trident_entity", "trident_in_hand", "trident_special")}
            }},
            "on_true": {{
{data_condition_text("trident_entity_throwing", "trident_throwing", "trident_special")}
            }},
            "property": "minecraft:using_item"
            }},
            "property": "minecraft:display_context\"""",

f"""
            "type": "minecraft:select",
            "cases": [
            {{
                "model": {{
{trim_condition_text("trident")}
                }},
                "when": [ "gui", "ground", "fixed" ]
            }}
            ],
            "fallback": {{
            "type": "minecraft:condition",
            "on_false": {{
{trim_condition_text("trident_entity", "trident_in_hand", "trident_special")}
            }},
            "on_true": {{
{trim_condition_text("trident_entity_throwing", "trident_throwing", "trident_special")}
            }},
            "property": "minecraft:using_item"
            }},
            "property": "minecraft:display_context\"""",
            
            ""
                    )
                )








            # spear exception
            elif "spear" in current_definition:
                file.write(
                    main_data_condition_text(
f"""
            "type": "minecraft:select",
            "cases": [
            {{
                "model": {{
{data_condition_text(current_definition)}
                }},
                "when": [ "gui", "ground", "fixed", "on_shelf" ]
            }}
            ],
            "fallback": {{
{data_condition_text(current_definition + "_in_hand")}
                }},
            "property": "minecraft:display_context\"""",

f"""
            "type": "minecraft:select",
            "cases": [
            {{
                "model": {{
{trim_condition_text(current_definition)}
                }},
                "when": [ "gui", "ground", "fixed", "on_shelf" ]
            }}
            ],
            "fallback": {{
{trim_condition_text(current_definition + "_in_hand")}
                }},
            "property": "minecraft:display_context\"""",

                    """,
    \"swap_animation_scale\": 1.95"""
                    )
                )
            





            

            # other tools
            else:
                file.write(
                    main_data_condition_text(
                f"""
{data_condition_text(current_definition)}""",
                f"""
{trim_condition_text(current_definition)}""",
                ""
                )
            )

if not(last_message_sent):
    print(last_message)