connect_four.py

#!/usr/bin/env python3
"""
Connect Four – a simple console implementation for one or two players.

The game uses a 6x7 board. Players take turns dropping a disc into one of the
seven columns. The disc occupies the lowest available row in that column.
The first player to connect four of their discs horizontally, vertically,
or diagonally wins. If the board fills without a winner, the game ends in a
draw.

Usage:
    python connect_four.py

The program will prompt for the number of players (1 or 2). In single‑player
mode the second player is a very simple AI that chooses the first available
column. Feel free to replace the AI logic with something more sophisticated.
"""

from __future__ import annotations
import sys
import random
from typing import List, Tuple

ROWS: int = 6
COLUMNS: int = 7
EMPTY: str = " "
PLAYER_SYMBOLS: Tuple[str, str] = ("X", "O")


def create_board() -> List[List[str]]:
    """Return a new empty board."""
    return [[EMPTY for _ in range(COLUMNS)] for _ in range(ROWS)]


def print_board(board: List[List[str]]) -> None:
    """Print the board to the console."""
    print("\n" + "  ".join(str(i + 1) for i in range(COLUMNS)))
    for row in board:
        print("|" + "|".join(row) + "|")
    print("-" * (COLUMNS * 2 + 1))


def is_valid_move(board: List[List[str]], col: int) -> bool:
    """Return True if a disc can be dropped into the given column."""
    return board[0][col] == EMPTY


def make_move(board: List[List[str]], col: int, symbol: str) -> None:
    """Drop a disc into the given column."""
    for row in reversed(range(ROWS)):
        if board[row][col] == EMPTY:
            board[row][col] = symbol
            break


def check_direction(
    board: List[List[str]],
    start_row: int,
    start_col: int,
    delta_row: int,
    delta_col: int,
    symbol: str,
) -> bool:
    """Check a line of four in a specific direction."""
    count = 0
    row, col = start_row, start_col
    while 0 <= row < ROWS and 0 <= col < COLUMNS:
        if board[row][col] == symbol:
            count += 1
            if count == 4:
                return True
        else:
            count = 0
        row += delta_row
        col += delta_col
    return False


def check_win(board: List[List[str]], symbol: str) -> bool:
    """Return True if the given symbol has a connect four."""
    # Horizontal
    for r in range(ROWS):
        if check_direction(board, r, 0, 0, 1, symbol):
            return True
    # Vertical
    for c in range(COLUMNS):
        if check_direction(board, 0, c, 1, 0, symbol):
            return True
    # Diagonal /
    for r in range(ROWS - 3):
        for c in range(COLUMNS - 3):
            if check_direction(board, r, c, 1, 1, symbol):
                return True
    # Diagonal \
    for r in range(3, ROWS):
        for c in range(COLUMNS - 3):
            if check_direction(board, r, c, -1, 1, symbol):
                return True
    return False


def board_full(board: List[List[str]]) -> bool:
    """Return True if the board has no empty cells."""
    return all(cell != EMPTY for cell in board[0])


def get_player_move(board: List[List[str]], player_num: int) -> int:
    """Prompt the human player for a column."""
    while True:
        try:
            col = int(input(f"Player {player_num} ({PLAYER_SYMBOLS[player_num - 1]}), choose column (1-{COLUMNS}): ")) - 1
            if 0 <= col < COLUMNS and is_valid_move(board, col):
                return col
            print(f"Column {col + 1} is full or out of range. Try again.")
        except ValueError:
            print("Invalid input. Enter a number.")


def get_ai_move(board: List[List[str]]) -> int:
    """Very simple AI: choose the first available column."""
    for col in range(COLUMNS):
        if is_valid_move(board, col):
            return col
    # Should never reach here if called correctly
    raise RuntimeError("No valid moves for AI")


def main() -> None:
    print("Welcome to Connect Four!")
    while True:
        try:
            num_players = int(input("Enter number of players (1 or 2): "))
            if num_players in (1, 2):
                break
            print("Please enter 1 or 2.")
        except ValueError:
            print("Invalid input. Enter a number.")

    board = create_board()
    current_player = 1
    while True:
        print_board(board)
        if num_players == 1 and current_player == 2:
            col = get_ai_move(board)
            print(f"AI chooses column {col + 1}")
        else:
            col = get_player_move(board, current_player)

        make_move(board, col, PLAYER_SYMBOLS[current_player - 1])

        if check_win(board, PLAYER_SYMBOLS[current_player - 1]):
            print_board(board)
            if num_players == 1 and current_player == 2:
                print("AI wins! Better luck next time.")
            else:
                print(f"Player {current_player} ({PLAYER_SYMBOLS[current_player - 1]}) wins!")
            break

        if board_full(board):
            print_board(board)
            print("It's a draw!")
            break

        current_player = 2 if current_player == 1 else 1

    print("Thanks for playing!")


if __name__ == "__main__":
    main()