from __future__ import annotations

import typing as tp
from dataclasses import dataclass
from functools import cache

Player = tp.Literal["X", "O"]
Winner = Player | tp.Literal["draw"]
Space = Player | tp.Literal[" "]
Board = tuple[Space, ...]
Score = int

C4_COLS = 7
C4_ROWS = 6


@dataclass(frozen=True)
class C4GameState:
    player: Player
    board: Board

    def __str__(self) -> str:
        b = self.board
        return f"───┼───┼───┼───┼───┼───┼───\n".join(
            f" {b[i+0]} │ {b[i+1]} │ {b[i+2]} │ {b[i+3]} │ {b[i+4]} │ {b[i+5]} │ {b[i+6]}\n"
            for i in range(0, C4_COLS * C4_ROWS, C4_COLS)
        )


@dataclass(frozen=True)
class Move:
    col: int

    def __str__(self) -> str:
        return str(self.col + 1)


@cache
def get_valid_moves(state: C4GameState) -> tp.Iterable[Move]:
    # NOTE: a slot must not be full if its top row is empty
    return tuple(Move(col) for col in range(C4_COLS) if state.board[col] == " ")


@cache
def apply_move(state: C4GameState, move: Move) -> C4GameState:
    new_player = "X" if state.player == "O" else "O"

    # B)
    for i in range(1, C4_ROWS):
        idx_above = (i - 1) * C4_COLS + move.col
        idx_below = (i - 0) * C4_COLS + move.col
        if state.board[idx_below] != " ":
            idx = idx_above
            break
    else:
        # bottom
        idx = (C4_ROWS - 1) * C4_COLS + move.col

    new_board = list(state.board)
    new_board[idx] = state.player
    new_board = tuple(new_board)

    return C4GameState(player=new_player, board=new_board)

@cache
def get_winner(state: C4GameState) -> Winner | None:
    pass

START = C4GameState(player="X", board=(" ",) * C4_COLS * C4_ROWS)

if __name__ == "__main__":
    print(START)

    print(apply_move(START, Move(col=6)))