#include <Arduino.h>
#include <FastLED.h>
#include <Encoder.h>

#define NUM_LEDS 64
const int COLS = 7;
const int ROWS = 6;
const int LOOK_AHEAD = 6; // Depth 6 is very stable and tough for C3

CRGB leds[NUM_LEDS];
Encoder *myEnc; 

int8_t board[COLS][ROWS]; 
bool winMask[NUM_LEDS];   
enum State { MENU, PLAYING, FINISHED_WIN, FINISHED_DRAW };
State gameState = MENU;

int8_t menuMode = 0;      // 0: P1-Yellow, 1: P1-Red, 2: PvP
int8_t currentPlayer = 1; // 1: Yellow, 2: Red
int8_t activeCol = 3;
long oldEncPos = -999;
uint8_t aiBrightness = 0;
bool aiFadeUp = true;

// --- Helper Functions ---
int getIdx(int x, int y) { return (y * 8) + x; }

void drawStaticUI() {
    FastLED.clear();
    for(int x = 0; x < 7; x++) leds[getIdx(x, 1)] = CRGB::Blue; 
    for(int y = 1; y < 8; y++) leds[getIdx(7, y)] = CRGB::Blue; 
}

void renderBoard() {
    drawStaticUI();
    for(int c=0; c<COLS; c++) {
        for(int r=0; r<ROWS; r++) {
            if(board[c][r] == 1) leds[getIdx(c, 7-r)] = CRGB::Yellow;
            if(board[c][r] == 2) leds[getIdx(c, 7-r)] = CRGB::Red;
        }
    }
}

int getFirstEmptyRow(int col) {
    for (int r = 0; r < ROWS; r++) { if (board[col][r] == 0) return r; }
    return -1;
}

bool isBoardFull() {
    for (int c = 0; c < COLS; c++) if (board[c][5] == 0) return false;
    return true;
}

// Scans board and fills winMask if 4+ connected
bool scanBoard(int8_t p) {
    bool found = false;
    memset(winMask, 0, sizeof(winMask));
    auto checkLine = [&](int x, int y, int dx, int dy) {
        int count = 0;
        for (int i = 0; i < 7; i++) {
            int nx = x + i * dx; int ny = y + i * dy;
            if (nx >= 0 && nx < COLS && ny >= 0 && ny < ROWS && board[nx][ny] == p) {
                count++;
            } else {
                if (count >= 4) {
                    for (int j = 1; j <= count; j++) 
                        winMask[getIdx(nx - j * dx, 7 - (ny - j * dy))] = true;
                    found = true;
                }
                count = 0;
            }
        }
    };
    for (int i = 0; i < ROWS; i++) checkLine(0, i, 1, 0); 
    for (int i = 0; i < COLS; i++) checkLine(i, 0, 0, 1); 
    for (int i = -5; i < 7; i++) { checkLine(i, 0, 1, 1); checkLine(i, 5, 1, -1); }
    return found;
}

// --- AI Thinking Visualization ---
void updateThinkingLED() {
    static uint32_t lastCycle = 0;
    if (millis() - lastCycle < 20) return; 
    lastCycle = millis();
    if (aiFadeUp) { aiBrightness += 15; if (aiBrightness >= 240) aiFadeUp = false; }
    else { aiBrightness -= 15; if (aiBrightness <= 15) aiFadeUp = true; }
    
    // Pulse in the computer's color
    CRGB compColor = (menuMode == 0) ? CRGB::Red : CRGB::Yellow;
    leds[getIdx(7, 0)] = compColor.nscale8(aiBrightness);
    FastLED.show();
}

// --- Minimax Logic ---
int minimax(int depth, int alpha, int beta, bool isMax, int8_t aiP, int8_t huP) {
    // Check wins within the simulation
    if (scanBoard(aiP)) return 1000 + depth;
    if (scanBoard(huP)) return -1000 - depth;
    if (depth == 0 || isBoardFull()) return 0;

    int order[] = {3, 2, 4, 1, 5, 0, 6};
    if (isMax) {
        int maxEval = -2000;
        for (int c : order) {
            int r = getFirstEmptyRow(c);
            if (r != -1) {
                board[c][r] = aiP;
                int eval = minimax(depth - 1, alpha, beta, false, aiP, huP);
                board[c][r] = 0;
                maxEval = max(maxEval, eval);
                alpha = max(alpha, eval);
                if (beta <= alpha) break;
            }
        }
        return maxEval;
    } else {
        int minEval = 2000;
        for (int c : order) {
            int r = getFirstEmptyRow(c);
            if (r != -1) {
                board[c][r] = huP;
                int eval = minimax(depth - 1, alpha, beta, true, aiP, huP);
                board[c][r] = 0;
                minEval = min(minEval, eval);
                beta = min(beta, eval);
                if (beta <= alpha) break;
            }
        }
        return minEval;
    }
}

void performAiMove() {
    int8_t aiP = (menuMode == 0) ? 2 : 1; // AI is Red if player chose Yellow
    int8_t huP = (menuMode == 0) ? 1 : 2;
    
    aiBrightness = 0; aiFadeUp = true;
    int bestScore = -30000;
    int bestCol = -1;

    // 1. Immediate Win/Block Check
    for(int c=0; c<COLS; c++) {
        int r = getFirstEmptyRow(c);
        if(r != -1) {
            board[c][r] = aiP; if(scanBoard(aiP)) { leds[getIdx(7, 0)] = CRGB::Black; return; }
            board[c][r] = huP; if(scanBoard(huP)) { board[c][r] = aiP; leds[getIdx(7, 0)] = CRGB::Black; return; }
            board[c][r] = 0;
        }
    }

    // 2. Recursive Search
    for (int c : {3, 2, 4, 1, 5, 0, 6}) {
        int r = getFirstEmptyRow(c);
        if (r != -1) {
            board[c][r] = aiP;
            int score = minimax(LOOK_AHEAD, -30000, 30000, false, aiP, huP);
            board[c][r] = 0;
            updateThinkingLED(); // Visual feedback
            if (score > bestScore) { bestScore = score; bestCol = c; }
        }
    }
    if (bestCol != -1) board[bestCol][getFirstEmptyRow(bestCol)] = aiP;
    leds[getIdx(7, 0)] = CRGB::Black;
}

void showMenu() {
    drawStaticUI();
    if (menuMode < 2) {
        CRGB p1Col = (menuMode == 1) ? CRGB::Red : CRGB::Yellow;
        for(int y=3; y<=6; y++) leds[getIdx(3, y)] = p1Col;
        leds[getIdx(2,3)] = p1Col; leds[getIdx(4,3)] = p1Col;
        leds[getIdx(2,6)] = p1Col; leds[getIdx(4,6)] = p1Col;
    } else {
        for(int y=3; y<=6; y++) { leds[getIdx(2, y)] = CRGB::Yellow; leds[getIdx(4, y)] = CRGB::Red; }
        leds[getIdx(1, 3)] = CRGB::Yellow; leds[getIdx(1, 6)] = CRGB::Yellow;
        leds[getIdx(5, 3)] = CRGB::Red; leds[getIdx(5, 6)] = CRGB::Red;
        leds[getIdx(3, 3)] = CRGB::Red; leds[getIdx(3, 6)] = CRGB::Yellow;
    }
    FastLED.show();
}

void setup() {
    Serial.begin(115200);
    myEnc = new Encoder(ENC_A, ENC_B);
    FastLED.addLeds<WS2812B, LED_PIN, GRB>(leds, NUM_LEDS);
    FastLED.setBrightness(BRIGHTNESS);
    pinMode(ENC_SW, INPUT_PULLUP);
    showMenu();
}

void loop() {
    long newPos = myEnc->read() / SENSITIVITY;
    bool pressed = (digitalRead(ENC_SW) == LOW);

    if (gameState == MENU) {
        if (newPos != oldEncPos) {
            menuMode = (newPos % 3 + 3) % 3;
            oldEncPos = newPos;
            showMenu();
        }
        if (pressed) {
            memset(board, 0, sizeof(board));
            gameState = PLAYING;
            // If Single Player RED, computer (1/Yellow) starts
            if (menuMode == 1) {
                currentPlayer = 1;
                renderBoard(); FastLED.show();
                performAiMove();
                currentPlayer = 2; // Set back to player
            } else {
                currentPlayer = 1; // Human starts
            }
            delay(300);
        }
    } 
    else if (gameState == PLAYING) {
        if (newPos != oldEncPos) {
            activeCol = (newPos % 7 + 7) % 7;
            oldEncPos = newPos;
        }
        renderBoard();
        leds[getIdx(activeCol, 0)] = (currentPlayer == 1) ? CRGB::Yellow : CRGB::Red;
        FastLED.show();

        if (pressed) {
            int row = getFirstEmptyRow(activeCol);
            if (row != -1) {
                board[activeCol][row] = currentPlayer;
                renderBoard(); FastLED.show(); 
                
                // 1. Check if the move just made ended the game
                if (scanBoard(currentPlayer)) {
                    gameState = FINISHED_WIN;
                } else if (isBoardFull()) {
                    gameState = FINISHED_DRAW;
                } else {
                    // 2. Handle Turn Switching
                    if (menuMode < 2) { // Single Player
                        int8_t aiP = (menuMode == 0) ? 2 : 1;
                        performAiMove();
                        if (scanBoard(aiP)) {
                            currentPlayer = aiP; // For the flashing color
                            gameState = FINISHED_WIN;
                        } else if (isBoardFull()) {
                            gameState = FINISHED_DRAW;
                        }
                    } else { // PvP
                        currentPlayer = (currentPlayer == 1) ? 2 : 1;
                    }
                }
            } else {
                for(int i=0; i<3; i++) {
                    leds[getIdx(activeCol, 0)] = CRGB::Black; FastLED.show(); delay(80);
                    leds[getIdx(activeCol, 0)] = (currentPlayer == 1) ? CRGB::Yellow : CRGB::Red; FastLED.show(); delay(80);
                }
            }
            delay(300);
        }
    } 
    else { 
        static unsigned long lastFlash = 0;
        static bool toggle = true;
        if (millis() - lastFlash > 300) {
            lastFlash = millis(); toggle = !toggle;
            renderBoard(); 
            for (int i = 0; i < NUM_LEDS; i++) {
                if (gameState == FINISHED_WIN) {
                    if (winMask[i]) leds[i] = toggle ? (currentPlayer == 1 ? CRGB::Yellow : CRGB::Red) : CRGB::Black;
                    else if (leds[i] && leds[i] != CRGB::Blue) leds[i].nscale8(40);
                } else { 
                    if (leds[i] && leds[i] != CRGB::Blue) leds[i] = toggle ? leds[i] : CRGB::Black;
                }
            }
            FastLED.show();
        }
        if (pressed) { gameState = MENU; showMenu(); delay(300); }
    }
}