From 6e11324182c330fdd0e26ed18e4bedec3e19081d Mon Sep 17 00:00:00 2001
From: Reldnahc <cwmlee00@outlook.com>
Date: Mon, 6 May 2024 02:21:58 -0500
Subject: [PATCH] have applyCheckTrappedAbAttrs() send the pokemon that is
 being prevented from switching also.

implement conditional arena traps.

implement magnet pull.

correct arena trap to only effect grounded pokemon.

correct shadow tag to only effect non-ghost types.
---
 src/data/ability.ts | 33 ++++++++++++++++++++++++---------
 src/phases.ts       |  4 ++--
 2 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/src/data/ability.ts b/src/data/ability.ts
index eec1240e1d3..23c7256660b 100644
--- a/src/data/ability.ts
+++ b/src/data/ability.ts
@@ -2190,13 +2190,13 @@ export class CheckTrappedAbAttr extends AbAttr {
     super(false);
   }
   
-  applyCheckTrapped(pokemon: Pokemon, passive: boolean, trapped: Utils.BooleanHolder, args: any[]): boolean | Promise<boolean> {
+  applyCheckTrapped(pokemon: Pokemon, target: Pokemon, passive: boolean, trapped: Utils.BooleanHolder, args: any[]): boolean | Promise<boolean> {
     return false;
   }
 }
 
 export class ArenaTrapAbAttr extends CheckTrappedAbAttr {
-  applyCheckTrapped(pokemon: Pokemon, passive: boolean, trapped: Utils.BooleanHolder, args: any[]): boolean {
+  applyCheckTrapped(pokemon: Pokemon, target: Pokemon, passive: boolean, trapped: Utils.BooleanHolder, args: any[]): boolean {
     trapped.value = true;
     return true;
   }
@@ -2206,6 +2206,23 @@ export class ArenaTrapAbAttr extends CheckTrappedAbAttr {
   }
 }
 
+export class ConditionalArenaTrapAbAttr extends ArenaTrapAbAttr {
+  private condition: (opponent: Pokemon) => boolean;
+
+  constructor(condition: (opponent: Pokemon) => boolean) {
+    super();
+    this.condition = condition;
+  }
+
+  applyCheckTrapped(pokemon: Pokemon, opponent: Pokemon, passive: boolean, trapped: Utils.BooleanHolder, args: any[]): boolean {
+    if (this.condition(opponent)) {
+      trapped.value = true;
+      return true;
+    }
+    return false;
+  }
+}
+
 export class MaxMultiHitAbAttr extends AbAttr {
   apply(pokemon: Pokemon, passive: boolean, cancelled: Utils.BooleanHolder, args: any[]): boolean {
     (args[0] as Utils.IntegerHolder).value = 0;
@@ -2638,8 +2655,8 @@ export function applyPostTerrainChangeAbAttrs(attrType: { new(...args: any[]): P
 }
 
 export function applyCheckTrappedAbAttrs(attrType: { new(...args: any[]): CheckTrappedAbAttr },
-  pokemon: Pokemon, trapped: Utils.BooleanHolder, ...args: any[]): Promise<void> {
-  return applyAbAttrsInternal<CheckTrappedAbAttr>(attrType, pokemon, (attr, passive) => attr.applyCheckTrapped(pokemon, passive, trapped, args), args, true);
+  pokemon: Pokemon, target: Pokemon, trapped: Utils.BooleanHolder, ...args: any[]): Promise<void> {
+  return applyAbAttrsInternal<CheckTrappedAbAttr>(attrType, pokemon, (attr, passive) => attr.applyCheckTrapped(pokemon, target, passive, trapped, args), args, true);
 }
 
 export function applyPostBattleAbAttrs(attrType: { new(...args: any[]): PostBattleAbAttr },
@@ -2733,7 +2750,7 @@ export function initAbilities() {
     new Ability(Abilities.INTIMIDATE, 3)
       .attr(PostSummonStatChangeAbAttr, BattleStat.ATK, -1, false, true),
     new Ability(Abilities.SHADOW_TAG, 3)
-      .attr(ArenaTrapAbAttr),
+      .attr(ConditionalArenaTrapAbAttr, (opponent: Pokemon) => !opponent.isOfType(Type.GHOST)),
     new Ability(Abilities.ROUGH_SKIN, 3)
       .attr(PostDefendContactDamageAbAttr, 8)
       .bypassFaint(),
@@ -2790,9 +2807,7 @@ export function initAbilities() {
       .attr(StatusEffectImmunityAbAttr, StatusEffect.BURN)
       .ignorable(),
     new Ability(Abilities.MAGNET_PULL, 3)
-      /*.attr(ArenaTrapAbAttr)
-      .condition((pokemon: Pokemon) => pokemon.getOpponent()?.isOfType(Type.STEEL))*/
-      .unimplemented(),
+      .attr(ConditionalArenaTrapAbAttr, (opponent: Pokemon) => opponent.isOfType(Type.STEEL)),
     new Ability(Abilities.SOUNDPROOF, 3)
       .attr(MoveImmunityAbAttr, (pokemon, attacker, move) => pokemon !== attacker && move.getMove().hasFlag(MoveFlags.SOUND_BASED))
       .ignorable(),
@@ -2866,7 +2881,7 @@ export function initAbilities() {
       .attr(PostSummonWeatherChangeAbAttr, WeatherType.SUNNY)
       .attr(PostBiomeChangeWeatherChangeAbAttr, WeatherType.SUNNY),
     new Ability(Abilities.ARENA_TRAP, 3)
-      .attr(ArenaTrapAbAttr)
+      .attr(ConditionalArenaTrapAbAttr, (opponent: Pokemon) => opponent.isGrounded())
       .attr(DoubleBattleChanceAbAttr),
     new Ability(Abilities.VITAL_SPIRIT, 3)
       .attr(StatusEffectImmunityAbAttr, StatusEffect.SLEEP)
diff --git a/src/phases.ts b/src/phases.ts
index fc6af354a2d..fb2367c2c96 100644
--- a/src/phases.ts
+++ b/src/phases.ts
@@ -1779,7 +1779,7 @@ export class CommandPhase extends FieldPhase {
           const trapped = new Utils.BooleanHolder(false);
           const batonPass = isSwitch && args[0] as boolean;
           if (!batonPass)
-            enemyField.forEach(enemyPokemon => applyCheckTrappedAbAttrs(CheckTrappedAbAttr, enemyPokemon, trapped));
+            enemyField.forEach(enemyPokemon => applyCheckTrappedAbAttrs(CheckTrappedAbAttr, enemyPokemon, playerPokemon, trapped));
           if (batonPass || (!trapTag && !trapped.value)) {
             this.scene.currentBattle.turnCommands[this.fieldIndex] = isSwitch
               ? { command: Command.POKEMON, cursor: cursor, args: args }
@@ -1877,7 +1877,7 @@ export class EnemyCommandPhase extends FieldPhase {
 
       const trapTag = enemyPokemon.findTag(t => t instanceof TrappedTag) as TrappedTag;
       const trapped = new Utils.BooleanHolder(false);
-      opponents.forEach(playerPokemon => applyCheckTrappedAbAttrs(CheckTrappedAbAttr, playerPokemon, trapped));
+      opponents.forEach(playerPokemon => applyCheckTrappedAbAttrs(CheckTrappedAbAttr, playerPokemon, enemyPokemon, trapped));
       if (!trapTag && !trapped.value) {
         const partyMemberScores = trainer.getPartyMemberMatchupScores(enemyPokemon.trainerSlot, true);