#!/bin/python

import random
from collections import namedtuple
from vector import vec2, vec2_from_text

Piece = namedtuple('Piece', 'side rank')

ranks = 'RNBQKBNR'

def next_turn(turn):
    return 'B' if turn == 'W' else 'W'

def get_turn_pieces(turn, whites, blacks):
    return (whites, blacks) if turn == 'W' else (blacks, whites)

def create_board():
    whites = {}
    whites.update({vec2(c, 0): rank for (c, rank) in zip(range(8), ranks)})
    whites.update({vec2(c, 1): 'P' for c in range(8)})

    blacks = {}
    blacks.update({vec2(c, 7): rank for (c, rank) in zip(range(8), ranks)})
    blacks.update({vec2(c, 6): 'P' for c in range(8)})

    return whites, blacks

def print_board(whites, blacks):
    for y in range(7, -1, -1):
        print(str(y + 1) + ' ', end='')
        for x in range(8):
            pos = vec2(x, y)
            if pos in whites:
                print('+' + whites[pos] + '+ ', end='')
            elif pos in blacks:
                print('-' + blacks[pos] + '- ', end='')
            else:
                print('    ', end='')
        print()
    print('   ' + '   '.join('ABCDEFGH'))

def get_human_move(turn):
    s = input(turn + ' to move: ')
    return tuple(map(vec2_from_text, s.split()))

def get_start_move(start_move):
    print(start_move)
    return tuple(map(vec2_from_text, start_move.split()))
    

def get_computer_move(turn, whites, blacks):
    return random.choice(get_valid_moves(turn, whites, blacks))

def get_valid_dsts(src, turn, whites, blacks):
    # todo: castling
    # todo: en passant
    # todo: promotion
    self_pieces, oppo_pieces = get_turn_pieces(turn, whites, blacks)

    piece = self_pieces[src]
    valid_dsts = set()

    straight_dirs = set([vec2(1, 0), vec2(-1, 0), vec2(0, 1), vec2(0, -1)])
    diagonal_dirs = set([vec2(1, 1), vec2(-1, 1), vec2(1, -1), vec2(-1, -1)])
    def add_dirs(dirs, *, single=False):
        for dir in dirs:
            delta = dir
            while (src + delta).is_valid():
                dst = src + delta
                if dst in self_pieces:
                    break
                valid_dsts.add(dst)
                if dst in oppo_pieces:
                    break
                if single:
                    break
                delta += dir
        
    if piece == 'P': # Pawn
        dy = 1 if turn == 'W' else -1
        sy = 1 if turn == 'W' else 6
        dst_fw1 = src + vec2(0, dy)
        dst_fw2 = src + vec2(0, dy * 2)
        dst_tkl = src + vec2(-1, dy)
        dst_tkr = src + vec2(1, dy)
        if dst_fw1.is_valid() and dst_fw1 not in self_pieces and dst_fw1 not in oppo_pieces:
            valid_dsts.add(dst_fw1)
            if src.y == sy and dst_fw2 not in self_pieces and dst_fw2 not in oppo_pieces:
                valid_dsts.add(dst_fw2)
        if dst_tkl in oppo_pieces:
            valid_dsts.add(dst_tkl)
        if dst_tkr in oppo_pieces:
            valid_dsts.add(dst_tkr)
    elif piece == 'N': # Knight
        knight_dirs = set([vec2(1, 2), vec2(2, 1), vec2(2, -1), vec2(1, -2), vec2(-1, -2), vec2(-2, -1), vec2(-2, 1), vec2(-1, 2)])
        add_dirs(knight_dirs, single=True)
    elif piece == 'R': # Rook
        add_dirs(straight_dirs)
    elif piece == 'B': # Bishop
        add_dirs(diagonal_dirs)
    elif piece == 'Q': # Queen
        add_dirs(straight_dirs | diagonal_dirs)
    elif piece == 'K': # King
        add_dirs(straight_dirs | diagonal_dirs, single=True)
    
    return valid_dsts

def get_valid_moves(turn, whites, blacks):
    self_pieces = whites if turn == 'W' else blacks
    return [(src, dst) for src in self_pieces for dst in get_valid_dsts(src, turn, whites, blacks)]

def next_board(src, dst, turn, whites, blacks):
    self_pieces, oppo_pieces = get_turn_pieces(turn, whites, blacks)
    self_pieces[dst] = self_pieces[src] # Move src -> dst
    del self_pieces[src] # Remove src
    if dst in oppo_pieces: # Remove dst if it was an opponent
        del oppo_pieces[dst]
    return whites, blacks

# Valid moves filtered such that the opponent cannot take the current player's king
def get_legal_moves(turn, whites, blacks):
    self_pieces, oppo_pieces = get_turn_pieces(turn, whites, blacks)
    valid_moves = get_valid_moves(turn, whites, blacks)
    legal_moves = []
    for src, dst in valid_moves:
        n_whites, n_blacks = next_board(src, dst, turn, whites.copy(), blacks.copy())
        n_turn = next_turn(turn)
        n_self_pieces, n_oppo_pieces = get_turn_pieces(n_turn, n_whites, n_blacks)
        n_valid_moves = get_valid_moves(n_turn, n_whites, n_blacks)
        legal = True
        for n_src, n_dst in n_valid_moves:
            if n_dst in n_oppo_pieces and n_oppo_pieces[n_dst] == 'K':
                legal = False
                break
        if legal:
            legal_moves.append((src, dst))
    return legal_moves

def is_in_check(turn, whites, blacks):
    self_pieces, oppo_pieces = get_turn_pieces(turn, whites, blacks)
    n_turn = next_turn(turn)
    threat_moves = get_valid_moves(n_turn, whites, blacks)
    for src, dst in threat_moves:
        if dst in self_pieces and self_pieces[dst] == 'K':
            return True
    return False

turn = 'W'
whites, blacks = create_board()
start_moves = [
    'D2 D4',
    'E7 E5',
    'D4 E5',
    'D8 E7',
    'E5 E6',
    'E7 E6',
    'E2 E4',
    'F7 F5',
    #'E4 F5' # Should not be legal
]
while True:
    print(turn, 'to move')
    print_board(whites, blacks)
    legal_moves = get_legal_moves(turn, whites, blacks)
    print(len(legal_moves), 'legal moves')
    self_pieces, oppo_pieces = get_turn_pieces(turn, whites, blacks)
    if len(legal_moves) == 0:
        print('game is over')
        break
#    if turn == 'W':
    if len(start_moves):
        src, dst = get_start_move(start_moves.pop(0))
    else:
        for src, dst in legal_moves:
            print(f"{self_pieces[src]} {src}{'x' if dst in oppo_pieces else ' '}{dst}")
        while True:
            try:
                src, dst = get_human_move(turn)
                if not (src, dst) in legal_moves:
                    print('not a legal move')
                else:
                    break
            except ValueError:
                print('invalid input')
#    else:
#        src, dst = get_computer_move(turn, whites, blacks)
    
    whites, blacks = next_board(src, dst, turn, whites, blacks)
    turn = next_turn(turn)