package io.gitlab.jfronny.betterwhitelist.server;

import com.mojang.authlib.GameProfile;
import com.mojang.brigadier.context.CommandContext;
import io.gitlab.jfronny.betterwhitelist.BetterWhitelist;
import io.gitlab.jfronny.betterwhitelist.server.mixin.ServerLoginNetworkHandlerAccessor;
import io.gitlab.jfronny.commons.StringFormatter;
import io.gitlab.jfronny.muscript.compiler.Parser;
import io.gitlab.jfronny.muscript.data.Scope;
import io.gitlab.jfronny.muscript.data.Script;
import io.gitlab.jfronny.muscript.data.dynamic.DNull;
import io.gitlab.jfronny.muscript.data.dynamic.Dynamic;
import io.gitlab.jfronny.muscript.error.LocationalException;
import net.fabricmc.api.DedicatedServerModInitializer;
import net.fabricmc.fabric.api.command.v2.CommandRegistrationCallback;
import net.fabricmc.fabric.api.networking.v1.*;
import net.fabricmc.loader.api.FabricLoader;
import net.minecraft.network.PacketByteBuf;
import net.minecraft.server.command.ServerCommandSource;
import net.minecraft.server.network.ServerLoginNetworkHandler;
import net.minecraft.text.Text;
import net.minecraft.util.Util;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.*;
import java.util.concurrent.FutureTask;

import static net.minecraft.server.command.CommandManager.literal;

public class BetterWhitelistServer implements DedicatedServerModInitializer {
    private String scriptSource;
    private Script script;
    private final Map<GameProfile, Challenge> challenges = new HashMap<>();

    private static class Challenge {
        public final Queue<Dynamic<?>> responses = new LinkedList<>();
        public final GameProfile profile;
        public PacketSender response;

        public Challenge(GameProfile profile, PacketSender response) {
            this.profile = profile;
            this.response = response;
        }
    }

    @Override
    public void onInitializeServer() {
        CommandRegistrationCallback.EVENT.register((dispatcher, registryAccess, environment) -> {
            dispatcher.register(
                    literal(BetterWhitelist.MOD_ID)
                            .requires(source -> source.hasPermissionLevel(4))
                            .executes(this::printVersion)
                            .then(literal("reload").executes(this::reloadScript))
            );
        });

        ServerLoginConnectionEvents.QUERY_START.register((handler, server, sender, synchronizer) -> {
            GameProfile gp = profile(handler);
            challenges.put(gp, new Challenge(gp, sender));
            Challenge challenge = challenges.get(gp);

            Scope fork = BetterWhitelist.SCOPE.fork();
            fork.set("assert", args -> {
                if (args.size() != 1 && args.size() != 2) throw new IllegalArgumentException("Invalid number of arguments for assert: expected 1 or 2 but got " + args.size());
                if (!args.get(0).asBool().getValue()) throw new AssertFail(args.size() > 1 ? args.get(1).asString().getValue() : "Failed Whitelist Check");
                return new DNull();
            }).set("challenge", args -> {
                if (args.size() == 0) throw new IllegalArgumentException("Invalid number of arguments for challenge: expected 1 or more but got 0");
                PacketByteBuf buf = PacketByteBufs.create();
                String challengeString = Dynamic.serialize(args.get(0).asCallable());
                BetterWhitelist.LOG.info("Sending challenge to " + gp.getName() + ": " + challengeString);
                buf.writeString(challengeString);
                List<Dynamic<?>> params = args.getValue().subList(1, args.size());
                buf.writeInt(params.size());
                params.forEach(p -> buf.writeString(Dynamic.serialize(p)));
                challenge.response.sendPacket(BetterWhitelist.LOGIN_CHANNEL, buf);
                int i = 0;
                while (challenge.responses.isEmpty() && handler.isConnectionOpen() && i++ < 100) sleep(10);
                if (!challenge.responses.isEmpty()) return challenge.responses.remove();
                else throw new AssertFail("Took too long to respond");
            });

            FutureTask<?> future = new FutureTask<>(() -> {
                try {
                    try {
                        script.run(fork);
                        BetterWhitelist.LOG.info("Completed challenge for " + gp.getName());
                    } catch (LocationalException le) {
                        for (Throwable t = le; t != null; t = t.getCause()) {
                            if (t instanceof AssertFail af) throw af;
                        }
                        throw le;
                    }
                } catch (AssertFail fail) {
                    BetterWhitelist.LOG.warn("Failed challenge for " + gp.getName() + ": " + fail.getMessage());
                    if (handler.isConnectionOpen()) handler.disconnect(Text.literal(fail.getMessage()));
                } catch (Throwable t) {
                    BetterWhitelist.LOG.error("Something went wrong while trying to execute a challenge\n"
                            + StringFormatter.toString(t, e ->
                            e instanceof LocationalException le
                                    ? le.asPrintable(scriptSource).toString()
                                    : e.toString()
                    ));
                    if (handler.isConnectionOpen()) handler.disconnect(Text.literal("Something went wrong"));
                }
                challenges.remove(gp);
                return null;
            });

            Util.getMainWorkerExecutor().execute(future);
            synchronizer.waitFor(future);
        });

        ServerLoginNetworking.registerGlobalReceiver(BetterWhitelist.LOGIN_CHANNEL, (server, handler, understood, buf, synchronizer, responseSender) -> {
            if (!understood) {
                handler.disconnect(Text.literal("This server requires better-whitelist to be installed"));
                return;
            }
            try {
                Challenge ch = challenges.get(profile(handler));
                ch.response = responseSender;
                String response = buf.readString();
                BetterWhitelist.LOG.info("Got response from " + ch.profile.getName() + ": " + response);
                ch.responses.add(Dynamic.deserialize(response));
            } catch (Throwable t) {
                BetterWhitelist.LOG.error("Failed login", t);
                handler.disconnect(Text.literal("Invalid dynamic"));
            }
        });

        try {
            reloadScript();
        } catch (IOException e) {
            throw new RuntimeException("Could not load whitelist script", e);
        }
    }

    private GameProfile profile(ServerLoginNetworkHandler handler) {
        GameProfile gp = ((ServerLoginNetworkHandlerAccessor) handler).getProfile();
        if (gp == null) throw new NullPointerException("Missing GameProfile");
        return gp;
    }

    private static void sleep(long millis) {
        try {
            Thread.sleep(millis);
        } catch (InterruptedException ignored) {
        }
    }

    private static class AssertFail extends RuntimeException {
        public AssertFail(String message) {
            super(message);
        }
    }

    private int printVersion(CommandContext<ServerCommandSource> context) {
        context.getSource().sendMessage(Text.literal("Loaded " + BetterWhitelist.MOD_METADATA.getName() + " " + BetterWhitelist.MOD_METADATA.getVersion()));
        return 1;
    }

    private int reloadScript(CommandContext<ServerCommandSource> context) {
        try {
            reloadScript();
            context.getSource().sendMessage(Text.literal("Successfully reloaded script"));
        } catch (Throwable t) {
            BetterWhitelist.LOG.error("Could not reload script", t);
            context.getSource().sendError(Text.literal("Could not reload script, check server log for details"));
        }
        return 1;
    }

    private void reloadScript() throws IOException {
        Path scriptPath = FabricLoader.getInstance()
                .getConfigDir()
                .resolve(BetterWhitelist.MOD_ID + ".mu");
        if (!Files.exists(scriptPath)) Files.writeString(scriptPath, """
                // Use this method to execute a closure on the client and get back the result
                clientVersion = challenge({ ->
                  // Note that closures sent to the client do not have access to things you declare elsewhere
                  mod('better-whitelist').version
                })
                
                println("You can, of course, use println-debugging")
                
                // You have access to the same methods on the server as you do on the client
                // You may use the assert method to short-circuit if you encounter a case where the client should not be allowed access
                // Assert can also have a second argument for the message to send if the assertion fails
                assert(mod('better-whitelist').version == clientVersion, 'You have the wrong mod version')
                
                // you can also send server-evaluated parameters with your challenge
                assert(challenge({ arg ->
                  arg::allMatch({ v -> mods::values()::anyMatch({ m -> v.id == m.id & v.version == m.version }) })
                }, mods::values()::filter({ v -> v.environment != 'server' })::map({ v -> { id = v.id, version = v.version } })))
                """);
        String s = Files.readString(scriptPath);
        this.script = Parser.parseScript(s);
        this.scriptSource = s;
    }
}