package io.papermc.paper.util;

import com.google.common.base.Charsets;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.mojang.datafixers.util.Pair;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.Map;
import java.util.Set;
import net.fabricmc.mapping.tree.ClassDef;
import net.fabricmc.mapping.tree.MethodDef;
import net.fabricmc.mapping.tree.TinyMappingFactory;
import net.fabricmc.mapping.tree.TinyTree;
import org.checkerframework.checker.nullness.qual.NonNull;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.checkerframework.framework.qual.DefaultQualifier;

@DefaultQualifier(NonNull.class)
public enum ObfHelper {
    INSTANCE;

    public static final String MOJANG_PLUS_YARN_NAMESPACE = "mojang+yarn";
    public static final String SPIGOT_NAMESPACE = "spigot";

    private final @Nullable Map<String, ClassMapping> mappingsByObfName;
    private final @Nullable Map<String, ClassMapping> mappingsByMojangName;

    ObfHelper() {
        final @Nullable Set<ClassMapping> maps = loadMappingsIfPresent();
        if (maps != null) {
            this.mappingsByObfName = maps.stream().collect(ImmutableMap.toImmutableMap(ClassMapping::obfName, map -> map));
            this.mappingsByMojangName = maps.stream().collect(ImmutableMap.toImmutableMap(ClassMapping::mojangName, map -> map));
        } else {
            this.mappingsByObfName = null;
            this.mappingsByMojangName = null;
        }
    }

    public @Nullable Map<String, ClassMapping> mappingsByObfName() {
        return this.mappingsByObfName;
    }

    public @Nullable Map<String, ClassMapping> mappingsByMojangName() {
        return this.mappingsByMojangName;
    }

    /**
     * Attempts to get the obf name for a given class by its Mojang name. Will
     * return the input string if mappings are not present.
     *
     * @param fullyQualifiedMojangName fully qualified class name (dotted)
     * @return mapped or original fully qualified (dotted) class name
     */
    public String reobfClassName(final String fullyQualifiedMojangName) {
        if (this.mappingsByMojangName == null) {
            return fullyQualifiedMojangName;
        }

        final ClassMapping map = this.mappingsByMojangName.get(fullyQualifiedMojangName);
        if (map == null) {
            return fullyQualifiedMojangName;
        }

        return map.obfName();
    }

    /**
     * Attempts to get the Mojang name for a given class by its obf name. Will
     * return the input string if mappings are not present.
     *
     * @param fullyQualifiedObfName fully qualified class name (dotted)
     * @return mapped or original fully qualified (dotted) class name
     */
    public String deobfClassName(final String fullyQualifiedObfName) {
        if (this.mappingsByObfName == null) {
            return fullyQualifiedObfName;
        }

        final ClassMapping map = this.mappingsByObfName.get(fullyQualifiedObfName);
        if (map == null) {
            return fullyQualifiedObfName;
        }

        return map.mojangName();
    }

    private static @Nullable Set<ClassMapping> loadMappingsIfPresent() {
        try (final @Nullable InputStream mappingsInputStream = ObfHelper.class.getClassLoader().getResourceAsStream("META-INF/mappings/reobf.tiny")) {
            if (mappingsInputStream == null) {
                return null;
            }
            final TinyTree tree = TinyMappingFactory.loadWithDetection(new BufferedReader(new InputStreamReader(mappingsInputStream, Charsets.UTF_8)));
            final var builder = ImmutableSet.<ClassMapping>builder();

            for (final ClassDef classDef : tree.getClasses()) {
                final String obfClassName = classDef.getName(SPIGOT_NAMESPACE).replace('/', '.');
                final var methodMappings = ImmutableMap.<Pair<String, String>, MethodMapping>builder();

                for (final MethodDef methodDef : classDef.getMethods()) {
                    final MethodMapping method = new MethodMapping(
                        methodDef.getName(SPIGOT_NAMESPACE),
                        methodDef.getName(MOJANG_PLUS_YARN_NAMESPACE),
                        methodDef.getDescriptor(SPIGOT_NAMESPACE)
                    );
                    methodMappings.put(
                        new Pair<>(method.obfName(), method.descriptor()),
                        method
                    );
                }

                final ClassMapping map = new ClassMapping(
                    obfClassName,
                    classDef.getName(MOJANG_PLUS_YARN_NAMESPACE).replace('/', '.'),
                    methodMappings.build()
                );
                builder.add(map);
            }

            return builder.build();
        } catch (final IOException ex) {
            System.err.println("Failed to load mappings for stacktrace deobfuscation.");
            ex.printStackTrace();
            return null;
        }
    }

    public record ClassMapping(
        String obfName,
        String mojangName,
        Map<Pair<String, String>, MethodMapping> methodMappings
    ) {}

    public record MethodMapping(
        String obfName,
        String mojangName,
        String descriptor
    ) {}
}
