import { describe, expect, test } from "vitest";

import {
  boardScore,
  cardsOnBoard,
  currentHandSorted,
  encodeAction,
  legalActionMask,
  previewPlacement,
  resetFromOrder,
  step,
} from "./engine";
import { DISCARD, DRAW_DECK, N_CARDS, PLAY } from "./types";
import { cardColor } from "./cards";
import { observation } from "./observation";
import fixture from "./parity-fixture.json";
import type { GameState } from "./types";

const orderedDeck = Array.from({ length: N_CARDS }, (_, index) => index);

describe("JAX-compatible game engine", () => {
  test("deals eight cards to each player in explicit deck order", () => {
    const state = resetFromOrder(orderedDeck);
    expect(currentHandSorted(state, 0)).toEqual([0, 1, 2, 3, 4, 5, 6, 7]);
    expect(currentHandSorted(state, 1)).toEqual([8, 9, 10, 11, 12, 13, 14, 15]);
  });

  test("cannot discard and redraw the same card", () => {
    const state = resetFromOrder(orderedDeck);
    const color = cardColor(currentHandSorted(state)[0]);
    const mask = legalActionMask(state);
    expect(mask[encodeAction(0, DISCARD, color + 1)]).toBe(false);
    expect(mask[encodeAction(0, DISCARD, DRAW_DECK)]).toBe(true);
  });

  test("previews a placed card without ending the turn", () => {
    const state = resetFromOrder(orderedDeck);
    const card = currentHandSorted(state)[3];
    const preview = previewPlacement(state, 3, PLAY);

    expect(currentHandSorted(preview, 0)).not.toContain(card);
    expect(cardsOnBoard(preview, 0, cardColor(card))).toContain(card);
    expect(preview.toMove).toBe(0);
    expect(preview.drawPtr).toBe(state.drawPtr);
    expect(preview.stepCount).toBe(state.stepCount);
    expect(currentHandSorted(state, 0)).toContain(card);
  });

  test("deck-only games terminate after exactly 44 plies", () => {
    let state = resetFromOrder(orderedDeck);
    let plies = 0;
    while (!state.done) {
      const action = legalActionMask(state).findIndex((legal, action) => legal && action % 6 === 0);
      state = step(state, action);
      plies += 1;
    }
    expect(plies).toBe(44);
    expect(state.drawPtr).toBe(N_CARDS);
  });

  test("observation has the trained model shape", () => {
    expect(observation(resetFromOrder(orderedDeck), 0)).toHaveLength(454);
  });

  test("illegal actions are no-ops", () => {
    const state = resetFromOrder(orderedDeck);
    expect(step(state, -1)).toBe(state);
    expect(boardScore(state)).toEqual([0, 0]);
  });

  test("matches deterministic JAX masks, observations, and transitions", () => {
    for (let index = 0; index < fixture.rows.length; index += 1) {
      const row = fixture.rows[index];
      const state = row.state as GameState;
      expect(legalActionMask(state)).toEqual(row.legalMask);
      const p0 = observation(state, 0);
      const p1 = observation(state, 1);
      row.observationP0.forEach((value, offset) => expect(p0[offset]).toBeCloseTo(value, 6));
      row.observationP1.forEach((value, offset) => expect(p1[offset]).toBeCloseTo(value, 6));
      if (index + 1 < fixture.rows.length) {
        expect(step(state, row.action)).toEqual(fixture.rows[index + 1].state);
      }
    }
  });
});
