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

import fixture from "./match-parity-fixture.json";
import {
  N_ROUNDS,
  matchAdvanceRound,
  matchFromOrders,
  matchLegalActionMask,
  matchScore,
  matchStep,
  roundScore,
  startingPlayer,
  type MatchState,
} from "./match";
import { MATCH_OBS_DIM, matchObservation } from "./matchObservation";
import type { GameState } from "./types";

interface FixtureRow {
  match: {
    round: GameState;
    deckOrders: number[][];
    coinFlips: number[];
    roundIdx: number;
    carry: number[];
    done: boolean;
  };
  player: number;
  observation: number[];
}

const rows = fixture.rows as unknown as FixtureRow[];

function toMatch(row: FixtureRow): MatchState {
  return {
    round: row.match.round,
    deckOrders: row.match.deckOrders,
    coinFlips: row.match.coinFlips,
    roundIdx: row.match.roundIdx,
    carry: [row.match.carry[0], row.match.carry[1]],
    // The fixture is generated from the JAX match, which is always three rounds.
    totalRounds: N_ROUNDS,
    // Not part of the observation; only present so the type is satisfied.
    roundComplete: false,
    roundHistory: [],
    done: row.match.done,
  };
}

function firstLegal(match: MatchState): number {
  const index = matchLegalActionMask(match).findIndex(Boolean);
  if (index < 0) throw new Error("no legal action");
  return index;
}

describe("match observation parity with JAX", () => {
  it("has the dimension the exported model expects", () => {
    expect(MATCH_OBS_DIM).toBe(fixture.obsDim);
    expect(MATCH_OBS_DIM).toBe(501);
  });

  it("reproduces every fixture observation", () => {
    expect(rows.length).toBeGreaterThan(100);
    let worst = 0;
    let worstAt = "";

    for (const [index, row] of rows.entries()) {
      const actual = matchObservation(toMatch(row), row.player);
      expect(actual.length).toBe(row.observation.length);

      for (let i = 0; i < actual.length; i += 1) {
        const delta = Math.abs(actual[i] - row.observation[i]);
        if (delta > worst) {
          worst = delta;
          worstAt = `row ${index}, feature ${i}`;
        }
      }
    }
    // A mismatch here throws nowhere: the ONNX policy consumes the wrong vector
    // and plays worse for reasons nobody can see. So the bar is float32 round-off,
    // not "close enough".
    expect(worst, `largest disagreement at ${worstAt}`).toBeLessThan(1e-5);
  });

  it("covers positions past a round roll-over, with a real carry", () => {
    expect(rows.some((row) => row.match.roundIdx > 0)).toBe(true);
    expect(rows.some((row) => row.match.carry[0] !== row.match.carry[1])).toBe(true);
  });
});

/** Play a full match, advancing past each round-complete pause. */
function playMatch(start: MatchState): { final: MatchState; pauses: number } {
  let match = start;
  let pauses = 0;
  for (let ply = 0; ply < 1400 && !match.done; ply += 1) {
    if (match.roundComplete) {
      pauses += 1;
      match = matchAdvanceRound(match);
      continue;
    }
    match = matchStep(match, firstLegal(match));
  }
  return { final: match, pauses };
}

describe("match rules", () => {
  it("pauses at each round end, then plays exactly three rounds", () => {
    const start = matchFromOrders(rows[0].match.deckOrders, rows[0].match.coinFlips);
    const { final, pauses } = playMatch(start);
    // Two roll-overs for three rounds; the third ends the match, not a pause.
    expect(pauses).toBe(N_ROUNDS - 1);
    expect(final.done).toBe(true);
    expect(final.roundIdx).toBe(N_ROUNDS - 1);
    expect(final.roundHistory.length).toBe(N_ROUNDS);
  });

  it("holds the finished board while round-complete, offering no legal move", () => {
    let match = matchFromOrders(rows[0].match.deckOrders, rows[0].match.coinFlips);
    while (!match.roundComplete && !match.done) match = matchStep(match, firstLegal(match));
    expect(match.roundComplete).toBe(true);
    expect(match.roundIdx).toBe(0); // not advanced yet
    expect(matchLegalActionMask(match).some(Boolean)).toBe(false);
    // The running total already includes the round being shown.
    expect(matchScore(match)).toEqual(roundScore(match));
  });

  it("banks each round exactly once into the history and carry", () => {
    const { final } = playMatch(matchFromOrders(rows[0].match.deckOrders, rows[0].match.coinFlips));
    const summed = final.roundHistory.reduce(
      (acc, [you, rival]) => [acc[0] + you, acc[1] + rival] as [number, number],
      [0, 0] as [number, number],
    );
    // carry banks rounds 1..2; the third is still on the board at match end.
    const board = roundScore(final);
    expect(matchScore(final)).toEqual([summed[0], summed[1]]);
    expect([final.carry[0] + board[0], final.carry[1] + board[1]]).toEqual(matchScore(final));
  });

  it("lets whoever has more points begin, and flips a coin when level", () => {
    expect(startingPlayer([60, 10], 1, [1, 1, 1])).toBe(0); // ahead leads, coin ignored
    expect(startingPlayer([10, 60], 1, [0, 0, 0])).toBe(1);
    expect(startingPlayer([30, 30], 1, [0, 0, 0])).toBe(0); // level falls back to the coin
    expect(startingPlayer([30, 30], 1, [1, 1, 1])).toBe(1);
  });
});
