from move import *
from type_defs import *
from utils import *


class Position:
    def __init__(self):
        self.pboard = PieceBoard()
        self.fen = None
        self.flipped = None
        self.moveList = None
        self.moveHistory = None
        self.stm = None
        self.init()

    def init(self):
        self.pboard.init()
        self.fen = config.startFEN
        self.flipped = False
        self.moveList = []
        self.moveHistory = []
        self.stm = 'w'

    def __repr__(self):
        return self.pboard.__repr__()

    def reset_variables(self):
        self.fen = config.startFEN
        self.flipped = False
        self.moveList = []
        self.moveHistory = []

    def opponent_stm(self):
        return 'w' if self.stm == 'b' else 'b'

    def set(self, fullFen: str, flip=False):
        self.reset_variables()
        self.flipped = flip

        parts = fullFen.split()
        if len(parts) < 2:
            return
        self.fen = parts[0] + ' ' + parts[1]

        self.setFen(parts[0])
        self.stm = parts[1]

        idx = 2
        if idx < len(parts):
            while parts[idx] == '-' or parts[idx].isdigit():
                idx += 1

        if idx < len(parts):
            if parts[idx] == 'moves':
                idx += 1
                while idx < len(parts):
                    self.do_move(parts[idx])
                    idx += 1

    def setFen(self, fenStr: str):
        counter = {}
        for name in config.piece_names:
            counter[name] = 0

        nx, ny = 0, 0
        for c in fenStr:
            if c.isdigit():
                for i in range(int(c)):
                    self.pboard[ny][nx] = None
                    nx += 1
            elif c == '/':
                ny += 1
                nx = 0
            elif c.isalpha():
                name = config.name_dict[c]
                self.pboard[ny][nx] = self.pboard.pieces[name + str(counter[name])]
                self.pboard[ny][nx].set_npos((nx, ny))
                counter[name] += 1
                nx += 1

    def toRawFen(self):
        board = []
        for i in range(10):
            board.append([])
            for j in range(9):
                board[i].append(None)
        for i in range(10):
            for j in range(9):
                if self.pboard[i][j]:
                    board[i][j] = self.pboard[i][j].name
        return board2fen(board) + ' ' + self.stm

    def toMoveFen(self):
        move_list_str = ''
        for m in self.moveList:
            move_list_str += m + ' '
        return self.fen + ' moves ' + move_list_str

    def toggle_stm(self):
        self.stm = 'b' if self.stm == 'w' else 'w'

    def do_move(self, move):
        (from_nx, from_ny), (to_nx, to_ny) = move2pos(move)

        if not self.pboard[from_ny][from_nx]:
            print(f'[Position::do_move] move {move} from_pos empty')
            return

        captured = None
        if self.pboard[to_ny][to_nx]:
            captured = self.pboard[to_ny][to_nx]
            captured.visible = False

        self.pboard[to_ny][to_nx] = self.pboard[from_ny][from_nx]
        self.pboard[to_ny][to_nx].set_npos((to_nx, to_ny))
        self.pboard[from_ny][from_nx] = None

        self.moveList.append(move)
        self.moveHistory.append({'from_nx': from_nx, 'from_ny': from_ny,
                                 'to_nx': to_nx, 'to_ny': to_ny,
                                 'move': move, 'captured': captured})

        self.toggle_stm()

    def undo_move(self):
        if len(self.moveHistory) == 0:
            return

        self.toggle_stm()

        last_move = self.moveHistory[-1]
        from_nx = last_move['from_nx']
        from_ny = last_move['from_ny']
        to_nx = last_move['to_nx']
        to_ny = last_move['to_ny']
        captured = last_move['captured']
        self.pboard[from_ny][from_nx] = self.pboard[to_ny][to_nx]
        self.pboard[from_ny][from_nx].set_npos((from_nx, from_ny))
        self.pboard[to_ny][to_nx] = None

        if captured:
            self.pboard[to_ny][to_nx] = captured
            self.pboard[to_ny][to_nx].set_npos((to_nx, to_ny))
            self.pboard[to_ny][to_nx].visible = True

        self.moveHistory.pop()
        self.moveList.pop()

    def is_move_legal(self, move):
        (from_nx, from_ny), (to_nx, to_ny) = move2pos(move)

        s_board = self.pboard.to_string_board()
        if s_board[to_ny][to_nx] and s_board[to_ny][to_nx][0] == self.stm:
            return False, f'piece on {move}\'s target pos is yours'

        legal_moves = []
        sqs = gen_tsq_by_nxy(s_board, from_nx, from_ny)
        for sq in sqs:
            legal_moves.append(config.idx2sq[from_ny][from_nx] + sq)
        if move not in legal_moves:
            return False, f'illegal move {move}'

        do_move(s_board, move)
        return not is_in_check_s(s_board, self.stm), f'in check after move {move}'

    def is_in_check(self, stm):
        return is_in_check_s(self.pboard.to_string_board(), stm)

    def gen_legal_moves(self, flags=None, banned_moves=None):
        legal_moves = []
        board = self.pboard.to_string_board()
        for ny in range(10):
            for nx in range(9):
                if board[ny][nx] and board[ny][nx][0] == self.stm and (flags is None or board[ny][nx][2] in flags):
                    sqs = gen_tsq_by_nxy(board, nx, ny)
                    for sq in sqs:
                        m = config.idx2sq[ny][nx] + sq
                        assert isinstance(m, str)
                        if self.is_move_legal(m):
                            if banned_moves is None or m not in banned_moves:
                                legal_moves.append(m)
        return legal_moves