package net.minestom.vanilla.loot;

import net.kyori.adventure.key.Key;
import net.kyori.adventure.nbt.CompoundBinaryTag;
import net.kyori.adventure.nbt.IntBinaryTag;
import net.kyori.adventure.nbt.NumberBinaryTag;
import net.minestom.server.codec.Codec;
import net.minestom.server.codec.StructCodec;
import net.minestom.server.item.enchant.LevelBasedValue;
import net.minestom.server.registry.DynamicRegistry;
import net.minestom.vanilla.loot.util.nbt.NBTPath;
import net.minestom.vanilla.loot.util.nbt.NBTReference;
import org.jetbrains.annotations.NotNull;

import java.util.List;
import java.util.Random;

/**
 * Generates numbers based on provided loot contexts.
 */
@SuppressWarnings("UnstableApiUsage")
public interface LootNumber {

    @NotNull Codec<LootNumber> CODEC = Codec.DOUBLE.<LootNumber>transform(Constant::new, a -> ((Constant) a).value()).orElse(Codec.RegistryTaggedUnion(registries -> {
        class Holder {
            static final @NotNull DynamicRegistry<StructCodec<? extends LootNumber>> CODEC = createDefaultRegistry();
        }
        return Holder.CODEC;
    }, LootNumber::codec, "type"));

    static @NotNull DynamicRegistry<StructCodec<? extends LootNumber>> createDefaultRegistry() {
        final DynamicRegistry<StructCodec<? extends LootNumber>> registry = DynamicRegistry.create(Key.key("loot_numbers"));
        registry.register("binomial", Binomial.CODEC);
        registry.register("constant", Constant.CODEC);
        registry.register("enchantment_level", EnchantmentLevel.CODEC);
        registry.register("score", Score.CODEC);
        registry.register("storage", Storage.CODEC);
        registry.register("uniform", Uniform.CODEC);
        return registry;
    }

    /**
     * Generates an integer depending on the information in the provided context.<br>
     * This is an explicitly impure method—it depends on state outside the given context.
     * @param context the context object, to use if required
     * @return the integer generated by this loot number for the provided context
     */
    int getInt(@NotNull LootContext context);

    /**
     * Generates a double depending on the information in the provided context.<br>
     * This is an explicitly impure method—it depends on state outside the given context.
     * @param context the context object, to use if required
     * @return the double generated by this loot number for the provided context
     */
    double getDouble(@NotNull LootContext context);

    /**
     * @return the codec that can encode this number
     */
    @NotNull StructCodec<? extends LootNumber> codec();

    record Binomial(@NotNull LootNumber trials, @NotNull LootNumber probability) implements LootNumber {
        public static final @NotNull StructCodec<Binomial> CODEC = StructCodec.struct(
                "n", LootNumber.CODEC, Binomial::trials,
                "p", LootNumber.CODEC, Binomial::probability,
                Binomial::new
        );

        @Override
        public int getInt(@NotNull LootContext context) {
            int trials = trials().getInt(context);
            double probability = probability().getDouble(context);
            Random random = context.require(LootContext.RANDOM);

            int successes = 0;
            for (int trial = 0; trial < trials; trial++) {
                if (random.nextDouble() < probability) {
                    successes++;
                }
            }
            return successes;
        }

        @Override
        public double getDouble(@NotNull LootContext context) {
            return getInt(context);
        }

        @Override
        public @NotNull StructCodec<? extends LootNumber> codec() {
            return CODEC;
        }
    }

    record Constant(@NotNull Double value) implements LootNumber {
        public static final @NotNull StructCodec<Constant> CODEC = StructCodec.struct(
                "value", Codec.DOUBLE, Constant::value,
                Constant::new
        );

        @Override
        public int getInt(@NotNull LootContext context) {
            return value.intValue();
        }

        @Override
        public double getDouble(@NotNull LootContext context) {
            return value;
        }

        @Override
        public @NotNull StructCodec<? extends LootNumber> codec() {
            return CODEC;
        }
    }

    record EnchantmentLevel(@NotNull LevelBasedValue amount) implements LootNumber {
        public static final @NotNull StructCodec<EnchantmentLevel> CODEC = StructCodec.struct(
                "amount", LevelBasedValue.CODEC, EnchantmentLevel::amount,
                EnchantmentLevel::new
        );

        @Override
        public int getInt(@NotNull LootContext context) {
            return (int) Math.round(getDouble(context));
        }

        @Override
        public double getDouble(@NotNull LootContext context) {
            return amount.calc(context.require(LootContext.ENCHANTMENT_LEVEL));
        }

        @Override
        public @NotNull StructCodec<? extends LootNumber> codec() {
            return CODEC;
        }
    }

    record Score(@NotNull LootScore target, @NotNull String objective, double scale) implements LootNumber {
        public static final @NotNull StructCodec<Score> CODEC = StructCodec.struct(
                "target", LootScore.CODEC, Score::target,
                "score", Codec.STRING, Score::objective,
                "scale", Codec.DOUBLE, Score::scale,
                Score::new
        );

        @Override
        public int getInt(@NotNull LootContext context) {
            return (int) Math.round(getDouble(context));
        }

        @Override
        public double getDouble(@NotNull LootContext context) {
            var score = target.apply(context).apply(objective);

            return score != null ? score * scale : 0;
        }

        @Override
        public @NotNull StructCodec<? extends LootNumber> codec() {
            return CODEC;
        }
    }

    record Storage(@NotNull Key storage, @NotNull NBTPath path) implements LootNumber {
        public static final @NotNull StructCodec<Storage> CODEC = StructCodec.struct(
                "storage", Codec.KEY, Storage::storage,
                "path", NBTPath.CODEC, Storage::path,
                Storage::new
        );

        @Override
        public int getInt(@NotNull LootContext context) {
            return get(context).intValue();
        }

        @Override
        public double getDouble(@NotNull LootContext context) {
            return get(context).doubleValue();
        }

        private NumberBinaryTag get(@NotNull LootContext context) {
            if (true) throw new UnsupportedOperationException("TODO: Implement entity scores (Entity entity -> String objective -> @Nullable Integer)");
            CompoundBinaryTag compound = null;

            List<NBTReference> refs = path.get(compound != null ? compound : CompoundBinaryTag.empty());
            if (refs.size() != 1) return IntBinaryTag.intBinaryTag(0);

            if (refs.getFirst().get() instanceof NumberBinaryTag number) {
                return number;
            } else {
                return IntBinaryTag.intBinaryTag(0);
            }
        }

        @Override
        public @NotNull StructCodec<? extends LootNumber> codec() {
            return CODEC;
        }
    }

    record Uniform(@NotNull LootNumber min, @NotNull LootNumber max) implements LootNumber {
        public static final @NotNull StructCodec<Uniform> CODEC = StructCodec.struct(
                "min", LootNumber.CODEC, Uniform::min,
                "max", LootNumber.CODEC, Uniform::max,
                Uniform::new
        );

        @Override
        public int getInt(@NotNull LootContext context) {
            return context.require(LootContext.RANDOM).nextInt(min().getInt(context), max().getInt(context) + 1);
        }

        @Override
        public double getDouble(@NotNull LootContext context) {
            return context.require(LootContext.RANDOM).nextDouble(min().getDouble(context), max().getDouble(context));
        }

        @Override
        public @NotNull StructCodec<? extends LootNumber> codec() {
            return CODEC;
        }
    }

}
