package net.mattias.blockarmor.datagen;

import com.google.common.hash.HashCode;
import net.mattias.blockarmor.BlockArmor;
import net.mattias.blockarmor.item.ModItems;
import net.minecraft.data.DataGenerator;
import net.minecraft.data.DataProvider;
import net.minecraft.data.CachedOutput;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.world.level.block.Block;
import net.minecraft.world.level.block.Blocks;
import net.minecraftforge.common.data.ExistingFileHelper;
import net.minecraftforge.registries.ForgeRegistries;
import net.minecraftforge.registries.RegistryObject;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;

public class ModTextureProvider implements DataProvider {
    private final DataGenerator generator;
    private final ExistingFileHelper existingFileHelper;

    private BufferedImage templateIronLayer1;
    private BufferedImage templateIronLayer2;

    private BufferedImage templateIronHelmet;
    private BufferedImage templateIronChestplate;
    private BufferedImage templateIronLeggings;
    private BufferedImage templateIronBoots;

    private final Map<Block, ResourceLocation> generatedArmorLayerTextures = new HashMap<>();
    private final Map<RegistryObject<?>, ResourceLocation> generatedItemIconTextures = new HashMap<>();

    private static class ArmorItemInfo {
        final RegistryObject<?> item;
        final String blockName;
        final int armorLayerType;

        ArmorItemInfo(RegistryObject<?> item, String blockName, int armorLayerType) {
            this.item = item;
            this.blockName = blockName;
            this.armorLayerType = armorLayerType;
        }
    }

    private static final List<ArmorItemInfo> ARMOR_ITEM_INFOS = new ArrayList<>() {{
        // Dirt Armor
        add(new ArmorItemInfo(ModItems.DIRT_HELMET, "dirt", 1));
        add(new ArmorItemInfo(ModItems.DIRT_CHESTPLATE, "dirt", 1));
        add(new ArmorItemInfo(ModItems.DIRT_LEGGINGS, "dirt", 2));
        add(new ArmorItemInfo(ModItems.DIRT_BOOTS, "dirt", 1));

        // Cobblestone Armor
        add(new ArmorItemInfo(ModItems.COBBLESTONE_HELMET, "cobblestone", 1));
        add(new ArmorItemInfo(ModItems.COBBLESTONE_CHESTPLATE, "cobblestone", 1));
        add(new ArmorItemInfo(ModItems.COBBLESTONE_LEGGINGS, "cobblestone", 2));
        add(new ArmorItemInfo(ModItems.COBBLESTONE_BOOTS, "cobblestone", 1));

        add(new ArmorItemInfo(ModItems.CRAFTING_TABLE_HELMET, "crafting_table", 1));
        add(new ArmorItemInfo(ModItems.CRAFTING_TABLE_CHESTPLATE, "crafting_table", 1));
        add(new ArmorItemInfo(ModItems.CRAFTING_TABLE_LEGGINGS, "crafting_table", 2));
        add(new ArmorItemInfo(ModItems.CRAFTING_TABLE_BOOTS, "crafting_table", 1));

        add(new ArmorItemInfo(ModItems.OAK_PLANKS_HELMET, "oak_planks", 1));
        add(new ArmorItemInfo(ModItems.OAK_PLANKS_CHESTPLATE, "oak_planks", 1));
        add(new ArmorItemInfo(ModItems.OAK_PLANKS_LEGGINGS, "oak_planks", 2));
        add(new ArmorItemInfo(ModItems.OAK_PLANKS_BOOTS, "oak_planks", 1));
    }};

    public ModTextureProvider(DataGenerator generator, ExistingFileHelper existingFileHelper) {
        this.generator = generator;
        this.existingFileHelper = existingFileHelper;
    }

    @Override
    public CompletableFuture<?> run(CachedOutput cache) {
        Path armorLayerOutputFolder = generator.getPackOutput().getOutputFolder()
                .resolve("assets/" + BlockArmor.MOD_ID + "/textures/models/armor");
        Path itemIconOutputFolder = generator.getPackOutput().getOutputFolder()
                .resolve("assets/" + BlockArmor.MOD_ID + "/textures/item");

        try {
            Files.createDirectories(armorLayerOutputFolder);
            Files.createDirectories(itemIconOutputFolder);

            loadTemplates();
        } catch (IOException e) {
            throw new RuntimeException("Failed to initialize directories or load templates", e);
        }

        for (Block block : ForgeRegistries.BLOCKS.getValues()) {
            if (isValidForArmor(block)) {
                try {
                    generateArmorLayerTextures(block, armorLayerOutputFolder, cache);
                } catch (IOException e) {
                    System.err.println("Skipping armor layer texture generation for block " +
                            ForgeRegistries.BLOCKS.getKey(block) + ": " + e.getMessage());
                }
            }
        }

        for (ArmorItemInfo info : ARMOR_ITEM_INFOS) {
            try {
                generateItemIconTexture(info, itemIconOutputFolder, cache);
            } catch (IOException e) {
                System.err.println("Skipping item icon texture generation for " +
                        info.item.getId() + ": " + e.getMessage());
            }
        }

        return CompletableFuture.completedFuture(null);
    }

    private void loadTemplates() throws IOException {

        templateIronLayer1 = loadTexture("assets/minecraft/textures/models/armor/iron_layer_1.png");
        templateIronLayer2 = loadTexture("assets/minecraft/textures/models/armor/iron_layer_2.png");

        templateIronHelmet = loadTexture("assets/minecraft/textures/item/iron_helmet.png");
        templateIronChestplate = loadTexture("assets/minecraft/textures/item/iron_chestplate.png");
        templateIronLeggings = loadTexture("assets/minecraft/textures/item/iron_leggings.png");
        templateIronBoots = loadTexture("assets/minecraft/textures/item/iron_boots.png");
    }

    private BufferedImage loadTexture(String path) throws IOException {
        InputStream stream = getClass().getClassLoader().getResourceAsStream(path);
        if (stream == null) {
            throw new IOException("Resource not found: " + path);
        }
        return ImageIO.read(stream);
    }

    private void generateArmorLayerTextures(Block block, Path outputFolder, CachedOutput cache) throws IOException {
        ResourceLocation blockRL = ForgeRegistries.BLOCKS.getKey(block);
        if (blockRL == null) return;

        String blockName = blockRL.getPath();
        String blockTexPath = "assets/minecraft/textures/block/" + blockName + ".png";
        BufferedImage blockTex;
        try {
            blockTex = loadTexture(blockTexPath);
        } catch (IOException e) {
            return;
        }

        BufferedImage layer1Output = createMaskedTexture(blockTex, templateIronLayer1);
        saveTexture(layer1Output, outputFolder, cache, blockName + "_layer_1.png");

        BufferedImage layer2Output = createMaskedTexture(blockTex, templateIronLayer2);
        saveTexture(layer2Output, outputFolder, cache, blockName + "_layer_2.png");

        generatedArmorLayerTextures.put(block, new ResourceLocation(BlockArmor.MOD_ID,
                "textures/models/armor/" + blockName));
    }

    private void generateItemIconTexture(ArmorItemInfo info, Path outputFolder, CachedOutput cache) throws IOException {
        String blockName = info.blockName;

        String itemPath = info.item.getId().getPath();

        String itemArmorType = itemPath.substring(blockName.length() + 1);

        String blockTexPath = "assets/minecraft/textures/block/" + blockName + ".png";
        BufferedImage blockTex = loadTexture(blockTexPath);

        BufferedImage template;
        switch (itemArmorType) {
            case "helmet":
                template = templateIronHelmet;
                break;
            case "chestplate":
                template = templateIronChestplate;
                break;
            case "leggings":
                template = templateIronLeggings;
                break;
            case "boots":
                template = templateIronBoots;
                break;
            default:
                return;
        }

        BufferedImage output = createMaskedTexture(blockTex, template);

        saveTexture(output, outputFolder, cache, itemPath + ".png");

        generatedItemIconTextures.put(info.item, new ResourceLocation(BlockArmor.MOD_ID,
                "textures/item/" + itemPath));
    }

    private BufferedImage createMaskedTexture(BufferedImage blockTex, BufferedImage template) {
        BufferedImage output = new BufferedImage(template.getWidth(), template.getHeight(), BufferedImage.TYPE_INT_ARGB);
        Graphics2D g2d = output.createGraphics();

        g2d.drawImage(template, 0, 0, null);

        g2d.setComposite(AlphaComposite.SrcIn);

        int tileCount = template.getWidth() / 16;
        for (int x = 0; x < tileCount; x++) {
            for (int y = 0; y < tileCount; y++) {
                g2d.drawImage(blockTex, x * 16, y * 16, 16, 16, null);
            }
        }

        g2d.dispose();
        return output;
    }

    private void saveTexture(BufferedImage image, Path outputFolder, CachedOutput cache, String fileName) throws IOException {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        ImageIO.write(image, "png", baos);
        byte[] bytes = baos.toByteArray();

        MessageDigest sha1;
        try {
            sha1 = MessageDigest.getInstance("SHA-1");
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException("SHA-1 algorithm not available", e);
        }
        byte[] hash = sha1.digest(bytes);

        Path outputPath = outputFolder.resolve(fileName);
        cache.writeIfNeeded(outputPath, bytes, HashCode.fromBytes(hash));
    }

    private static boolean isValidForArmor(Block block) {
        ResourceLocation loc = ForgeRegistries.BLOCKS.getKey(block);
        return loc != null && loc.getNamespace().equals("minecraft") &&
                block != Blocks.AIR && block.defaultBlockState().isSolidRender(null, null);
    }

    @Override
    public String getName() {
        return "BlockArmor Texture Provider";
    }
}

Please make it so it turns crafting tables, furnaces, smokers etc into model armor and item images also