package cn.hamster3.mc.auth.mixin;

import cn.hamster3.mc.auth.config.NeteaseServerConfig;
import cn.hamster3.mc.auth.data.NeteaseAuthResponse;
import com.google.common.primitives.Longs;
import com.google.gson.JsonObject;
import com.velocitypowered.api.proxy.crypto.IdentifiedKey;
import com.velocitypowered.api.util.GameProfile;
import com.velocitypowered.proxy.VelocityServer;
import com.velocitypowered.proxy.connection.MinecraftConnection;
import com.velocitypowered.proxy.connection.client.AuthSessionHandler;
import com.velocitypowered.proxy.connection.client.InitialLoginSessionHandler;
import com.velocitypowered.proxy.connection.client.LoginInboundConnection;
import com.velocitypowered.proxy.crypto.IdentifiedKeyImpl;
import com.velocitypowered.proxy.protocol.StateRegistry;
import com.velocitypowered.proxy.protocol.packet.EncryptionResponsePacket;
import com.velocitypowered.proxy.protocol.packet.ServerLoginPacket;
import net.kyori.adventure.text.Component;
import net.kyori.adventure.text.format.NamedTextColor;
import org.apache.logging.log4j.Logger;
import org.spongepowered.asm.mixin.Final;
import org.spongepowered.asm.mixin.Mixin;
import org.spongepowered.asm.mixin.Shadow;
import org.spongepowered.asm.mixin.injection.At;
import org.spongepowered.asm.mixin.injection.Inject;
import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;

import java.net.InetSocketAddress;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.security.GeneralSecurityException;
import java.security.KeyPair;
import java.security.MessageDigest;
import java.util.ArrayList;

import static com.velocitypowered.proxy.VelocityServer.GENERAL_GSON;
import static com.velocitypowered.proxy.crypto.EncryptionUtils.decryptRsa;
import static com.velocitypowered.proxy.crypto.EncryptionUtils.generateServerId;

@Mixin(InitialLoginSessionHandler.class)
public abstract class InitialLoginSessionHandlerMixin {
    @Shadow
    protected abstract void assertState(InitialLoginSessionHandler.LoginState expectedState);

    @Shadow
    private InitialLoginSessionHandler.LoginState currentState;

    @Shadow
    private ServerLoginPacket login;

    @Shadow
    private byte[] verify;

    @Shadow
    @Final
    private VelocityServer server;

    @Shadow
    @Final
    private LoginInboundConnection inbound;

    @Shadow
    @Final
    private MinecraftConnection mcConnection;

    @Shadow
    @Final
    private static Logger logger;

    @Inject(at = @At(value = "HEAD"), method = "handle(Lcom/velocitypowered/proxy/protocol/packet/EncryptionResponsePacket;)Z", cancellable = true)
    public void injectHandle(EncryptionResponsePacket packet, CallbackInfoReturnable<Boolean> cir) {
        assertState(InitialLoginSessionHandler.LoginState.ENCRYPTION_REQUEST_SENT);
        currentState = InitialLoginSessionHandler.LoginState.ENCRYPTION_RESPONSE_RECEIVED;
        ServerLoginPacket login = this.login;
        if (login == null) {
            throw new IllegalStateException("No ServerLogin packet received yet.");
        }
        if (verify.length == 0) {
            throw new IllegalStateException("No EncryptionRequest packet sent yet.");
        }
        try {
            KeyPair serverKeyPair = server.getServerKeyPair();
            if (inbound.getIdentifiedKey() != null) {
                IdentifiedKey playerKey = inbound.getIdentifiedKey();
                if (!playerKey.verifyDataSignature(packet.getVerifyToken(), verify,
                        Longs.toByteArray(packet.getSalt()))) {
                    throw new IllegalStateException("Invalid client public signature.");
                }
            } else {
                byte[] decryptedVerifyToken = decryptRsa(serverKeyPair, packet.getVerifyToken());
                if (!MessageDigest.isEqual(verify, decryptedVerifyToken)) {
                    throw new IllegalStateException("Unable to successfully decrypt the verification token.");
                }
            }
            final byte[] decryptedSharedSecret = decryptRsa(serverKeyPair, packet.getSharedSecret());
            final String serverId = generateServerId(decryptedSharedSecret, serverKeyPair.getPublic());
            final String playerIp = ((InetSocketAddress) mcConnection.getRemoteAddress()).getHostString();

            JsonObject data = new JsonObject();
            data.addProperty("username", login.getUsername());
            data.addProperty("serverId", serverId);
            data.addProperty("gameID", NeteaseServerConfig.getGameId());

            HttpRequest httpRequest = HttpRequest.newBuilder()
                    .headers("Content-Type", "application/json")
                    .uri(URI.create(NeteaseServerConfig.getAuthUrl()))
                    .POST(HttpRequest.BodyPublishers.ofString(data.toString()))
                    .build();

            @SuppressWarnings("resource")
            final HttpClient httpClient = server.createHttpClient();
            httpClient.sendAsync(httpRequest, HttpResponse.BodyHandlers.ofString())
                    .whenCompleteAsync((response, throwable) -> {
                        if (mcConnection.isClosed()) {
                            // The player disconnected after we authenticated them.
                            return;
                        }

                        if (throwable != null) {
                            logger.error("Unable to authenticate player", throwable);
                            inbound.disconnect(Component.translatable("multiplayer.disconnect.authservers_down"));
                            return;
                        }

                        // Go ahead and enable encryption. Once the client sends EncryptionResponse, encryption
                        // is enabled.
                        try {
                            mcConnection.enableEncryption(decryptedSharedSecret);
                        } catch (GeneralSecurityException e) {
                            logger.error("Unable to enable encryption for connection", e);
                            // At this point, the connection is encrypted, but something's wrong on our side and
                            // we can't do anything about it.
                            mcConnection.close(true);
                            return;
                        }

                        try {
                            if (response.statusCode() == 200) {
                                final NeteaseAuthResponse authResponse = GENERAL_GSON.fromJson(response.body(), NeteaseAuthResponse.class);
                                if (authResponse.getCode() != 0) {
                                    inbound.disconnect(Component.translatable("multiplayer.disconnect.authservers_down"));
                                    logger.error("Error authenticating {} with netease", login.getUsername());
                                } else {
                                    NeteaseAuthResponse.ResponseEntity entity = authResponse.getEntity();
                                    if (entity.getName() == null || entity.getName().isEmpty()) {
                                        entity.setName(login.getUsername());
                                    }
                                    if (entity.getProperties() == null) {
                                        entity.setProperties(new ArrayList<>());
                                    }
                                    if (entity.getId() == null) {
                                        inbound.disconnect(
                                                Component.translatable("velocity.error.online-mode-only", NamedTextColor.RED)
                                        );
                                    } else {
                                        GameProfile profile = new GameProfile(entity.getId(), entity.getName(), entity.getProperties());
                                        // Not so fast, now we verify the public key for 1.19.1+
                                        if (inbound.getIdentifiedKey() != null
                                                && inbound.getIdentifiedKey().getKeyRevision() == IdentifiedKey.Revision.LINKED_V2
                                                && inbound.getIdentifiedKey() instanceof final IdentifiedKeyImpl key) {
                                            if (!key.internalAddHolder(profile.getId())) {
                                                inbound.disconnect(
                                                        Component.translatable("multiplayer.disconnect.invalid_public_key"));
                                            }
                                        }
                                        // All went well, initialize the session.
                                        mcConnection.setActiveSessionHandler(StateRegistry.LOGIN,
                                                new AuthSessionHandler(server, inbound, profile, true));
                                    }
                                }
                            } else if (response.statusCode() == 204) {
                                // Apparently an offline-mode user logged onto this online-mode proxy.
                                inbound.disconnect(
                                        Component.translatable("velocity.error.online-mode-only", NamedTextColor.RED));
                            } else {
                                // Something else went wrong
                                logger.error(
                                        "Got an unexpected error code {} whilst contacting Mojang to log in {} ({})",
                                        response.statusCode(), login.getUsername(), playerIp);
                                inbound.disconnect(Component.translatable("multiplayer.disconnect.authservers_down"));
                            }
                        } catch (Exception e) {
                            logger.error("Got an unexpected error", e);
                        }
                    }, mcConnection.eventLoop())
                    .thenRun(() -> {
                        //noinspection ConstantValue
                        if (httpClient instanceof final AutoCloseable closeable) {
                            try {
                                closeable.close();
                            } catch (Exception e) {
                                // In Java 21, the HttpClient does not throw any Exception
                                // when trying to clean its resources, so this should not happen
                                logger.error("An unknown error occurred while trying to close an HttpClient", e);
                            }
                        }
                    });
        } catch (GeneralSecurityException e) {
            logger.error("Unable to enable encryption", e);
            mcConnection.close(true);
        }
        cir.cancel();
        cir.setReturnValue(true);
    }
}
