import { test, expect, type Page } from "@playwright/test";
import { existsSync } from "node:fs";
import { C } from "../src/game/battle/constants";
import {
  romPath,
  symbolsPath,
  loadSymbols,
  openReference,
} from "./reference-fixture";

type Input = {
  rate: number;
  ball: number;
  hp: number;
  maxHp: number;
  status: number;
  heldItem: number;
  tutorial: boolean;
  initialRate: number;
  caught: boolean;
  rolls: number[];
  /** wEnemyMonSpecies; Fast, Moon and Love Balls use wTempEnemyMonSpecies. */
  species: number;
  tempSpecies: number;
  enemyLevel: number;
  enemyDvs: number;
  playerSpecies: number;
  playerLevel: number;
  playerDvs: number;
  fish: boolean;
};
const make = (overrides: Partial<Input> = {}): Input => ({
  rate: 45,
  ball: C.POKE_BALL,
  hp: 20,
  maxHp: 20,
  status: 0,
  heldItem: 0,
  tutorial: false,
  initialRate: 173,
  caught: false,
  rolls: [255, 0, 128, 255],
  species: 19,
  tempSpecies: 19,
  enemyLevel: 5,
  enemyDvs: 0x8888,
  playerSpecies: 19,
  playerLevel: 5,
  playerDvs: 0x8888,
  fish: false,
  ...overrides,
});

/** Capture calls hardware Random rather than BattleRandom. Replace only
 * Random with a byte-stream reader in the in-page reference ROM. Enter
 * PokeBallEffect after its UI setup and return just after wWildMon is set;
 * GetPokeBallWobble itself is unmodified. No captured-mon loading or
 * animation timing is included in this comparison. */
async function compare(page: Page, mode: "capture" | "wobble", cases: Input[]) {
  const { at, routine } = loadSymbols();
  const addresses = {
    start: routine("PokeBallEffect.room_in_party"),
    multiplier: at("PokeBallEffect.get_multiplier_loop"),
    finish: routine("PokeBallEffect.fail_to_catch"),
    wobble: routine("GetPokeBallWobble"),
    random: at("Random"),
    romBank: at("hROMBank"),
    rate: at("wEnemyMonCatchRate"),
    finalRate: at("wFinalCatchRate"),
    wild: at("wWildMon"),
    wobbles: at("wThrownBallWobbleCount"),
    item: at("wCurItem"),
    battleType: at("wBattleType"),
    hp: at("wEnemyMonHP"),
    maxHp: at("wEnemyMonMaxHP"),
    status: at("wEnemyMonStatus"),
    species: at("wEnemyMonSpecies"),
    heldItem: at("wBattleMonItem"),
    tempSpecies: at("wTempEnemyMonSpecies"),
    tempBattleSpecies: at("wTempBattleMonSpecies"),
    enemyLevel: at("wEnemyMonLevel"),
    enemyDvs: at("wEnemyMonDVs"),
    playerLevel: at("wBattleMonLevel"),
    partyDvs: at("wPartyMon1DVs"),
    curBattleMon: at("wCurBattleMon"),
  };
  await openReference(page);
  const results = await page.evaluate(
    async ({ addresses: a, mode, cases }) => {
      const { BinjgbEmulator } = await import("/src/emulator/binjgb.ts");
      const { CrystalAssets } = await import("/src/game/rom/assets.ts");
      const { Battle } = await import("/src/game/battle/battle.ts");
      const file = (
        document.querySelector("#reference-fixture") as HTMLInputElement
      ).files![0];
      const rom = new Uint8Array(await file.arrayBuffer());
      const assets = new CrystalAssets(rom.slice());
      const offset = (bank: number, address: number) =>
        bank * 0x4000 + (address & 0x3fff);
      // Locate the first calculation instruction, with both WRAM operands
      // checked against symbols rather than assuming a fixed UI byte count.
      const pattern = [
        0xfa,
        a.rate & 255,
        a.rate >> 8,
        0x47,
        0xfa,
        a.battleType & 255,
        a.battleType >> 8,
      ];
      const starts: number[] = [];
      for (let p = a.start.address; p < a.multiplier; p++) {
        if (
          pattern.every((byte, i) => rom[offset(a.start.bank, p) + i] === byte)
        )
          starts.push(p);
      }
      if (starts.length !== 1)
        throw new Error("Capture calculation entry changed");
      const finish = offset(a.finish.bank, a.finish.address);
      if (
        ![0xea, a.wild & 255, a.wild >> 8, 0x0e, 20].every(
          (b, i) => rom[finish + i] === b,
        )
      )
        throw new Error("Capture calculation exit changed");
      rom[finish + 3] = 0xc9; // RET before DelayFrames and animation setup.
      if (
        rom[a.random] !== 0xc5 ||
        rom[a.random + 1] !== 0xf0 ||
        rom[a.random + 2] !== 4
      )
        throw new Error("Unexpected hardware Random entry");
      rom.set([0xc3, 0x00, 0xc1], a.random); // JP $c100 (controlled input).
      const reference = new BinjgbEmulator();
      await reference.load(rom, document.createElement("canvas"));
      const write = (address: number, value: number) =>
        reference.writeMemory(address, value & 255);
      const word = (address: number, value: number) => {
        write(address, value >> 8);
        write(address + 1, value);
      };
      // Preserve BC/HL, fetch stream[cursor], increment cursor, return in A.
      // DE is untouched. The original consumers immediately CP the result.
      [
        0xe5, 0xc5, 0xfa, 0x80, 0xc1, 0x4f, 0x06, 0x00, 0x21, 0x90, 0xc1, 0x09,
        0x46, 0x3c, 0xea, 0x80, 0xc1, 0x78, 0xc1, 0xe1, 0xc9,
      ].forEach((byte, i) => write(0xc100 + i, byte));
      const call = (target: { bank: number; address: number }) => {
        // DI; SP=$c7fe; select bank; CALL; save C; sentinel; loop.
        const program = [
          0xf3,
          0x31,
          0xfe,
          0xc7,
          0x3e,
          target.bank,
          0xe0,
          a.romBank & 255,
          0xea,
          0x00,
          0x20,
          0xcd,
          target.address & 255,
          target.address >> 8,
          0x79,
          0xea,
          0xfe,
          0xc6,
          0x3e,
          0x42,
          0xea,
          0xff,
          0xc6,
          0x18,
          0xfe,
        ];
        write(0xc6ff, 0);
        program.forEach((byte, i) => write(0xc000 + i, byte));
        reference.setProgramCounter(0xc000);
        reference.advanceFrames(1);
        if (reference.readMemory(0xc6ff) !== 0x42)
          throw new Error("Capture routine did not return");
        return reference.readMemory(0xc6fe);
      };
      try {
        return cases.map((input, index) => {
          let draws = 0;
          const battle = new Battle({
            data: assets.pokemon,
            party: [{ species: input.playerSpecies, dvs: input.playerDvs }] as never,
            pokedex: { seen: new Set(), caught: new Set() },
            playerId: 1,
            playerName: "CHRIS",
            timeOfDay: 1,
            text: () => ({}) as never,
            random: () => {
              if (draws >= input.rolls.length)
                throw new Error(`Exhausted capture stream in case ${index}`);
              return input.rolls[draws++];
            },
          });
          const internal = battle as unknown as {
            finalCatchRate: number;
            enemySpecies: number;
            battleType: number;
            rollCapture(
              item: number,
              tutorial: boolean,
            ): { caught: boolean; wobbles: number };
            ballWobbles(caught: boolean): number;
          };
          battle.enemyBase.catchRate = input.rate;
          battle.enemy.mon = {
            hp: input.hp,
            maxHp: input.maxHp,
            status: input.status,
            species: input.species,
            level: input.enemyLevel,
            dvs: input.enemyDvs,
          } as never;
          battle.player.mon = { item: input.heldItem, level: input.playerLevel } as never;
          internal.enemySpecies = input.tempSpecies;
          if (input.fish) internal.battleType = 4;
          internal.finalCatchRate = input.initialRate;
          const native =
            mode === "capture"
              ? internal.rollCapture(input.ball, input.tutorial)
              : {
                  caught: input.caught,
                  wobbles: internal.ballWobbles(input.caught),
                };
          write(0xff70, 1);
          write(a.rate, input.rate);
          write(a.finalRate, input.initialRate);
          write(a.wild, mode === "wobble" && input.caught ? 19 : 0);
          write(a.wobbles, 0);
          write(a.item, input.ball);
          write(a.battleType, input.tutorial ? 3 : input.fish ? 4 : 0);
          write(a.species, input.species);
          write(a.tempSpecies, input.tempSpecies);
          write(a.tempBattleSpecies, input.playerSpecies);
          write(a.enemyLevel, input.enemyLevel);
          write(a.playerLevel, input.playerLevel);
          word(a.enemyDvs, input.enemyDvs);
          word(a.partyDvs, input.playerDvs);
          write(a.curBattleMon, 0);
          write(a.status, input.status);
          write(a.heldItem, input.heldItem);
          word(a.hp, input.hp);
          word(a.maxHp, input.maxHp);
          write(0xc180, 0);
          // Poison unused bytes so an extra reference draw cannot pass silently.
          for (let i = 0; i < 16; i++)
            write(0xc190 + i, input.rolls[i] ?? 0xee);
          if (mode === "capture")
            call({ bank: a.start.bank, address: starts[0] });
          let outcome = 0;
          for (let i = 0; i < 4 && outcome === 0; i++) outcome = call(a.wobble);
          if (outcome !== 1 && outcome !== 2)
            throw new Error("Wobble loop did not finish");
          return {
            native: { ...native, rate: internal.finalCatchRate, draws },
            original: {
              caught: reference.readMemory(a.wild) !== 0,
              wobbles: reference.readMemory(a.wobbles),
              rate: reference.readMemory(a.finalRate),
              draws: reference.readMemory(0xc180),
            },
            outcome,
          };
        });
      } finally {
        reference.dispose();
      }
    },
    { addresses, mode, cases },
  );
  results.forEach((result, i) => {
    expect(
      result.native,
      `${mode} case ${i}: ${JSON.stringify(cases[i])}`,
    ).toEqual(result.original);
    expect(result.outcome).toBe(result.native.caught ? 1 : 2);
  });
}

test.beforeEach(() => {
  test.skip(
    !existsSync(romPath) || !existsSync(symbolsPath),
    "Requires the local user ROM and private PRET rebuild.",
  );
});

test("capture formula, success and random consumption match the original", async ({
  page,
}) => {
  const cases: Input[] = [];
  // Every catch-rate byte, all supported multipliers, saturation and odd rounding.
  for (const ball of [C.POKE_BALL, C.GREAT_BALL, C.ULTRA_BALL])
    for (let rate = 0; rate < 256; rate++) cases.push(make({ ball, rate }));
  // HP scaling boundary, byte wrapping above 341, and the largest legal HP.
  // 342 and 683 are omitted: the original divides by zero and hangs.
  for (const maxHp of [
    1, 2, 84, 85, 86, 87, 100, 255, 256, 340, 341, 343, 400, 511, 512, 682, 684,
    703, 713,
  ])
    for (const hp of new Set(
      [1, 2, Math.floor(maxHp / 2), maxHp - 1, maxHp].filter(
        (hp) => hp > 0 && hp <= maxHp,
      ),
    ))
      for (const ball of [C.POKE_BALL, C.GREAT_BALL, C.ULTRA_BALL])
        cases.push(make({ maxHp, hp, ball, rate: 255 }));
  // Sleep counts, freeze, and the intentionally ineffective poison/burn/paralysis bits.
  for (let status = 0; status < 128; status++)
    for (const rate of [1, 45, 255])
      cases.push(make({ status, rate, hp: 1, maxHp: 100 }));
  // Exhaust the success draw at several final rates, including equality and 255.
  for (const rate of [1, 2, 3, 45, 127, 254, 255])
    for (let roll = 0; roll < 256; roll++)
      cases.push(make({ rate, hp: 0, rolls: [roll, 0, 128, 255] }));
  for (const heldItem of [0, 0x49, 0x8c, 0xac]) cases.push(make({ heldItem }));
  // Neither guaranteed-catch path updates the previous final rate or draws RNG.
  for (const initialRate of [0, 1, 173, 255]) {
    cases.push(make({ ball: C.MASTER_BALL, initialRate, rolls: [] }));
    cases.push(make({ tutorial: true, initialRate, rolls: [] }));
  }
  await compare(page, "capture", cases);
});

test("ball shakes match original table boundaries, outcomes and draw counts", async ({
  page,
}) => {
  const cases: Input[] = [];
  // Every final rate with early escape, each later escape and all three shakes.
  for (let initialRate = 0; initialRate < 256; initialRate++) {
    for (const rolls of [[255], [0, 255], [0, 0, 255], [0, 0, 0]])
      cases.push(make({ initialRate, rolls }));
    cases.push(make({ initialRate, caught: true, rolls: [] }));
  }
  // Every random byte at each table entry; neighboring rates are covered by the loop above.
  for (const rate of [
    1, 2, 3, 4, 5, 7, 10, 15, 20, 30, 40, 50, 60, 80, 100, 120, 140, 160, 180,
    200, 220, 240, 254, 255,
  ])
    for (let roll = 0; roll < 256; roll++)
      cases.push(make({ initialRate: rate, rolls: [roll, 255] }));
  await compare(page, "wobble", cases);
});

test("Kurt's and Park Balls match the original multipliers and their bugs", async ({
  page,
}) => {
  const cases: Input[] = [];
  const all = Array.from({ length: 251 }, (_, i) => i + 1);
  // Heavy Ball: every species' weight (Transform changes wEnemyMonSpecies,
  // not the temp species), around the light-mon floor and the 255 cap.
  // Kadabra's and Sunflora's reads reach VRAM and are refused natively.
  for (const species of all.filter((species) => species !== 64 && species !== 192))
    for (const rate of [1, 20, 21, 45, 240])
      cases.push(make({ ball: C.HEAVY_BALL, species, tempSpecies: 19, rate, hp: 1, maxHp: 100 }));
  // Level Ball: each halving boundary and saturation; it skips HP and status.
  const levels = [1, 2, 3, 4, 5, 7, 8, 9, 15, 16, 17, 31, 32, 33, 63, 64, 65, 100];
  for (const playerLevel of levels)
    for (const enemyLevel of levels)
      for (const rate of [1, 31, 32, 64, 255])
        cases.push(make({ ball: C.LEVEL_BALL, playerLevel, enemyLevel, rate, status: 1 }));
  // Lure Ball only on a fishing rod battle; Park Ball everywhere.
  for (let rate = 0; rate < 256; rate++) {
    cases.push(make({ ball: C.LURE_BALL, rate, fish: true, hp: 1, maxHp: 100 }));
    cases.push(make({ ball: C.LURE_BALL, rate, hp: 1, maxHp: 100 }));
    cases.push(make({ ball: C.PARK_BALL, rate, hp: 1, maxHp: 100 }));
    cases.push(make({ ball: C.FRIEND_BALL, rate, hp: 1, maxHp: 100 }));
  }
  // Fast and Moon Balls over every species (temp species, not wEnemyMonSpecies).
  for (const tempSpecies of all)
    for (const ball of [C.FAST_BALL, C.MOON_BALL])
      for (const rate of [3, 64])
        cases.push(make({ ball, tempSpecies, species: 19, rate, hp: 1, maxHp: 100 }));
  // Love Ball: same and different species, each gender pairing, genderless.
  for (const tempSpecies of all)
    for (const [playerDvs, enemyDvs] of [[0x0000, 0x0000], [0xffff, 0xffff], [0x0000, 0xffff], [0xffff, 0x0000], [0x7777, 0x8888]])
      cases.push(make({ ball: C.LOVE_BALL, tempSpecies, playerSpecies: tempSpecies, playerDvs, enemyDvs, rate: 30, hp: 1, maxHp: 100 }));
  for (const playerSpecies of [1, 29, 32, 81, 251])
    cases.push(make({ ball: C.LOVE_BALL, tempSpecies: 19, playerSpecies, rate: 30, hp: 1, maxHp: 100 }));
  await compare(page, "capture", cases);
});