2 files changed, 10 insertions, 11 deletions
diff --git a/src/main/java/world/anhgelus/floodhunt/Floodhunt.java b/src/main/java/world/anhgelus/floodhunt/Floodhunt.java
index 5dcaf6e..a5f545f 100644
--- a/src/main/java/world/anhgelus/floodhunt/Floodhunt.java
+++ b/src/main/java/world/anhgelus/floodhunt/Floodhunt.java
@@ -4,7 +4,6 @@ import com.mojang.brigadier.Command;
 import com.mojang.brigadier.exceptions.SimpleCommandExceptionType;
 import net.fabricmc.api.ModInitializer;
 import net.fabricmc.fabric.api.command.v2.CommandRegistrationCallback;
-import net.fabricmc.fabric.api.entity.event.v1.ServerLivingEntityEvents;
 import net.fabricmc.fabric.api.entity.event.v1.ServerPlayerEvents;
 import net.fabricmc.fabric.api.event.lifecycle.v1.ServerLifecycleEvents;
 import net.fabricmc.fabric.api.gamerule.v1.GameRuleBuilder;
@@ -17,11 +16,9 @@ import net.minecraft.network.packet.s2c.play.OverlayMessageS2CPacket;
 import net.minecraft.server.MinecraftServer;
 import net.minecraft.server.command.CommandManager;
 import net.minecraft.server.command.ServerCommandSource;
-import net.minecraft.server.network.ServerPlayerEntity;
 import net.minecraft.text.Text;
 import net.minecraft.util.Formatting;
 import net.minecraft.util.Identifier;
-import net.minecraft.world.GameMode;
 import net.minecraft.world.rule.GameRule;
 import net.minecraft.world.rule.GameRuleCategory;
 import org.slf4j.Logger;
@@ -175,16 +172,10 @@ public class Floodhunt implements ModInitializer {
 
 		ServerMessageEvents.ALLOW_CHAT_MESSAGE.register((message, sender, params) -> false);
 
-		ServerLivingEntityEvents.AFTER_DEATH.register((entity, damageSource) -> {
-			if (!(entity instanceof ServerPlayerEntity) || game == null) return;
-			if (!game.started()) return;
-			if (game.wonByInfected()) game.end();
-		});
-
 		ServerPlayerEvents.AFTER_RESPAWN.register((oldPlayer, newPlayer, alive) -> {
 			if (game == null) return;
 			if (!game.started()) return;
-			newPlayer.changeGameMode(GameMode.SPECTATOR);
+			game.onRespawn(newPlayer);
 		});
 
 		ServerPlayConnectionEvents.JOIN.register((handler, sender, server) -> {
diff --git a/src/main/java/world/anhgelus/floodhunt/game/Game.java b/src/main/java/world/anhgelus/floodhunt/game/Game.java
index 868b145..b04b1e9 100644
--- a/src/main/java/world/anhgelus/floodhunt/game/Game.java
+++ b/src/main/java/world/anhgelus/floodhunt/game/Game.java
@@ -121,7 +121,7 @@ public class Game {
 		end();
 	}
 
-	public void end() {
+	private void end() {
 		final var timer = TimerAccess.getTimerFromOverworld(server);
 		timer.floodhunt_cancel();
 
@@ -198,4 +198,12 @@ public class Game {
 		final var payload = new GamePayload(hasStarted);
 		server.getPlayerManager().sendToAll(ServerPlayNetworking.createS2CPacket(payload));
 	}
+
+	public void onRespawn(ServerPlayerEntity player) {
+		infected.add(player.getUuid());
+		if (wonByInfected()) end();
+		player.changeGameMode(GameMode.SURVIVAL);
+		player.sendMessage(Text.translatable("floodhunt.game.infected"));
+		player.sendMessage(Text.translatable("floodhunt.game.start.mole.subtitle"));
+	}
 }