#include "Board.hpp"
#include "BoardState.hpp"

#include <ostream>
#include <cassert>
#include <cmath>
#include <bitset>
#include <algorithm>

void Board::setPiece(const Square &square, const Piece::Optional &piece) {
    if (!piece.has_value())
        return;

    auto index = square.index();
    for (int i = 0; i < BB_NUM; i++) {
        mPieceBBs[i].clear(index);
    }

    mPieceBBs[toIndex(piece->type())].set(index);
    mPieceBBs[toIndex(piece->color())].set(index);
    mOccupiedBB.set(index);
}

Piece::Optional Board::piece(const Square &square) const {
    BitBoard mask = BitBoard::fromIndex(square.index());
    if (!(mOccupiedBB & mask)) {
        return std::nullopt;
    }

    auto c = pieceColor(mask);
    auto t = pieceType(mask);
    return Piece(c, t);
}

int Board::evaluate() const {
    int score = evaluate(mTurn);
    score -= evaluate(!mTurn);

    return score;
}

int Board::evaluate(const PieceColor color) const {
    // Pawns
    int score = 0;

    BitBoard bb;
    BitBoard colorMask = mPieceBBs[toIndex(color)];
    for (int i = 2; i < 8; i++) {
        bb = mPieceBBs[i] & colorMask;
        while (bb) {
            bb.pop();
            score += mPieceValue[i];
        }
    }

    return score;
}

bool Board::isCheckMate() const {
    auto moves = MoveVec();
    pseudoLegalMoves(moves);

    if (!moves.empty())
        return false;

    BoardState bs = BoardState(&mPieceBBs, mOccupiedBB, mTurn, mCR, mEPS);
    auto attacked = MoveGenerator::generateAttackedSquares(bs, ~mOccupiedBB, !mTurn);

    return attacked & mPieceBBs[toIndex(PieceType::King)] & mPieceBBs[toIndex(mTurn)];

    return moves.empty();
}

void Board::setTurn(PieceColor turn) {
    mTurn = turn;
}

PieceColor Board::turn() const {
    return mTurn;
}

void Board::setCastlingRights(CastlingRights cr) {
    mCR = cr;
}

CastlingRights Board::castlingRights() const {
    return mCR;
}

void Board::setEnPassantSquare(const Square::Optional &square) {
    if (!square.has_value())
        return;
    mEPS = square;
}

Square::Optional Board::enPassantSquare() const {
    return mEPS;
}

void Board::makeMove(const Move &move) {

    BitBoard fromBB = BitBoard::fromIndex(move.from().index());
    BitBoard toBB = BitBoard::fromIndex(move.to().index());
    BitBoard changeBB = fromBB ^ toBB;

    // If Piece is captured
    if (mOccupiedBB & toBB) {
        auto capturedPiece = Piece(!mTurn, pieceType(toBB));
        mPieceBBs[toIndex(capturedPiece.color())] ^= toBB;
        mPieceBBs[toIndex(capturedPiece.type())] ^= toBB;
        mOccupiedBB ^= fromBB;

        if (toBB & CastlingRanks) { // Check castling rights
            handleCastlingRights(capturedPiece, toBB);
        }
    } else {
        mOccupiedBB ^= changeBB; // update occupied bitboard
    }

    auto movedPiece = Piece(mTurn, pieceType(fromBB));
    mPieceBBs[toIndex(movedPiece.color())] ^= changeBB; // update color bitboard

    // TODO verify if we should check if piece == Pawn
    if (move.promotion().has_value()) {
        mPieceBBs[toIndex(movedPiece.type())] ^= fromBB; // Remove old piece
        mPieceBBs[toIndex(move.promotion().value())] ^= toBB; // Set new piece
    } else {
        mPieceBBs[toIndex(movedPiece.type())] ^= changeBB; // remove and set piece bitboard
    }

    if ((fromBB & CastlingRanks)) {
        handleCastlingRights(movedPiece, fromBB);

        if (isMoveCastling(fromBB, toBB, movedPiece)) {
            BitBoard rookBB;
            if (toBB & GFile) { // Kingside
                rookBB = (fromBB << 3) | (fromBB << 1);
            } else {
                rookBB = (fromBB >> 4) | (fromBB >> 1);
            }

            mPieceBBs[toIndex(PieceType::Rook)] ^= rookBB;
            mPieceBBs[toIndex(movedPiece.color())] ^= rookBB;
            mOccupiedBB ^= rookBB;
        }
    }

    handleEnPassant(move, movedPiece);

    handlePawnDoubleAdvance(move, toBB, movedPiece);


    // change turn
    mTurn = !mTurn;
}
void Board::handleEnPassant(const Move &move, const Piece &movedPiece) {
    if (movedPiece.type() == PieceType::Pawn && mEPS.has_value() && move.from().file() != move.to().file()) {
        auto epBB = BitBoard::fromIndex(Square::fromCoordinates(move.to().file(), move.from().rank())->index());
        auto capturedPiece = Piece(!mTurn, PieceType::Pawn);

        mPieceBBs[toIndex(capturedPiece.color())] ^= epBB;
        mPieceBBs[toIndex(capturedPiece.type())] ^= epBB;
        mOccupiedBB ^= epBB;
    }
}

void Board::handlePawnDoubleAdvance(const Move &move, BitBoard bb, const Piece &movedPiece) {
    if (movedPiece.type() == PieceType::Pawn) {
        auto fromR = move.from().rank();
        auto toR = move.to().rank();
        if (abs(static_cast<int>(fromR) - static_cast<int>(toR)) == 2) {
            if (mPieceBBs[toIndex(PieceType::Pawn)] & mPieceBBs[toIndex(!mTurn)] // Get all other color pawns
                    & (bb.west() | bb.east())) { // is pawn left or right?
                mEPS = Square::fromCoordinates(move.to().file(), std::max(fromR, toR) - 1);
                return;
            }
        }

    }

    mEPS = std::nullopt;
}

void Board::pseudoLegalMoves(MoveVec &moves) const {
    auto pieces = mPieceBBs[toIndex(mTurn)];

    while (pieces) {
        auto sq = Square(pieces.pop());
        pseudoLegalMovesFrom(sq, moves);
    }
}

void Board::pseudoLegalMovesFrom(const Square &from, Board::MoveVec &moves) const {
    auto fromBB = BitBoard::fromIndex(from.index());

    if (!(fromBB & mPieceBBs[toIndex(mTurn)])) {
        return;
    }

    BoardState bs = BoardState(&mPieceBBs, mOccupiedBB, mTurn, mCR, mEPS);

    switch (pieceType(fromBB)) {
    case PieceType::Pawn: MoveGenerator::generatePawnMoves(bs, from, moves);
        break;
    case PieceType::Knight: MoveGenerator::generateKnightMoves(bs, from, moves);
        break;
    case PieceType::Bishop: MoveGenerator::generateBishopMoves(bs, from, moves);
        break;
    case PieceType::Rook: MoveGenerator::generateRookMoves(bs, from, moves);
        break;
    case PieceType::Queen: MoveGenerator::generateQueenMoves(bs, from, moves);
        break;
    case PieceType::King: MoveGenerator::generateKingMoves(bs, from, moves);
        break;
    }
}
void Board::handleCastlingRights(const Piece &piece, const BitBoard &bb) {
    if (piece.type() != PieceType::King && piece.type() != PieceType::Rook) {
        return;
    }

    auto crColor = CastlingRights::White;
    auto crKingside = CastlingRights::WhiteKingside;
    auto crQueenside = CastlingRights::WhiteQueenside;

    if (piece.color() == PieceColor::Black) {
        crColor = CastlingRights::Black;
        crKingside = CastlingRights::BlackKingside;
        crQueenside = CastlingRights::BlackQueenside;
    }

    if ((mCR & crColor) == CastlingRights::None) {
        return;
    }

    if (piece.type() == PieceType::King) {
        mCR &= ~crColor;
        return;
    }

    if (piece.type() == PieceType::Rook) {
        if (bb & (HFile & CastlingRanks)) { // Kingside
            mCR &= ~crKingside;
        } else if (bb & (AFile & CastlingRanks)) { // Queenside
            mCR &= ~crQueenside;
        }
    }
}

bool Board::isMoveCastling(const BitBoard &from, const BitBoard &to, const Piece &piece) {
    if (piece.type() != PieceType::King)
        return false;

    if (!(to & CastlingSquares)) {
        return false;
    }

    if (piece.color() == PieceColor::White) {
        return (from & BitBoard::fromIndex(E1));
    }

    return from & BitBoard::fromIndex(E8);
}

std::ostream &operator<<(std::ostream &os, const Board &board) {
    // For debugging only, performance isn't important
    for (int i = 7; i >= 0; i--) {
        int rank = i * 8;
        for (int j = 0; j < 8; j++) {
            // Get the piece for this index. Assume it exists.
            auto piece = board.piece(Square::fromIndex(rank + j).value());

            // Print piece, otherwise '.';
            if (piece.has_value()) {
                os << piece.value();
            } else {
                os << '.';
            }

            os << ' ';
        }
        os << '\n';
    }

    return os;
}