from prefect import Task, flow, task, get_run_logger
from prefect.blocks.core import Block
from prefect.orion.schemas.states import Completed, Cancelled, Failed, StateType
import asyncio
import random
from prefect.blocks.system import JSON
from typing import List


# Initialize board

wins = (
    "012",
    "345",
    "678",
    "036",
    "147",
    "258",
    "048",
    "246",
)


def build_square(i: int) -> Task:
    @task(name=f"square({i})")
    async def square():
        return i

    return square


square_ids = list(range(9))
squares = {i: build_square(i) for i in square_ids}


class Board(Block):
    next_player: int
    o_moves: List[int]
    x_moves: List[int]

    def display(self) -> str:
        moves = []
        for i in square_ids:
            if i in self.o_moves:
                moves.append("🍩")
            elif i in self.x_moves:
                moves.append("⚔️")
            else:
                moves.append("  ")
        return "\n--+--+--\n".join("|".join(moves[x : x + 3]) for x in (0, 3, 6))

    def get_avaliable_squares(self):
        squares_taken = set(self.o_moves + self.x_moves)
        return list(set(square_ids) - squares_taken)

    def update_o(self, square):
        o_moves = self.o_moves
        self.o_moves = o_moves + [square]
        self.next_player = 2

    def update_x(self, square):
        x_moves = self.x_moves
        self.x_moves = x_moves + [square]
        self.next_player = 1

    def check_win(self) -> int:
        o_moves = "".join([str(x) for x in sorted(self.o_moves)])
        x_moves = "".join([str(x) for x in sorted(self.x_moves)])
        if any(win in o_moves for win in wins):
            # O wins
            return 1
        if any(win in x_moves for win in wins):
            # X wins
            return 2
        elif len(o_moves) + len(x_moves) == 9:
            # Draw
            return 0
        else:
            # Continue game
            return -1


@task
def create_board(game_no: int, o_start: bool, starting_square: int) -> Board:

    if o_start:
        board = Board(next_player=2, o_moves=[starting_square], x_moves=[])
    else:
        board = Board(next_player=1, x_moves=[starting_square], o_moves=[])

    tally = JSON(value={"n_games": game_no})
    tally.save("tally", overwrite=True)
    board.save(f"n{game_no}", overwrite=True)

    logger = get_run_logger()
    <http://logger.info|logger.info>("\n(Game %s) Starting move:\n%s", game_no, board.display())


@flow
async def nought(game_no: int):
    logger = get_run_logger()

    while True:

        board = await Board.load(f"n{game_no}")
        game_status = board.check_win()

        if game_status == -1 and board.next_player == 1:
            possible_moves = board.get_avaliable_squares()
            square = await squares[random.choice(possible_moves)]()
            board.update_o(square)
            await board.save(f"n{game_no}", overwrite=True)
            <http://logger.info|logger.info>("\n(Game %s) O moves:\n%s", game_no, board.display())

        game_status = board.check_win()
        if game_status == 1:
            return Completed(message=f"(Game {game_no} Over) 🎉 NOUGHT WON")
        elif game_status == 2:
            return Failed(message=f"(Game {game_no} Over) 👻 NOUGHT LOST")
        elif game_status == 0:
            return Cancelled(message=f"(Game {game_no} Over) 🤯 DRAW")

        await asyncio.sleep(1)


@flow
async def crosses(game_no: int):
    logger = get_run_logger()

    while True:

        board = await Board.load(f"n{game_no}")
        game_status = board.check_win()

        if game_status == -1 and board.next_player == 2:
            possible_moves = board.get_avaliable_squares()
            square = await squares[random.choice(possible_moves)]()
            board.update_x(square)
            await board.save(f"n{game_no}", overwrite=True)
            <http://logger.info|logger.info>("\n(Game %s) X moves:\n%s", game_no, board.display())

        game_status = board.check_win()
        if game_status == 2:
            return Completed(message=f"(Game {game_no} Over) 🎉 CROSSES WON")
        elif game_status == 1:
            return Failed(message=f"(Game {game_no} Over) 👻 CROSSES LOST")
        elif game_status == 0:
            return Cancelled(message=f"(Game {game_no} Over) 🤯 DRAW")

        await asyncio.sleep(1)


tally = JSON(value={"n_games": 0})
tally.save("tally", overwrite=True)


@flow
async def tic_tac_toe(o_start: bool, starting_square: int, max_games: int = 100):
    """
    Rules:
    1. First move and player is selected via flow parameters
    2. Loser moves first in next game
    3. Draw spawns two next games
    Subflow states
    - Win = COMPLETED
    - Lose = FAILED
    - Draw = CANCELLED
    """

    # Initialize game
    tally = await JSON.load("tally")
    game_no = tally.value["n_games"] + 1

    if game_no > max_games:
        return Completed(f"All {max_games} games complete.")

    # Play
    create_board(game_no, o_start, starting_square)
    players = [nought(game_no, return_state=True), crosses(game_no, return_state=True)]
    o_state, x_state = await asyncio.gather(*players)

    # Rematch
    new_games = []
    if o_state.message in ["NOUGHT LOST", "DRAW"]:
        new_games.append(
            tic_tac_toe(o_start=False, starting_square=random.choice(square_ids))
        )
    if x_state.message in ["CROSSES LOST", "DRAW"]:
        new_games.append(
            tic_tac_toe(o_start=True, starting_square=random.choice(square_ids))
        )
    await asyncio.gather(*new_games)

    return Completed(message=f"GAME {game_no} OVER")


if __name__ == "__main__":

    while True:
        asyncio.run(tic_tac_toe(o_start=True, starting_square=0))