#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <ctype.h>
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
Adafruit_SSD1306 display(SCREEN_WIDTH,SCREEN_HEIGHT,&Wire,-1);
// Button pins for input
#define BTN_UP 6
#define BTN_DOWN 7
#define BTN_LEFT 8
#define BTN_RIGHT 9
#define BTN_SELECT 10
// 8x8 Chess board: Uppercase=White,lowercase=Black
char board[8][8]={
    {'r','n','b','q','k','b','n','r'},
    {'p','p','p','p','p','p','p','p'},
    {' ',' ',' ',' ',' ',' ',' ',' '},
    {' ',' ',' ',' ',' ',' ',' ',' '},
    {' ',' ',' ',' ',' ',' ',' ',' '},
    {' ',' ',' ',' ',' ',' ',' ',' '},
    {'P','P','P','P','P','P','P','P'},
    {'R','N','B','Q','K','B','N','R'}
};
int cursorX=4,cursorY=4;
bool pieceSelected=false;
int selX,selY;
bool whiteTurn=true; // Current player
bool gameOver=false;
// Castling rights: [White/Black][KingSide/QueenSide]
bool canCastle[2][2]={{true,true},{true,true}};
// --- Piece Bitmaps ---
// White Pawn 'P'
static const unsigned char PROGMEM whitePawn_8x8[]={
    0x00,0x00,0x1C,0x3E,0x3E,0x1C,0x3E,0x00
};
// White Rook 'R'
static const unsigned char PROGMEM whiteRook_8x8[]={
    0x55,0x7F,0x3E,0x3E,0x3E,0x7F,0x7F,0x00
};
// White Knight 'N'
static const unsigned char PROGMEM whiteKnight_8x8[]={
    0x1C,0x3E,0x7E,0x6E,0x0E,0x3F,0x7F,0x00
};
// White Bishop 'B'
static const unsigned char PROGMEM whiteBishop_8x8[]={
    0x3E,0x5F,0x7F,0x7F,0x3E,0x1C,0x7F,0x00
};
// White Queen 'Q'
static const unsigned char PROGMEM whiteQueen_8x8[]={
    0x2A,0x3E,0x3E,0x1C,0x3E,0x1C,0x7F,0x00
};
// White King 'K'
static const unsigned char PROGMEM whiteKing_8x8[]={
    0x08,0x1C,0x49,0x7F,0x3E,0x1C,0x7F,0x00
};
// Black Pawn 'p'
static const unsigned char PROGMEM blackPawn_8x8[]={
    0x00,0x00,0x1C,0x22,0x22,0x1C,0x7F,0x00
};
// Black Rook 'r'
static const unsigned char PROGMEM blackRook_8x8[]={
    0x55,0x7F,0x22,0x22,0x22,0x41,0x7F,0x00
};
// Black Knight 'n'
static const unsigned char PROGMEM blackKnight_8x8[]={
    0x1C,0x22,0x72,0x6A,0x0A,0x31,0x7F,0x00
};
// Black Bishop 'b'
static const unsigned char PROGMEM blackBishop_8x8[]={
    0x3E,0x41,0x41,0x41,0x22,0x14,0x7F,0x00
};
// Black Queen 'q'
static const unsigned char PROGMEM blackQueen_8x8[]={
    0x2A,0x77,0x22,0x14,0x22,0x14,0x7F,0x00
};
// Black King 'k'
static const unsigned char PROGMEM blackKing_8x8[]={
    0x08, 0x1C, 0x49, 0x77, 0x22, 0x14, 0x7F, 0x00
};
// --- Function Prototypes ---
bool isLegalMove(int x1,int y1,int x2,int y2);
bool isCastlingLegal(int x1,int y1,int x2,int y2);
bool isPathClear(int x1,int y1,int x2,int y2);
bool doesMoveLeaveKingSafe(int x1,int y1,int x2,int y2);
void findKing(bool white,int *kx,int *ky);
bool isSquareAttacked(int tx,int ty,bool byWhite);
bool noLegalMoves(bool white);
void applyMove(int x1,int y1,int x2,int y2);
void updateCastlingRights(int x1,int y1);
void drawBoard(const char* msg);
void setup(){
    // Set up buttons with internal pull-up resistors
    pinMode(BTN_UP, INPUT_PULLUP);
    pinMode(BTN_DOWN, INPUT_PULLUP);
    pinMode(BTN_LEFT, INPUT_PULLUP);
    pinMode(BTN_RIGHT, INPUT_PULLUP);
    pinMode(BTN_SELECT, INPUT_PULLUP);
    // Initialize the OLED display
    if(!display.begin(SSD1306_SWITCHCAPVCC,0x3C)){
        for(;;); // Halt if display not found
    }
    drawBoard("White to move");
}
void loop(){
    static unsigned long lastButtonTime=0;
    unsigned long now=millis();
    if(now-lastButtonTime>150){ // Debounce button presses
        if(!gameOver){
            bool moved=false;
            // Move cursor
            if(!digitalRead(BTN_UP) && cursorY>0){cursorY--; moved=true;}
            else if(!digitalRead(BTN_DOWN) && cursorY<7){cursorY++; moved=true;}
            else if(!digitalRead(BTN_LEFT) && cursorX>0){cursorX--; moved=true;}
            else if(!digitalRead(BTN_RIGHT) && cursorX<7){cursorX++; moved=true;}
            if(moved){
                lastButtonTime=now;
            }
            if(!digitalRead(BTN_SELECT)){
                lastButtonTime=now;
                if(!pieceSelected){
                    // Selection attempt
                    char p=board[cursorY][cursorX];
                    // Check if piece matches current player's color
                    if((whiteTurn && isupper(p)) || (!whiteTurn && islower(p))){
                        selX=cursorX;
                        selY=cursorY;
                        pieceSelected=true;
                    }
                }
                else{
                    // Move attempt: check legality and king safety
                    if(isLegalMove(selX,selY,cursorX,cursorY) && doesMoveLeaveKingSafe(selX,selY,cursorX,cursorY)){
                        applyMove(selX, selY,cursorX,cursorY);
                        whiteTurn=!whiteTurn;
                        // Check for Check, Checkmate, or Stalemate
                        int kx,ky;
                        findKing(whiteTurn,&kx,&ky);
                        if(isSquareAttacked(kx,ky,!whiteTurn)){
                            if(noLegalMoves(whiteTurn)){
                                // Checkmate: Opponent wins
                                drawBoard(whiteTurn?"Black Wins!":"White Wins!");
                                gameOver=true;
                            }
                            else{
                                drawBoard("Check!");
                            }
                        }
                        else if(noLegalMoves(whiteTurn)){
                            // Stalemate: Draw
                            drawBoard("STALEMATE");
                            gameOver=true;
                        }
                        else{
                            drawBoard(whiteTurn?"White to move":"Black to move");
                        }
                    }
                    pieceSelected=false;
                }
            }
        }
    }
    if(!gameOver){
        drawBoard(NULL);
    }
    delay(50);
}
// Executes the move, including special cases like castling and promotion.
void applyMove(int x1,int y1,int x2,int y2){
    char piece=board[y1][x1];
    // Handle Castling: King moves 2 squares horizontally
    if(tolower(piece)=='k' && abs(x2-x1)==2){
        board[y2][x2]=piece; // Move King
        board[y1][x1]=' ';
        // Move Rook
        if(x2==6){ // King-side (O-O)
            board[y2][5]=board[y1][7];
            board[y1][7]=' ';
        }
        else if(x2==2){ // Queen-side (O-O-O)
            board[y2][3]=board[y1][0];
            board[y1][0]=' ';
        }
    }
    else{
        // Normal move
        board[y2][x2]=piece;
        board[y1][x1]=' ';
        // Handle Pawn Promotion (to Queen automatically)
        if(tolower(piece)=='p'){
            if(isupper(piece) && y2==0){ // White pawn to rank 8
                board[y2][x2]='Q';
            }
            else if(islower(piece) && y2==7){ // Black pawn to rank 1
                board[y2][x2]='q';
            }
        }
    }
    // Update castling rights (King or Rook moved)
    updateCastlingRights(x1,y1);
}
// Updates castling flags if the piece at (x1, y1) was a King or Rook.
void updateCastlingRights(int x1,int y1){
    char piece=board[y1][x1];
    int player_index=isupper(piece)?0:1;
    if(tolower(piece)=='k'){
        canCastle[player_index][0]=false;
        canCastle[player_index][1]=false;
    }
    else if(tolower(piece)=='r'){
        // Check which Rook moved (King-side or Queen-side)
        if(x1==0) canCastle[player_index][1]=false; // Queen-side Rook moved
        if(x1==7) canCastle[player_index][0]=false; // King-side Rook moved
    }
}
// --- Drawing the Board ---
void drawBoard(const char* msg){
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    // Draw all pieces
    for(int y=0;y<8;y++){
        for(int x=0;x<8;x++){
            char piece=board[y][x];
            const unsigned char* bitmap=NULL;
            switch(piece){
                case 'P': bitmap=whitePawn_8x8; break;
                case 'R': bitmap=whiteRook_8x8; break;
                case 'N': bitmap=whiteKnight_8x8; break;
                case 'B': bitmap=whiteBishop_8x8; break;
                case 'Q': bitmap=whiteQueen_8x8; break;
                case 'K': bitmap=whiteKing_8x8; break;
                case 'p': bitmap=blackPawn_8x8; break;
                case 'r': bitmap=blackRook_8x8; break;
                case 'n': bitmap=blackKnight_8x8; break;
                case 'b': bitmap=blackBishop_8x8; break;
                case 'q': bitmap=blackQueen_8x8; break;
                case 'k': bitmap=blackKing_8x8; break;
            }
            // Draw the piece bitmap at the correct (x, y) coordinates
            if(bitmap!=NULL){
                display.drawBitmap(x*8,y*8,bitmap,8,8,SSD1306_WHITE);
            }
        }
    }
    // Draw cursor border
    if(!gameOver){
        display.drawRect(cursorX*8,cursorY*8,8,8,SSD1306_WHITE);
        display.drawRect(cursorX*8-1,cursorY*8-1,10,10,SSD1306_WHITE);
    }
    // Draw selection highlight
    if(pieceSelected){
        display.drawRect(selX*8+1,selY*8+1,6,6,SSD1306_WHITE);
    }
    // Display game message/status
    display.setCursor(66,0);
    if(msg!=NULL){
        display.println(msg);
    }
    else{
        display.println(whiteTurn?"White":"Black");
        display.setCursor(66,8);
        display.println("to move");
    }
    display.display();
}
// --- Chess Logic ---
// Checks if a move from (x1, y1) to (x2, y2) is legal based on piece type.
bool isLegalMove(int x1,int y1,int x2,int y2){
    if(x1==x2 && y1==y2) return false;
    char piece=board[y1][x1];
    char target=board[y2][x2];
    bool white=isupper(piece);
    bool target_is_friendly=(white && isupper(target)) || (!white && islower(target));
    if(target!=' ' && target_is_friendly) return false; // Cannot capture friendly piece
    int dx=x2-x1,dy=y2-y1;
    int adx=abs(dx),ady=abs(dy);
    // Castling is handled as a special King move (dx=2)
    if(tolower(piece)=='k' && adx==2 && dy==0){
        return isCastlingLegal(x1,y1,x2,y2);
    }
    switch(tolower(piece)){
        case 'p': // Pawn movement
            if(isupper(piece)){ // White
                if(dx==0 && target==' ' && dy==-1) return true;
                if(dx==0 && target==' ' && dy==-2 && y1==6 && board[5][x1]==' ') return true;
                if(adx==1 && dy==-1 && target!=' ') return true; // Capture
            }
            else{ // Black
                if(dx==0 && target==' ' && dy==1) return true;
                if(dx==0 && target==' ' && dy==2 && y1==1 && board[2][x1]==' ') return true;
                if(adx==1 && dy==1 && target!=' ') return true; // Capture
            }
            return false;
        case 'r': // Rook movement
            if(dx==0 || dy==0) return isPathClear(x1,y1,x2,y2);
            return false;
        case 'n': // Knight movement
            return (adx==1&&ady==2) || (adx==2&&ady==1);
        case 'b': // Bishop movement
            if(adx==ady) return isPathClear(x1,y1,x2,y2);
            return false;
        case 'q': // Queen movement
            if((adx==ady) || (dx==0 || dy==0)) return isPathClear(x1,y1,x2,y2);
            return false;
        case 'k': // King movement
            return (adx<=1 && ady<=1);
    }
    return false;
}
// Checks if castling move is legal (rights, path clear, not in check).
bool isCastlingLegal(int x1,int y1,int x2,int y2){
    bool white=isupper(board[y1][x1]);
    int row=white?7:0;
    int king_index=white?0:1;
    int kx=4; // King's starting column
    if(y1!=row || x1!=kx || !canCastle[king_index][(x2==6)?0:1]) return false;
    // Must not be in check
    if(isSquareAttacked(x1,y1,!white)) return false;
    if(x2==6){ // King-side (O-O)
        // Check squares f/g are empty and not attacked
        if(board[row][5]!=' ' || board[row][6]!=' ') return false;
        if(isSquareAttacked(5,row,!white) || isSquareAttacked(6,row,!white)) return false;
        return true;
    }
    else if(x2==2){ // Queen-side (O-O-O)
        // Check squares b/c/d are empty and c/d are not attacked
        if(board[row][1]!=' ' || board[row][2]!=' ' || board[row][3]!=' ') return false;
        if(isSquareAttacked(2,row,!white) || isSquareAttacked(3,row,!white)) return false;
        return true;
    }
    return false;
}
// Checks if path between two squares is clear (for R, B, Q).
bool isPathClear(int x1,int y1,int x2,int y2){
    int dx=(x2>x1)?1:(x2<x1?-1:0);
    int dy=(y2>y1)?1:(y2<y1?-1:0);
    int cx=x1+dx,cy=y1+dy;
    while(cx!=x2 || cy!=y2){
        if(board[cy][cx]!=' ') return false;
        cx+=dx;
        cy+=dy;
    }
    return true;
}
// Finds the coordinates of the specified King (K or k).
void findKing(bool white,int *kx,int *ky){
    for(int y=0;y<8;y++){
        for(int x=0;x<8;x++){
          char p=board[y][x];
          if(white && p=='K'){*kx=x;*ky=y;return;}
          if(!white && p=='k'){*kx=x;*ky=y;return;}
        }
    }
    *kx=-1;*ky=-1;
}
// Checks if the square (tx, ty) is attacked by the specified color (byWhite).
bool isSquareAttacked(int tx,int ty,bool byWhite){
    for(int y=0;y<8;y++){
        for(int x=0;x<8;x++){
            char p=board[y][x];
            if(p==' ') continue;
            // Check if the piece is the correct attacking color
            if((byWhite && isupper(p)) || (!byWhite && islower(p))){
                // Special check for King-on-King distance (always 1 square)
                if(tolower(p)=='k'){
                    int adx=abs(tx-x),ady=abs(ty-y);
                    if(adx<=1 && ady<=1) return true;
                }
                else{
                    // Simulate an opponent piece at target to check for attack path
                    char temp_target=board[ty][tx];
                    board[ty][tx]=byWhite?'p':'P'; 
                    bool attacks=isLegalMove(x,y,tx,ty);
                    board[ty][tx]=temp_target; // Restore target square
                    if(attacks) return true;
                }
            }
        }
    }
    return false;
}
// Simulates a move to ensure the King is not left in check.
bool doesMoveLeaveKingSafe(int x1,int y1,int x2,int y2){
    char src=board[y1][x1];
    char dst=board[y2][x2];
    bool white=isupper(src);
    int kx_orig,ky_orig;
    findKing(white,&kx_orig,&ky_orig);
    // If Castling, the move is safe as isCastlingLegal already checked path safety
    if(tolower(src)=='k' && abs(x2-x1)==2 && y1==y2) return true;
    // Simulate normal move
    board[y2][x2]=src;board[y1][x1]=' ';
    int kx,ky;
    // Determine King's new position
    if(tolower(src)=='k'){kx=x2; ky=y2;}else{kx=kx_orig; ky=ky_orig;}
    bool safe=!isSquareAttacked(kx,ky,!white);
    // Undo the move
    board[y1][x1]=src;
    board[y2][x2]=dst;
    return safe;
}
// Checks if the current player has any legal, safe moves left (for checkmate/stalemate).
bool noLegalMoves(bool white){
    for(int y=0;y<8;y++){
        for(int x=0;x<8;x++){
            char p=board[y][x];
            if(p==' ') continue;
            // Check only the current player's pieces
            if((white&&isupper(p)) || (!white&&islower(p))){
                for(int yy=0;yy<8;yy++){
                    for(int xx=0;xx<8;xx++){
                      if(isLegalMove(x,y,xx,yy) && doesMoveLeaveKingSafe(x,y,xx,yy)) return false; // Found a legal move
                    }
                }
            }
        }
    }
    return true; // No legal moves found
}