import {
  CARDS_PER_COLOR,
  DISCARD,
  DRAW_DECK,
  HAND_SIZE,
  INITIAL_DEAL,
  LOC_DECK,
  LOC_DISCARD,
  LOC_P0_BOARD,
  LOC_P0_HAND,
  MAX_STEPS,
  N_ACTIONS,
  N_CARDS,
  N_COLORS,
  PLAY,
  type DecodedAction,
  type GameState,
  type PlaceType,
} from "./types";
import { cardColor, cardRank, isHandshake } from "./cards";

function emptyMatrix(rows: number, columns: number): number[][] {
  return Array.from({ length: rows }, () => Array(columns).fill(0));
}

export function shuffledDeck(random: () => number = Math.random): number[] {
  const deck = Array.from({ length: N_CARDS }, (_, index) => index);
  for (let index = deck.length - 1; index > 0; index -= 1) {
    const swap = Math.floor(random() * (index + 1));
    [deck[index], deck[swap]] = [deck[swap], deck[index]];
  }
  return deck;
}

export function resetFromOrder(deckOrder: number[]): GameState {
  if (deckOrder.length !== N_CARDS || new Set(deckOrder).size !== N_CARDS) {
    throw new Error("deckOrder must be a permutation of 0..59");
  }
  const cardLoc = Array(N_CARDS).fill(LOC_DECK);
  for (const card of deckOrder.slice(0, HAND_SIZE)) cardLoc[card] = LOC_P0_HAND;
  for (const card of deckOrder.slice(HAND_SIZE, INITIAL_DEAL)) cardLoc[card] = LOC_P0_HAND + 1;
  return {
    deckOrder: [...deckOrder],
    drawPtr: INITIAL_DEAL,
    cardLoc,
    handPublic: Array(N_CARDS).fill(false),
    colTop: emptyMatrix(2, N_COLORS),
    colHandshakes: emptyMatrix(2, N_COLORS),
    colLength: emptyMatrix(2, N_COLORS),
    piles: Array.from({ length: N_COLORS }, () => []),
    toMove: 0,
    stepCount: 0,
    done: false,
  };
}

export function newGame(random: () => number = Math.random): GameState {
  return resetFromOrder(shuffledDeck(random));
}

export function cloneState(state: GameState): GameState {
  return {
    ...state,
    deckOrder: [...state.deckOrder],
    cardLoc: [...state.cardLoc],
    handPublic: [...state.handPublic],
    colTop: state.colTop.map((row) => [...row]),
    colHandshakes: state.colHandshakes.map((row) => [...row]),
    colLength: state.colLength.map((row) => [...row]),
    piles: state.piles.map((pile) => [...pile]),
  };
}

export function currentHandSorted(state: GameState, player = state.toMove): number[] {
  const location = LOC_P0_HAND + player;
  return state.cardLoc.flatMap((loc, card) => (loc === location ? [card] : [])).sort((a, b) => a - b);
}

export function decodeAction(action: number): DecodedAction {
  return {
    handSlot: Math.floor(action / 12),
    placeType: Math.floor((action % 12) / 6) as PlaceType,
    drawSource: action % 6,
  };
}

export function encodeAction(handSlot: number, placeType: PlaceType, drawSource: number): number {
  return handSlot * 12 + placeType * 6 + drawSource;
}

export function legalActionMask(state: GameState): boolean[] {
  const mask = Array(N_ACTIONS).fill(false);
  if (state.done) return mask;
  const hand = currentHandSorted(state);
  for (let handSlot = 0; handSlot < hand.length; handSlot += 1) {
    const card = hand[handSlot];
    const color = cardColor(card);
    const playLegal = isHandshake(card)
      ? state.colTop[state.toMove][color] === 0
      : cardRank(card) > state.colTop[state.toMove][color];
    for (const placeType of [PLAY, DISCARD] as const) {
      if (placeType === PLAY && !playLegal) continue;
      mask[encodeAction(handSlot, placeType, DRAW_DECK)] = true;
      for (let pileColor = 0; pileColor < N_COLORS; pileColor += 1) {
        const pileHasCardAfterPlace = state.piles[pileColor].length > 0 ||
          (placeType === DISCARD && pileColor === color);
        const drawingJustDiscarded = placeType === DISCARD && pileColor === color;
        if (pileHasCardAfterPlace && !drawingJustDiscarded) {
          mask[encodeAction(handSlot, placeType, pileColor + 1)] = true;
        }
      }
    }
  }
  return mask;
}

function placeCard(next: GameState, player: number, card: number, placeType: PlaceType): void {
  const color = cardColor(card);
  next.handPublic[card] = false;
  if (placeType === PLAY) {
    next.cardLoc[card] = LOC_P0_BOARD + player;
    next.colLength[player][color] += 1;
    if (isHandshake(card)) next.colHandshakes[player][color] += 1;
    else next.colTop[player][color] = cardRank(card);
  } else {
    next.cardLoc[card] = LOC_DISCARD;
    next.piles[color].push(card);
  }
}

export function previewPlacement(state: GameState, handSlot: number, placeType: PlaceType): GameState {
  const mask = legalActionMask(state);
  const hasLegalDraw = Array.from(
    { length: N_COLORS + 1 },
    (_, drawSource) => mask[encodeAction(handSlot, placeType, drawSource)],
  ).some(Boolean);
  if (!hasLegalDraw) return state;

  const next = cloneState(state);
  const card = currentHandSorted(state)[handSlot];
  placeCard(next, state.toMove, card, placeType);
  return next;
}

export function step(state: GameState, action: number): GameState {
  const mask = legalActionMask(state);
  if (!Number.isInteger(action) || action < 0 || action >= N_ACTIONS || !mask[action]) {
    return state;
  }
  const next = cloneState(state);
  const { handSlot, placeType, drawSource } = decodeAction(action);
  const player = state.toMove;
  const hand = currentHandSorted(state, player);
  const card = hand[handSlot];
  placeCard(next, player, card, placeType);

  let drawnCard: number;
  if (drawSource === DRAW_DECK) {
    drawnCard = next.deckOrder[next.drawPtr];
    next.drawPtr += 1;
    next.handPublic[drawnCard] = false;
  } else {
    const pile = next.piles[drawSource - 1];
    drawnCard = pile.pop()!;
    next.handPublic[drawnCard] = true;
  }
  next.cardLoc[drawnCard] = LOC_P0_HAND + player;
  next.stepCount += 1;
  next.done = (drawSource === DRAW_DECK && next.drawPtr >= N_CARDS) || next.stepCount >= MAX_STEPS;
  next.toMove = (1 - player) as 0 | 1;
  return next;
}

export function boardScore(state: GameState): [number, number] {
  const scores: [number, number] = [0, 0];
  for (let player = 0; player < 2; player += 1) {
    for (let color = 0; color < N_COLORS; color += 1) {
      const cards = state.cardLoc.flatMap((loc, card) =>
        loc === LOC_P0_BOARD + player && cardColor(card) === color ? [card] : [],
      );
      if (cards.length === 0) continue;
      const rankSum = cards.reduce((sum, card) => sum + cardRank(card), 0);
      const handshakes = cards.filter(isHandshake).length;
      scores[player] += (rankSum - 20) * (1 + handshakes) + (cards.length >= 8 ? 20 : 0);
    }
  }
  return scores;
}

export function cardsOnBoard(state: GameState, player: number, color: number): number[] {
  return state.cardLoc.flatMap((loc, card) =>
    loc === LOC_P0_BOARD + player && cardColor(card) === color ? [card] : [],
  );
}
