package de.winniepat.citrus.cosmetics.halo;

import net.minecraft.client.network.AbstractClientPlayerEntity;
import net.minecraft.client.render.RenderLayer;
import net.minecraft.client.render.VertexConsumerProvider;
import net.minecraft.client.render.entity.feature.FeatureRenderer;
import net.minecraft.client.render.entity.feature.FeatureRendererContext;
import net.minecraft.client.render.entity.model.PlayerEntityModel;
import net.minecraft.client.util.math.MatrixStack;
import net.minecraft.entity.player.PlayerEntity;
import net.minecraft.util.Identifier;
import net.minecraft.util.math.RotationAxis;
import software.bernie.geckolib.renderer.GeoObjectRenderer;

import java.util.WeakHashMap;

public class HaloFeatureRenderer extends FeatureRenderer<AbstractClientPlayerEntity, PlayerEntityModel<AbstractClientPlayerEntity>> {
    private final GeoObjectRenderer<CosmeticHalo> haloRenderer;
    private static final WeakHashMap<PlayerEntity, CosmeticHalo> HALO_CACHE = new WeakHashMap<>();

    public HaloFeatureRenderer(FeatureRendererContext<AbstractClientPlayerEntity, PlayerEntityModel<AbstractClientPlayerEntity>> context) {
        super(context);
        this.haloRenderer = new GeoObjectRenderer<>(new HaloModel());
    }

    @Override
    public void render(MatrixStack matrices, VertexConsumerProvider vertexConsumers, int light, AbstractClientPlayerEntity player, float limbAngle, float limbDistance, float tickDelta, float animationProgress, float headYaw, float headPitch) {

        CosmeticHalo halo = HALO_CACHE.computeIfAbsent(player, CosmeticHalo::new);
        boolean shouldRender = halo.shouldRender();
        HaloType type = halo.getHaloType();
        if (!shouldRender) return;

        Identifier texture = this.haloRenderer.getGeoModel().getTextureResource(halo);
        if (texture == null) return;

        matrices.push();

        matrices.multiply(RotationAxis.POSITIVE_Y.rotationDegrees(headYaw * 0.3f));
        matrices.multiply(RotationAxis.POSITIVE_X.rotationDegrees(headPitch * 0.3f));

        matrices.translate(0, -2.0, 0);

        matrices.translate(type.getPosX(), type.getPosY(), type.getPosZ());

        float scale = type.getScale();
        matrices.scale(scale, scale, scale);

        matrices.multiply(RotationAxis.POSITIVE_X.rotationDegrees(180f));

        try {
            this.haloRenderer.render(
                    matrices, halo, vertexConsumers,
                    RenderLayer.getEntityTranslucent(texture),
                    vertexConsumers.getBuffer(RenderLayer.getEntityTranslucent(texture)),
                    0xF000F0,
                    tickDelta
            );
        } catch (Exception e) {
            System.err.println("Failed to render halo: " + e.getMessage());
        } finally {
            matrices.pop();
        }
    }

    public static void invalidateHaloCache(PlayerEntity player) {
        HALO_CACHE.remove(player);
    }
}