import network
import time
import sys
from machine import Pin, I2C
from umqtt.simple import MQTTClient
from ssd1306 import SSD1306_I2C

# MQTT Server Parameters
MQTT_CLIENT_ID      = "tic_tac_toe_board"
MQTT_BROKER         = "broker.mqttdashboard.com"
MQTT_GAME_TOPIC     = "tic_tac_toe/moves"

# OLED Display Parameters
OLED_WIDTH = 128
OLED_HEIGHT = 64
I2C_SDA_PIN = 21
I2C_SCL_PIN = 22

# Game Board
board = [[" " for _ in range(3)] for _ in range(3)]

# Player Turn
current_player = "Player 1"

# Initialize I2C and OLED
i2c = I2C(sda=Pin(I2C_SDA_PIN), scl=Pin(I2C_SCL_PIN))
oled = SSD1306_I2C(OLED_WIDTH, OLED_HEIGHT, i2c)

def connect_to_wifi():
    # Connect to Wi-Fi
    print("Connecting to Wi-Fi", end="")
    sta_if = network.WLAN(network.STA_IF)
    sta_if.active(True)
    sta_if.connect('Wokwi-GUEST', '')
    while not sta_if.isconnected():
        print(".", end="")
        time.sleep(0.1)
    print(" Connected!")

def connect_to_mqtt_broker():
    # Connect to MQTT broker
    print("Connecting to MQTT server... ", end="")
    client.connect()
    print("Connected!")
    client.subscribe(MQTT_GAME_TOPIC)
    print("Subscribed to {}".format(MQTT_GAME_TOPIC))

def draw_board():
    oled.fill(0)
    for row in range(3):
        for col in range(3):
            oled.text(board[row][col], col * 40, row * 20)
    oled.show()

def handle_move(topic, message):
    global current_player

    move = eval(message.decode("utf-8"))
    player, (row, col) = move

    if player == current_player and board[row][col] == " ":
        board[row][col] = "X" if current_player == "Player 1" else "O"
        draw_board()

        # Check for a winner or a tie and update the current player
        if check_winner() or check_tie():
            print("Game over!")
            reset_game()
        else:
            current_player = "Player 2" if current_player == "Player 1" else "Player 1"
            print("It's {}'s turn.".format(current_player))
            send_turn_message()

def check_winner():
    # Check rows, columns, and diagonals for a winner
    for i in range(3):
        if board[i][0] == board[i][1] == board[i][2] != " " or \
           board[0][i] == board[1][i] == board[2][i] != " ":
            print("{} wins!".format(current_player))
            return True
    if board[0][0] == board[1][1] == board[2][2] != " " or \
       board[0][2] == board[1][1] == board[2][0] != " ":
        print("{} wins!".format(current_player))
        return True
    return False

def check_tie():
    # Check for a tie
    if all(board[i][j] != " " for i in range(3) for j in range(3)):
        print("It's a tie!")
        return True
    return False

def reset_game():
    global board, current_player
    board = [[" " for _ in range(3)] for _ in range(3)]
    current_player = "Player 1"
    draw_board()
    send_turn_message()

def send_turn_message():
    message = "{}'s turn. Make your move: ".format(current_player)
    client.publish(MQTT_GAME_TOPIC, message)

def get_player_move():
    try:
        move = eval(input("Enter your move (format: [row, col]): "))
        if not (isinstance(move, list) and len(move) == 2 and
                isinstance(move[0], int) and isinstance(move[1], int)):
            print("Invalid move format. Please enter again.")
            return get_player_move()
        return move
    except:
        print("Invalid input. Please enter again.")
        return get_player_move()

# set up MQTT client
client = MQTTClient(MQTT_CLIENT_ID, MQTT_BROKER)
client.set_callback(handle_move)

# Connect to Wi-Fi and MQTT broker
connect_to_wifi()
connect_to_mqtt_broker()

# Draw initial board and send the first turn message
draw_board()
send_turn_message()

try:
    while True:
        if current_player == "Player 1":
            move = get_player_move()
            message = "{}: {}".format(current_player, move)
            client.publish(MQTT_GAME_TOPIC, message)
            handle_move(None, message.encode("utf-8"))
        client.check_msg()
finally:
    client.disconnect()