#include "base.h"

#include <stdio.h>
#include <assert.h>

/* Tests are mostly generated by ChatGPT */

static void print_rook_test(const char *label,
                            enum square_index sq,
                            bitboard all_occ,
                            bitboard own_occ)
{
    printf("\n%s\n", label);
    printf("All occ:\n");
    bitboard_print(all_occ, stdout);
    printf("Own occ:\n");
    bitboard_print(own_occ, stdout);

    const bitboard attacks = rook_attacks_from_index(sq, all_occ) & ~own_occ;
    printf("Rook attacks:\n");
    bitboard_print(attacks, stdout);
}

static void test_rooks()
{
    {
        const enum square_index sq = SQUARE_INDEX_A1;
        const bitboard rook       = SQUARE_MASK_A1;
        const bitboard all_occ    = rook;
        const bitboard own_occ    = rook;

        /* Expected: full rank 1 and file A, except A1 */
        bitboard expected = (FILEMASK_A | RANKMASK_1) & ~SQUARE_MASK_A1;

        print_rook_test("Test 1: Rook at A1, empty board", sq, all_occ, own_occ);

        const bitboard attacks = rook_attacks_from_index(sq, all_occ) & ~own_occ;
        assert(attacks == expected);
    }

    {
        const enum square_index sq = SQUARE_INDEX_A1;
        const bitboard rook        = SQUARE_MASK_A1;
        const bitboard own_block   = SQUARE_MASK_A2 | SQUARE_MASK_B1;
        const bitboard all_occ     = rook | own_block;
        const bitboard own_occ     = all_occ;

        /* Expected: no legal moves (immediately blocked both directions) */
        const bitboard expected = 0ULL;

        print_rook_test("Test 2: Rook at A1, own blockers A2, B1", sq, all_occ, own_occ);

        const bitboard attacks = rook_attacks_from_index(sq, all_occ) & ~own_occ;
        assert(attacks == expected);
    }

    {
        const enum square_index sq = SQUARE_INDEX_A1;
        const bitboard rook        = SQUARE_MASK_A1;
        const bitboard enemies     = SQUARE_MASK_A3 | SQUARE_MASK_C1;
        const bitboard all_occ     = rook | enemies;
        const bitboard own_occ     = rook;

        /*
         * Expected:
         *  - Along file A: A2, A3 (enemy at A3 is capturable, stop there)
         *  - Along rank 1: B1, C1 (enemy at C1 capturable, stop there)
         */
        bitboard expected = 0ULL;
        expected |= SQUARE_MASK_A2 | SQUARE_MASK_A3;
        expected |= SQUARE_MASK_B1 | SQUARE_MASK_C1;

        print_rook_test("Test 3: Rook at A1, enemy blockers A3, C1", sq, all_occ, own_occ);

        const bitboard attacks = rook_attacks_from_index(sq, all_occ) & ~own_occ;
        assert(attacks == expected);
    }

        /* Rook Test 6: center rook on empty board */
    {
        const enum square_index sq = SQUARE_INDEX_E5;
        const bitboard rook        = SQUARE_MASK_E5;
        const bitboard all_occ     = rook;
        const bitboard own_occ     = rook;

        /* Full rank 5 and file E, except E5 itself */
        bitboard expected = (FILEMASK_E | RANKMASK_5) & ~SQUARE_MASK_E5;

        print_rook_test("Rook Test 6: E5, empty board", sq, all_occ, own_occ);

        bitboard attacks = rook_attacks_from_index(sq, all_occ) & ~own_occ;
        assert(attacks == expected);
    }

    /* Rook Test 7: center rook, mixed blockers on rays */
    {
        const enum square_index sq = SQUARE_INDEX_E5;
        const bitboard rook        = SQUARE_MASK_E5;

        /* Friendly: E7 and C5; Enemy: E3 and H5 */
        const bitboard friends     = rook | SQUARE_MASK_E7 | SQUARE_MASK_C5;
        const bitboard enemies     = SQUARE_MASK_E3 | SQUARE_MASK_H5;
        const bitboard all_occ     = friends | enemies;
        const bitboard own_occ     = friends;

        /*
         * From E5:
         *  Up:   E6, then friendly E7 (stop before E7)
         *  Down: E4, E3 (enemy, included, then stop)
         *  Left: D5, then friendly C5 (stop before C5)
         *  Right:F5, G5, H5 (enemy, included, then stop)
         */
        bitboard expected = 0ULL;
        expected |= SQUARE_MASK_E6;
        expected |= SQUARE_MASK_E4 | SQUARE_MASK_E3;
        expected |= SQUARE_MASK_D5;
        expected |= SQUARE_MASK_F5 | SQUARE_MASK_G5 | SQUARE_MASK_H5;

        print_rook_test("Rook Test 7: E5, friends E7/C5, enemies E3/H5", sq, all_occ, own_occ);

        bitboard attacks = rook_attacks_from_index(sq, all_occ) & ~own_occ;
        assert(attacks == expected);
    }

    /* Rook Test 8: edge rook on empty board (top edge, not corner) */
    {
        const enum square_index sq = SQUARE_INDEX_C8;
        const bitboard rook        = SQUARE_MASK_C8;
        const bitboard all_occ     = rook;
        const bitboard own_occ     = rook;

        /*
         * From C8:
         *  Down file C: C7..C1
         *  Across rank 8: A8,B8,D8,E8,F8,G8,H8
         */
        bitboard expected = 0ULL;
        expected |= (FILEMASK_C & ~SQUARE_MASK_C8);
        expected |= (RANKMASK_8 & ~SQUARE_MASK_C8);

        print_rook_test("Rook Test 8: C8, empty board", sq, all_occ, own_occ);

        bitboard attacks = rook_attacks_from_index(sq, all_occ) & ~own_occ;
        assert(attacks == expected);
    }

    /* Rook Test 9: rook completely boxed in by friendly orthogonal neighbors */
    {
        const enum square_index sq = SQUARE_INDEX_D4;
        const bitboard rook        = SQUARE_MASK_D4;
        const bitboard friends     = rook |
                                     SQUARE_MASK_D5 | SQUARE_MASK_D3 |
                                     SQUARE_MASK_C4 | SQUARE_MASK_E4;
        const bitboard enemies     = 0ULL;
        const bitboard all_occ     = friends | enemies;
        const bitboard own_occ     = friends;

        /* All four rays are immediately blocked by own pieces */
        const bitboard expected = 0ULL;

        print_rook_test("Rook Test 9: D4, boxed by own pieces at D5/D3/C4/E4",
                        sq, all_occ, own_occ);

        bitboard attacks = rook_attacks_from_index(sq, all_occ) & ~own_occ;
        assert(attacks == expected);
    }

    /* Rook Test 10: rook on file with non-interfering off-ray pieces */
    {
        const enum square_index sq = SQUARE_INDEX_A4;
        const bitboard rook        = SQUARE_MASK_A4;

        /* Pieces placed off the rook's rank/file; they should have no effect */
        const bitboard off_ray     = SQUARE_MASK_C1 | SQUARE_MASK_F6 | SQUARE_MASK_H8;
        (void)off_ray;
        const bitboard friends     = rook | SQUARE_MASK_C1;
        const bitboard enemies     = SQUARE_MASK_F6 | SQUARE_MASK_H8;
        const bitboard all_occ     = friends | enemies;
        const bitboard own_occ     = friends;

        /*
         * From A4:
         *  File A: A1..A8 except A4
         *  Rank 4: B4..H4
         * Pieces not on file A or rank 4 must not change attacks.
         */
        bitboard expected = 0ULL;
        expected |= (FILEMASK_A | RANKMASK_4) & ~SQUARE_MASK_A4;

        print_rook_test("Rook Test 10: A4, random off-ray pieces C1/F6/H8",
                        sq, all_occ, own_occ);

        bitboard attacks = rook_attacks_from_index(sq, all_occ) & ~own_occ;
        assert(attacks == expected);
    }
}


static void print_bishop_test(const char *label,
                              enum square_index sq,
                              bitboard all_occ,
                              bitboard own_occ)
{
    printf("\n%s\n", label);
    printf("All occ:\n");
    bitboard_print(all_occ, stdout);
    printf("Own occ:\n");
    bitboard_print(own_occ, stdout);

    const bitboard attacks = bishop_attacks_from_index(sq, all_occ) & ~own_occ;
    printf("Bishop attacks:\n");
    bitboard_print(attacks, stdout);
}

static void test_bishops(void)
{
    /* Test 1: Bishop at D4 on empty board (only bishop present) */
    {
        const enum square_index sq = SQUARE_INDEX_D4;
        const bitboard bishop      = SQUARE_MASK_D4;
        const bitboard all_occ     = bishop;
        const bitboard own_occ     = bishop;

        /*
         * Expected diagonals from D4:
         *  NE: E5, F6, G7, H8
         *  NW: C5, B6, A7
         *  SE: E3, F2, G1
         *  SW: C3, B2, A1
         */
        bitboard expected = 0ULL;
        expected |= SQUARE_MASK_E5 | SQUARE_MASK_F6 | SQUARE_MASK_G7 | SQUARE_MASK_H8;
        expected |= SQUARE_MASK_C5 | SQUARE_MASK_B6 | SQUARE_MASK_A7;
        expected |= SQUARE_MASK_E3 | SQUARE_MASK_F2 | SQUARE_MASK_G1;
        expected |= SQUARE_MASK_C3 | SQUARE_MASK_B2 | SQUARE_MASK_A1;

        print_bishop_test("Bishop Test 1: D4, empty board", sq, all_occ, own_occ);

        const bitboard attacks = bishop_attacks_from_index(sq, all_occ) & ~own_occ;
        assert(attacks == expected);
    }

    /* Test 2: Bishop at C1 on empty board (only bishop present) */
    {
        const enum square_index sq = SQUARE_INDEX_C1;
        const bitboard bishop      = SQUARE_MASK_C1;
        const bitboard all_occ     = bishop;
        const bitboard own_occ     = bishop;

        /*
         * From C1, diagonals:
         *  NE: D2, E3, F4, G5, H6
         *  NW: B2, A3
         *  SE / SW: none (edge of board)
         */
        bitboard expected = 0ULL;
        expected |= SQUARE_MASK_D2 | SQUARE_MASK_E3 | SQUARE_MASK_F4 |
                    SQUARE_MASK_G5 | SQUARE_MASK_H6;
        expected |= SQUARE_MASK_B2 | SQUARE_MASK_A3;

        print_bishop_test("Bishop Test 2: C1, empty board", sq, all_occ, own_occ);

        const bitboard attacks = bishop_attacks_from_index(sq, all_occ) & ~own_occ;
        assert(attacks == expected);
    }

    /* Test 3: Bishop at D4, friendly blockers at F6 and B2 (no enemies) */
    {
        const enum square_index sq = SQUARE_INDEX_D4;
        const bitboard bishop      = SQUARE_MASK_D4;
        const bitboard friends     = bishop | SQUARE_MASK_F6 | SQUARE_MASK_B2;
        const bitboard enemies     = 0ULL;
        const bitboard all_occ     = friends | enemies;
        const bitboard own_occ     = friends;

        /*
         * From D4:
         *  NE: E5, then blocked by friendly F6 (F6 not included)
         *  SW: C3, then blocked by friendly B2 (B2 not included)
         *  NW: C5, B6, A7 (no blockers)
         *  SE: E3, F2, G1 (no blockers)
         */
        bitboard expected = 0ULL;
        expected |= SQUARE_MASK_E5;
        expected |= SQUARE_MASK_C3;
        expected |= SQUARE_MASK_C5 | SQUARE_MASK_B6 | SQUARE_MASK_A7;
        expected |= SQUARE_MASK_E3 | SQUARE_MASK_F2 | SQUARE_MASK_G1;

        print_bishop_test("Bishop Test 3: D4, friendly blockers F6, B2", sq, all_occ, own_occ);

        const bitboard attacks = bishop_attacks_from_index(sq, all_occ) & ~own_occ;
        assert(attacks == expected);
    }

    /* Test 4: Bishop at D4, enemy blockers at F6 and B2 (no other friends) */
    {
        const enum square_index sq = SQUARE_INDEX_D4;
        const bitboard bishop      = SQUARE_MASK_D4;
        const bitboard friends     = bishop;
        const bitboard enemies     = SQUARE_MASK_F6 | SQUARE_MASK_B2;
        const bitboard all_occ     = friends | enemies;
        const bitboard own_occ     = friends;

        /*
         * From D4:
         *  NE: E5, F6 (enemy, included, then stop)
         *  SW: C3, B2 (enemy, included, then stop)
         *  NW: C5, B6, A7
         *  SE: E3, F2, G1
         */
        bitboard expected = 0ULL;
        expected |= SQUARE_MASK_E5 | SQUARE_MASK_F6;
        expected |= SQUARE_MASK_C3 | SQUARE_MASK_B2;
        expected |= SQUARE_MASK_C5 | SQUARE_MASK_B6 | SQUARE_MASK_A7;
        expected |= SQUARE_MASK_E3 | SQUARE_MASK_F2 | SQUARE_MASK_G1;

        print_bishop_test("Bishop Test 4: D4, enemy blockers F6, B2", sq, all_occ, own_occ);

        const bitboard attacks = bishop_attacks_from_index(sq, all_occ) & ~own_occ;
        assert(attacks == expected);
    }

    /* Test 5: Bishop at D4, mixed friend/enemy + another friendly bishop elsewhere */
    {
        const enum square_index sq = SQUARE_INDEX_D4;
        const bitboard bishop1     = SQUARE_MASK_D4;  /* tested bishop */
        const bitboard bishop2     = SQUARE_MASK_F4;  /* another friendly bishop */
        const bitboard friends     = bishop1 | bishop2 | SQUARE_MASK_F6;
        const bitboard enemies     = SQUARE_MASK_B2;
        const bitboard all_occ     = friends | enemies;
        const bitboard own_occ     = friends;

        /*
         * From D4:
         *  NE: E5, then friendly F6 (stop; F6 not included)
         *  SW: C3, B2 (enemy, included, then stop)
         *  NW: C5, B6, A7
         *  SE: E3, F2, G1
         *  Bishop at F4 is irrelevant; it does not sit on a diagonal from D4.
         */
        bitboard expected = 0ULL;
        expected |= SQUARE_MASK_E5;
        expected |= SQUARE_MASK_C3 | SQUARE_MASK_B2;
        expected |= SQUARE_MASK_C5 | SQUARE_MASK_B6 | SQUARE_MASK_A7;
        expected |= SQUARE_MASK_E3 | SQUARE_MASK_F2 | SQUARE_MASK_G1;

        print_bishop_test("Bishop Test 5: D4, mixed friend/enemy + extra bishop F4", sq, all_occ, own_occ);

        const bitboard attacks = bishop_attacks_from_index(sq, all_occ) & ~own_occ;
        assert(attacks == expected);
    }

    printf("\nAll bishop_attacks_from_index tests passed.\n");
}

int main()
{
    test_rooks();
    test_bishops();

    for (int i = 40; i < 47; i++) {
        fprintf(stdout, "\033[30;%dm ", i);
    }
    fprintf(stdout, "\033[0m\n"); /* reset background color */
    struct board board = BOARD_INITIAL;

    board_print(&board, stdout);

    printf("\nFile C3:\n");
    bitboard_print(FILEMASK_C, stdout);

    printf("\nRank 2:\n");
    bitboard_print(RANKMASK_2, stdout);

    printf("\nKnight at C3 moves:\n");
    bitboard_print(knight_attacks(SQUARE_MASK_C3), stdout);

    printf("\nKnight at B2 moves:\n");
    bitboard_print(knight_attacks(SQUARE_MASK_B2), stdout);

    printf("\nKnight at A1, H8 moves:\n");
    bitboard_print(knight_attacks(SQUARE_MASK_A1 | SQUARE_MASK_H8), stdout);

    printf("\nWhite pawn at E4 moves:\n");
    bitboard_print(
            pawn_white_moves(
                /* p */               SQUARE_MASK_E4,
                /* empty */          ~0ULL 
                ),
            stdout);

    printf("\nWhite pawns at E2, F4, A3 moves:\n");
    bitboard_print(
            pawn_white_moves(
                /* p */      SQUARE_MASK_E2 | SQUARE_MASK_F4 | SQUARE_MASK_A3,
                /* empty */ ~0ULL 
                ),
            stdout);

    return EXIT_SUCCESS;

}