#include <assert.h>
#include <ctype.h>
#include <stdio.h>
#include <stdlib.h> // atoi
#include <string.h>
#include "ramjet.h"
#include "bitboard.h"
#include "quadboard.h"
#include "state.h"

STATE st_castle_pl[2] = {0x00FF00, 0XFF0000};
STATE st_castle_sq[64] = {0};

// initialise the constants
void st_init()
{
    int i, j;

    for(i = 0; i < 8; ++i) {
        j = 1 << i;
        st_castle_sq[i] = j << 16;
        st_castle_sq[i + 56] = j << 8;
    }
}

// true when at least one enemy pawn attacks the ep square
static int is_ep_atk(const QUADBOARD* qb, SQUARE sq, PLAYER foe)
{
    assert(!IS_SQ_BAD(sq));
    assert(!IS_PL_BAD(foe));
    BITBOARD foe_p = qb_get_bb_pt(qb, PAWN) & BB_MASK_PL(qb, foe);
    return (bb_ep_atk[sq] & foe_p) != BB_NULL;
}

// set the ep square if it is pseudolegal
STATE st_set_true_ep(STATE st, const QUADBOARD* qb, SQUARE sq)
{
    assert(!IS_SQ_BAD(sq));
    if(is_ep_atk(qb, sq, ST_GET_PL(st)))
        return ST_SET_EP(st, sq);
    else
        return st;
}

// return the first castle target square while erasing it
SQUARE st_castle_sq_next(STATE* st)
{
    assert((*st & ~0xFFFF00) == 0);
    BITBOARD mask = *st;
    int lsb = bb_lsb_next(&mask);
    *st = mask;
    return (lsb + 48) & 0x3F;
}

// scan the side to move
static void st_scan_pl(const char* fen, STATE* st)
{
    PLAYER pl = (tolower(fen[0]) == 'b')? BLACK : WHITE;
    *st = ST_SET_PL(*st, pl);
}

// print the side to move
static void st_print_pl(char* fen, STATE st)
{
    int c = (ST_GET_PL(st) == WHITE)? 'w' : 'b';
    sprintf(fen, "%c", c);
}

// return the lsb or 63 if none is present
static int safe_lsb_sq(BITBOARD bb)
{
    return bb_lsb_sq(bb | BB_SQ(63));
}

// return the msb or 0 if none is present
static int safe_msb_sq(BITBOARD bb)
{
    int sq = 0;
    while(bb)
        sq = bb_lsb_next(&bb);
    return sq;
}

// scan the castling rights
static void st_scan_castle(const char* fen, const QUADBOARD* qb, STATE* st)
{
    static const BITBOARD rank_18 = BB_CONST(0xFF000000000000FF);
    static const BITBOARD rank_1 = BB_CONST(0xFF00000000000000);
    int i, c;

    for(i = 0; fen[i] != '\0'; ++i) {
        c = tolower(fen[i]);

        PLAYER pl = (isupper(fen[i]))? WHITE : BLACK;
        BITBOARD mask = qb_get_bb_pt(qb, ROOK) & BB_MASK_PL(qb, pl) &
                        rank_18 & (rank_1 ^ BB_ALL(pl));

        // X-FEN (outermost rook)
        if(c == 'q')
            *st = ST_SET_CASTLE_SQ(*st, safe_lsb_sq(mask));
        else if(c == 'k')
            *st = ST_SET_CASTLE_SQ(*st, safe_msb_sq(mask));

        // Shredder-FEN (or inner rook)
        else if(c >= 'a' && c <= 'h')
            *st = ST_SET_CASTLE_SQ(*st, SQ_MIRROR(c - 'a', FOE(pl)));
    }
}

// print the castling rights
static void st_print_castle(char* fen, const QUADBOARD* qb, STATE st)
{
    static const BITBOARD rank_18 = BB_CONST(0xFF000000000000FF);
    static const BITBOARD rank_1 = BB_CONST(0xFF00000000000000);
    int flag = 0;
    PLAYER pl;

    // print each player's castling rights
    for(pl = WHITE; pl <= BLACK; ++pl) {
        int k_sq = qb_get_k_sq(qb, pl);
        STATE castle = ST_GET_CASTLE_PL(st, pl);
        BITBOARD mask = qb_get_bb_pt(qb, ROOK) & BB_MASK_PL(qb, pl) &
                        rank_18 & (rank_1 ^ BB_ALL(pl));
        while(castle) {
            int c, r_sq = st_castle_sq_next(&castle);
            if(k_sq > r_sq)
                c = (r_sq == safe_lsb_sq(mask))? 'q': 'a' + SQ_FILE(r_sq); // queenside
            else
                c = (r_sq == safe_msb_sq(mask))? 'k': 'a' + SQ_FILE(r_sq); // kingside
            c = (pl == WHITE)? toupper(c) : tolower(c);
            sprintf(fen, "%c", c);
            fen = &fen[1];
            flag = 1;
        }
    }

    // no castling
    if(!flag)
        sprintf(fen, "-");

    return;
}

// scan the en-passant target square
static void st_scan_ep(const char* fen, const QUADBOARD* qb, STATE* st)
{
    int sq = STR_2_SQ(fen);
    if(!IS_SQ_BAD(sq))
        *st = st_set_true_ep(*st, qb, sq);
}

// print the en-passant target square
static void st_print_ep(char* fen, STATE st)
{
    int sq = ST_GET_EP(st);
    if(IS_SQ_BAD(sq))
        sprintf(fen, "-");
    else
        sprintf(fen, "%s", sq_2_str[sq]);
}

// scan the 50-move rule counter
static void st_scan_50(const char* fen, STATE* st)
{
    int c = atoi(fen);
    if(c >= 0 && c <= 100)
        *st = ST_SET_50(*st, c);
}

// print the 50-move rule counter
static void st_print_50(char* fen, STATE st)
{
    sprintf(fen, "%d", ST_GET_50(st));
}

// read the state from a FEN string
STATE st_scan_fen(char* fen, const QUADBOARD* qb)
{
    assert(!is_qb_bad(qb));

    STATE st = ST_NULL;
    st_scan_pl(grab_token(&fen), &st);
    st_scan_castle(grab_token(&fen), qb, &st);
    st_scan_ep(grab_token(&fen), qb, &st);
    st_scan_50(grab_token(&fen), &st);

    return st;
}

// convert the state to a FEN string
void st_print_fen(char* fen, const QUADBOARD* qb, STATE st)
{
    assert(!is_qb_bad(qb));
    assert(!is_st_bad(qb, st));

    st_print_pl(fen, st);
    st_print_castle(append_space(fen), qb, st);
    st_print_ep(append_space(fen), st);
    st_print_50(append_space(fen), st);
}

// compute a hash digest of the state
HASHKEY st_hash(STATE st)
{
    // random number from www.random.org
    static const HASHKEY chaos = HK_CONST(0xf4ad100e1ec3d386);
    return ST_REP(st) * chaos;
}

// sanity check on the state
int is_st_bad(const QUADBOARD* qb, STATE st)
{
    assert(!is_qb_bad(qb));

    static const BITBOARD rank_1 = BB_CONST(0xFF00000000000000);
    static const BITBOARD rank_8 = BB_CONST(0x00000000000000FF);
    BITBOARD mask_w, mask_b;

    // the opponent's king can be captured
    PLAYER pl = ST_GET_PL(st);
    if(bb_atk_sq(qb, qb_get_k_sq(qb, FOE(pl)), pl) != BB_NULL)
        return 1;

    // out of range en-passant target square
    int ep = ST_GET_EP(st);
    if(IS_SQ_BAD(ep) && ep != SQ_NULL)
        return 2;

    // no victim or attackers for this en-passant target square
    if(!IS_SQ_BAD(ep)) {
        BITBOARD victim = qb_get_bb_pt(qb, PAWN) & BB_MASK_PL(qb, FOE(pl)) & bb_front(BB_SQ(ep), FOE(pl));
        if(victim == BB_NULL || !is_ep_atk(qb, ep, pl))
            return 3;
    }

    // kings out of the first rank
    BITBOARD all_k = qb_get_bb_pt(qb, KING);
    mask_w = all_k & BB_MASK_PL(qb, WHITE) & ~rank_1;
    mask_b = all_k & BB_MASK_PL(qb, BLACK) & ~rank_8;
    if((mask_w && ST_GET_CASTLE_PL(st, WHITE)) ||
       (mask_b && ST_GET_CASTLE_PL(st, BLACK)))
        return 4;

    // no rook on a castle target square
    BITBOARD all_r = qb_get_bb_pt(qb, ROOK);
    mask_w = ((BITBOARD) ST_GET_CASTLE_PL(st, WHITE) << 48) & ~(all_r & BB_MASK_PL(qb, WHITE));
    mask_b = ((BITBOARD) ST_GET_CASTLE_PL(st, BLACK) >> 16) & ~(all_r & BB_MASK_PL(qb, BLACK));
    if(mask_w || mask_b)
        return 5;

    // fifty-move rule counter above 100
    if(ST_GET_50(st) > 100)
        return 6;

    return 0;
}

// test the functions above
int st_test()
{
    QUADBOARD qb, qb2; STATE st, st2;

    // side to move an en-passant target square
    st = ST_SET_PL(ST_NULL, WHITE);
    st = ST_TOGGLE_PL(st);
    if(ST_GET_PL(st) != BLACK || ST_GET_EP(st) != SQ_NULL)
        return 1;

    // castling
    st = ST_SET_CASTLE_SQ(ST_SET_CASTLE_SQ(ST_SET_CASTLE_SQ(ST_NULL, E8), B1), F7);
    if(ST_GET_CASTLE_PL(st, WHITE) != 0x0200 || ST_GET_CASTLE_PL(st, BLACK) != 0x100000)
        return 2;

    // 50-move counter
    st = ST_SET_50(ST_NULL, 34);
    st = ST_INC_50(st);
    if(ST_GET_50(st) != 35)
        return 3;

    // conditional en-passant setter
    qb_scan_fen("r1bq1r2/pp2n1p1/4p1k1/3pPpN1/1b1n2QP/2N5/PP3PP1/R1B1K2R", &qb);
    st = st_set_true_ep(ST_SET_PL(ST_NULL, WHITE), &qb, F6);
    st2 = st_set_true_ep(ST_SET_PL(ST_NULL, BLACK), &qb, H4);
    if(ST_GET_EP(st) != F6 || ST_GET_EP(st2) != SQ_NULL)
        return 4;

    // castle serialiser
    st = st_castle_sq[E8] | st_castle_sq[B1];
    if(st_castle_sq_next(&st) != B1 || st_castle_sq_next(&st) != E8)
        return 5;

    // input-output in FEN format
    char fen[16] = "b Gkq a3 4 11";
    qb_scan_fen("rn2k1r1/ppp1pp1p/3p2p1/5bn1/P7/2N2B2/1PPPPP2/2BNK1RR", &qb);
    st = st_scan_fen(fen, &qb);
    st_print_fen(fen, &qb, st);
    if(strcmp(fen, "b Gqk - 4"))
        return 6;

    // sanity check
    st = ST_SET_EP(ST_NULL, 65);
    qb_scan_fen("7k/8/1p2b3/pP3R2/P4P1q/n7/8/2K5", &qb);
    qb_scan_fen("7r/8/1p2b3/1P3k2/P4P1K/n7/8/2q5", &qb2);
    if(is_st_bad(&qb, ST_SET_EP(ST_NULL, A6)) != 0 ||
       is_st_bad(&qb2, ST_SET_PL(ST_NULL, BLACK)) != 1 ||
       is_st_bad(&qb2, ST_SET_EP(ST_NULL, 65)) != 2 ||
       is_st_bad(&qb2, ST_SET_EP(ST_NULL, A6)) != 3 ||
       is_st_bad(&qb2, ST_SET_CASTLE_SQ(ST_NULL, H8)) != 4 ||
       is_st_bad(&qb, ST_SET_CASTLE_SQ(ST_NULL, H8)) != 5 ||
       is_st_bad(&qb, ST_SET_50(ST_NULL, 123)) != 6)
        return 7;

    return 0;
}
