/*
 * This file is part of Sponge, licensed under the MIT License (MIT).
 *
 * Copyright (c) SpongePowered <https://www.spongepowered.org>
 * Copyright (c) contributors
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */
package org.spongepowered.common.mixin.inventory.api;

import com.google.common.base.Preconditions;
import org.spongepowered.api.data.type.HandType;
import org.spongepowered.api.item.inventory.ArmorEquipable;
import org.spongepowered.api.item.inventory.ItemStack;
import org.spongepowered.api.item.inventory.equipment.EquipmentType;
import org.spongepowered.api.item.inventory.equipment.EquipmentTypes;
import org.spongepowered.asm.mixin.Mixin;
import org.spongepowered.common.bridge.world.inventory.InventoryBridge;
import org.spongepowered.common.inventory.adapter.InventoryAdapter;
import org.spongepowered.common.inventory.lens.Lens;
import org.spongepowered.common.inventory.lens.impl.comp.EquipmentInventoryLens;
import org.spongepowered.common.inventory.lens.impl.minecraft.PlayerInventoryLens;
import org.spongepowered.common.item.util.ItemStackUtil;

import java.util.Optional;

import javax.annotation.Nullable;
import net.minecraft.world.InteractionHand;
import net.minecraft.world.entity.EquipmentSlot;
import net.minecraft.world.entity.LivingEntity;
import net.minecraft.world.entity.Mob;
import net.minecraft.world.entity.decoration.ArmorStand;
import net.minecraft.world.entity.player.Player;

// All living implementors of ArmorEquipable
@Mixin({ArmorStand.class, Mob.class, Player.class})
public abstract class TraitMixin_ArmorEquipable_Inventory_API implements ArmorEquipable {

    // TODO can we implement canEquip?
    // We might want to allow plugins to set any item
    // but we should least expose checks if an item can be equipped normally

    @Override
    public boolean canEquip(final EquipmentType type) {
        return true;
    }

    @Override
    public boolean canEquip(final EquipmentType type, @Nullable final ItemStack equipment) {
        return true;
    }

    @Override
    public Optional<ItemStack> getEquipped(final EquipmentType type) {
        final InventoryAdapter inv = ((InventoryBridge) this).bridge$getAdapter();
        final Lens rootLens = inv.inventoryAdapter$getRootLens();
        final EquipmentInventoryLens lens;
        if (rootLens instanceof EquipmentInventoryLens) {
            lens = (EquipmentInventoryLens) rootLens;
        } else if (rootLens instanceof PlayerInventoryLens) {
            lens = ((PlayerInventoryLens) rootLens).getEquipmentLens();
        } else {
            throw new IllegalStateException("Unexpected lens for Equipable Inventory " + rootLens.getClass().getName());
        }
        return Optional.of(ItemStackUtil.fromNative(lens.getStack(inv.inventoryAdapter$getFabric(), ((EquipmentSlot) (Object) type).getIndex())));
    }

    @Override
    public boolean equip(final EquipmentType type, @Nullable final ItemStack equipment) {
        final InventoryAdapter inv = ((InventoryBridge) this).bridge$getAdapter();
        final EquipmentInventoryLens lens = (EquipmentInventoryLens) inv.inventoryAdapter$getRootLens();
        return lens.setStack(inv.inventoryAdapter$getFabric(), ((EquipmentSlot) (Object) type).getIndex(), ItemStackUtil.toNative(equipment));
    }

    @Override
    public ItemStack getItemInHand(HandType handType) {
        Preconditions.checkNotNull(handType);
        final net.minecraft.world.item.ItemStack nmsItem = ((LivingEntity) (Object)this).getItemInHand((InteractionHand) (Object) handType);
        return ItemStackUtil.fromNative(nmsItem);
    }

    @Override
    public void setItemInHand(HandType handType, @Nullable ItemStack itemInHand) {
        Preconditions.checkNotNull(handType);
        ((LivingEntity) (Object)this).setItemInHand((InteractionHand) (Object) handType, ItemStackUtil.toNative(itemInHand).copy());
    }

    @Override
    public ItemStack getHead() {
        return this.getEquipped(EquipmentTypes.HEAD).get();
    }

    @Override
    public void setHead(ItemStack head) {
        this.equip(EquipmentTypes.HEAD, head);
    }

    @Override
    public ItemStack getChest() {
        return this.getEquipped(EquipmentTypes.CHEST).get();
    }

    @Override
    public void setChest(ItemStack chest) {
        this.equip(EquipmentTypes.CHEST, chest);
    }

    @Override
    public ItemStack getLegs() {
        return this.getEquipped(EquipmentTypes.LEGS).get();
    }

    @Override
    public void setLegs(ItemStack legs) {
        this.equip(EquipmentTypes.LEGS, legs);
    }

    @Override
    public ItemStack getFeet() {
        return this.getEquipped(EquipmentTypes.FEET).get();
    }

    @Override
    public void setFeet(ItemStack feet) {
        this.equip(EquipmentTypes.FEET, feet);
    }
}

@@ -3228,496 +3228,72 @@
 nal 
-Lens rootLens = inv.inventoryAdapter$getRootLens();%0A        final EquipmentInventoryLens lens;%0A        if (rootLens instanceof EquipmentInventoryLens) %7B%0A            lens = (EquipmentInventoryLens) rootLens;%0A        %7D else if (rootLens instanceof PlayerInventoryLens) %7B%0A            lens = ((PlayerInventoryLens) rootLens).getEquipmentLens();%0A        %7D else %7B%0A            throw new IllegalStateException(%22Unexpected lens for Equipable Inventory %22 + rootLens.getClass().getName());%0A        %7D
+EquipmentInventoryLens lens = this.impl$equipmentInventory(inv);
 %0A   
@@ -3674,18 +3674,27 @@
  lens = 
-(E
+this.impl$e
 quipment
@@ -3706,47 +3706,12 @@
 tory
-Lens) inv.inventoryAdapter$getRootLens(
+(inv
 );%0A 
@@ -5356,11 +5356,575 @@
 );%0A    %7D
+%0A%0A    private EquipmentInventoryLens impl$equipmentInventory(final InventoryAdapter adapter) %7B%0A        final Lens rootLens = adapter.inventoryAdapter$getRootLens();%0A        if (rootLens instanceof EquipmentInventoryLens) %7B%0A            return (EquipmentInventoryLens) rootLens;%0A        %7D else if (rootLens instanceof PlayerInventoryLens) %7B%0A            return ((PlayerInventoryLens) rootLens).getEquipmentLens();%0A        %7D else %7B%0A            throw new IllegalStateException(%22Unexpected lens for Equipable Inventory %22 + rootLens.getClass().getName());%0A        %7D%0A    %7D
 %0A%7D%0A

Fix finding equipable inventory