/*
 *   This file is a part of the SteamWar software.
 *
 *   Copyright (C) 2025 SteamWar.de-Serverteam
 *
 *   This program is free software: you can redistribute it and/or modify
 *   it under the terms of the GNU Affero General Public License as published by
 *   the Free Software Foundation, either version 3 of the License, or
 *   (at your option) any later version.
 *
 *   This program is distributed in the hope that it will be useful,
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 *   GNU Affero General Public License for more details.
 *
 *   You should have received a copy of the GNU Affero General Public License
 *   along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

package de.steamwar.core;

import com.comphenix.tinyprotocol.TinyProtocol;
import de.steamwar.Reflection;
import de.steamwar.sql.internal.Statement;
import io.netty.channel.ChannelFuture;
import net.minecraft.server.MinecraftServer;
import net.minecraft.server.network.ServerConnectionListener;
import org.bukkit.Bukkit;
import org.bukkit.craftbukkit.CraftWorld;
import org.eclipse.openj9.criu.CRIUSupport;
import org.eclipse.openj9.criu.JVMCRIUException;
import sun.misc.Signal;

import java.io.DataInputStream;
import java.io.File;
import java.io.IOException;
import java.net.InetAddress;
import java.nio.file.FileSystems;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Comparator;
import java.util.List;
import java.util.logging.Level;
import java.util.stream.Stream;

class CheckpointUtilsJ9 {
    private CheckpointUtilsJ9() {
    }

    static void signalHandler() {
        Signal.handle(new Signal("USR1"), signal -> Bukkit.getScheduler().runTask(Core.getInstance(), CheckpointUtils::freeze));
    }

    static void freeze() {
        String checkpointFile = System.getProperty("checkpoint");
        if (!CRIUSupport.isCheckpointAllowed() || checkpointFile == null) {
            Bukkit.shutdown();
            return;
        }

        Bukkit.getOnlinePlayers().forEach(player -> player.kickPlayer(null));

        List<?> networkManagers = TinyProtocol.instance.networkManagers;
        if (!Bukkit.getOnlinePlayers().isEmpty() || !networkManagers.isEmpty()) {
            Core.getInstance().getLogger().log(Level.INFO, "Waiting for players to disconnect for checkpointing");
            Bukkit.getScheduler().runTaskLater(Core.getInstance(), CheckpointUtils::freeze, 1);
            return;
        }

        Path path = FileSystems.getDefault().getPath(checkpointFile);

        try {
            freezeInternal(path);
        } catch (Exception e) {
            String message = e.getMessage() != null ? e.getMessage() : "";
            if (message.contains("Connected TCP socket")) {
                Core.getInstance().getLogger().log(Level.INFO, "Connected TCP socket, waiting for checkpointing");
                Bukkit.getScheduler().runTaskLater(Core.getInstance(), CheckpointUtils::freeze, 1);
                return;
            }

            Bukkit.shutdown();

            if (!message.contains("Can't dump ghost file") && !message.contains("Can't create link remap")) // File/Jar has been updated
                throw new SecurityException(e);
        } finally {
            // Delete checkpoint
            try (Stream<Path> stream = Files.walk(path)) {
                stream.sorted(Comparator.reverseOrder()).map(Path::toFile).forEach(File::delete);
            } catch (IOException e) {
                //ignore
            }
        }
    }


    private static final Reflection.Field<List> channelFutures = Reflection.getField(ServerConnectionListener.class, List.class, 0, ChannelFuture.class);

    private static void freezeInternal(Path path) throws Exception {
        Bukkit.getPluginManager().callEvent(new CRIUSleepEvent());

        Bukkit.getWorlds().forEach(world -> {
            ((CraftWorld) world).getHandle().save(null, true, false);
        });
        Statement.closeAll();

        // Close socket
        ServerConnectionListener serverConnection = MinecraftServer.getServer().getConnection();
        List<?> channels = channelFutures.get(serverConnection);
        for (Object future : channels) {
            ((ChannelFuture) future).channel().close().syncUninterruptibly();
        }
        channels.clear();

        System.runFinalization();
        System.gc();

        // Do the checkpoint
        path.toFile().mkdirs();
        CRIUSupport criu = new CRIUSupport(path);
        criu.setGhostFileLimit(10L * 1024L * 1024L);
        criu.setAutoDedup(true);
        criu.setFileLocks(true);
        criu.setShellJob(true);
        criu.setLogFile("criu.log");
        try {
            criu.checkpointJVM();
        } catch (JVMCRIUException e) {
            Path logfile = path.resolve("criu.log");
            if (logfile.toFile().exists()) {
                throw new IllegalStateException("Could not create checkpoint, criu log:\n" + new String(Files.readAllBytes(logfile)), e);
            }

            throw e;
        }

        // Get new port
        int port;
        try (DataInputStream stream = new DataInputStream(Files.newInputStream(path.resolve("port").toFile().toPath()))) {
            port = stream.readInt();
        }

        // Reopen socket
        serverConnection.startTcpServerListener(InetAddress.getLoopbackAddress(), port);
        for (Object future : channels) {
            ((ChannelFuture) future).channel().config().setAutoRead(true);
        }

        Bukkit.getPluginManager().callEvent(new CRIUWakeupEvent());
        Core.getInstance().getLogger().log(Level.INFO, "Checkpoint restored");
    }
}
