package com.coasteen.armcontainers.create;

import com.simibubi.create.content.kinetics.mechanicalArm.ArmBlockEntity;
import com.simibubi.create.content.kinetics.mechanicalArm.ArmInteractionPoint;
import com.simibubi.create.content.kinetics.mechanicalArm.ArmInteractionPointType;
import net.minecraft.core.BlockPos;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.sounds.SoundEvent;
import net.minecraft.sounds.SoundEvents;
import net.minecraft.sounds.SoundSource;
import net.minecraft.world.item.ItemStack;
import net.minecraft.world.level.Level;
import net.minecraft.world.level.block.BarrelBlock;
import net.minecraft.world.level.block.ChestBlock;
import net.minecraft.world.level.block.entity.BarrelBlockEntity;
import net.minecraft.world.level.block.entity.BlockEntity;
import net.minecraft.world.level.block.entity.ChestBlockEntity;
import net.minecraft.world.level.block.entity.EnderChestBlockEntity;
import net.minecraft.world.level.block.entity.ShulkerBoxBlockEntity;
import net.minecraft.world.level.block.state.BlockState;
import net.minecraft.world.level.block.state.properties.ChestType;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

public class GenericContainerPoint extends ArmInteractionPoint {

    // Track active interactions per container position to allow overlapping arm access
    private static final Map<BlockPos, AtomicInteger> CONTAINER_REF_COUNTS = new ConcurrentHashMap<>();

    public GenericContainerPoint(ArmInteractionPointType type, Level level, BlockPos pos, BlockState state) {
        super(type, level, pos, state);
    }

    @Override
    public ItemStack extract(ArmBlockEntity arm, int slot, int amount, boolean simulate) {
        ItemStack result = super.extract(arm, slot, amount, simulate);
        if (!simulate && !result.isEmpty()) {
            animateContainer();
        }
        return result;
    }

    @Override
    public ItemStack extract(ArmBlockEntity arm, int amount, boolean simulate) {
        ItemStack result = super.extract(arm, amount, simulate);
        if (!simulate && !result.isEmpty()) {
            animateContainer();
        }
        return result;
    }

    @Override
    public ItemStack insert(ArmBlockEntity arm, ItemStack stack, boolean simulate) {
        int originalCount = stack.getCount();
        ItemStack remainder = super.insert(arm, stack, simulate);
        if (!simulate && remainder.getCount() < originalCount) {
            animateContainer();
        }
        return remainder;
    }

    private void animateContainer() {
        if (level == null || level.isClientSide()) return;

        BlockEntity be = level.getBlockEntity(pos);
        BlockState state = level.getBlockState(pos);

        AtomicInteger refCount = CONTAINER_REF_COUNTS.computeIfAbsent(pos, k -> new AtomicInteger(0));
        
        // If the container isn't currently held open by an arm, trigger the open animation
        if (refCount.getAndIncrement() == 0) {
            openContainer(be, state);
        }

        if (level instanceof ServerLevel serverLevel) {
            // Extended delay (~1.2 seconds) to keep the container open longer
            CompletableFuture.delayedExecutor(1200, TimeUnit.MILLISECONDS).execute(() -> {
                serverLevel.getServer().execute(() -> {
                    // Only close if all pending arm actions are complete
                    if (refCount.decrementAndGet() <= 0) {
                        CONTAINER_REF_COUNTS.remove(pos);
                        closeContainer(be);
                    }
                });
            });
        }
    }

    private void openContainer(BlockEntity be, BlockState state) {
        if (be instanceof ChestBlockEntity) {
            triggerChestEvent(state, 1);
            playSound(SoundEvents.CHEST_OPEN);
        } else if (be instanceof EnderChestBlockEntity) {
            level.blockEvent(pos, state.getBlock(), 1, 1);
            playSound(SoundEvents.ENDER_CHEST_OPEN);
        } else if (be instanceof ShulkerBoxBlockEntity) {
            level.blockEvent(pos, state.getBlock(), 1, 1);
            playSound(SoundEvents.SHULKER_BOX_OPEN);
        } else if (be instanceof BarrelBlockEntity) {
            if (state.hasProperty(BarrelBlock.OPEN) && !state.getValue(BarrelBlock.OPEN)) {
                level.setBlock(pos, state.setValue(BarrelBlock.OPEN, true), 3);
                playSound(SoundEvents.BARREL_OPEN);
            }
        }
    }

    private void closeContainer(BlockEntity be) {
        if (level == null || level.isClientSide()) return;
        BlockState currentState = level.getBlockState(pos);

        if (be instanceof ChestBlockEntity) {
            triggerChestEvent(currentState, 0);
            playSound(SoundEvents.CHEST_CLOSE);
        } else if (be instanceof EnderChestBlockEntity) {
            level.blockEvent(pos, currentState.getBlock(), 1, 0);
            playSound(SoundEvents.ENDER_CHEST_CLOSE);
        } else if (be instanceof ShulkerBoxBlockEntity) {
            level.blockEvent(pos, currentState.getBlock(), 1, 0);
            playSound(SoundEvents.SHULKER_BOX_CLOSE);
        } else if (be instanceof BarrelBlockEntity) {
            if (currentState.hasProperty(BarrelBlock.OPEN) && currentState.getValue(BarrelBlock.OPEN)) {
                level.setBlock(pos, currentState.setValue(BarrelBlock.OPEN, false), 3);
                playSound(SoundEvents.BARREL_CLOSE);
            }
        }
    }

    private void triggerChestEvent(BlockState state, int param) {
        level.blockEvent(pos, state.getBlock(), 1, param);
        if (state.hasProperty(ChestBlock.TYPE) && state.getValue(ChestBlock.TYPE) != ChestType.SINGLE) {
            BlockPos neighborPos = pos.relative(ChestBlock.getConnectedDirection(state));
            level.blockEvent(neighborPos, state.getBlock(), 1, param);
        }
    }

    private void playSound(SoundEvent sound) {
        level.playSound(null, pos, sound, SoundSource.BLOCKS, 0.5F, level.random.nextFloat() * 0.1F + 0.9F);
    }
}