#include <Arduino.h>
#include <FastLED.h>
#include <Encoder.h>
#include <WiFi.h>
#include <WebServer.h>
#include <Preferences.h>

#define NUM_LEDS 64
const int COLS = 7;
const int ROWS = 6;

CRGB leds[NUM_LEDS];
Encoder myEnc(ENC_A, ENC_B);
WebServer server(80);
Preferences prefs;

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

int8_t menuMode = 0;
int8_t currentPlayer = 1;
int8_t activeCol = 3;
long oldEncPos = -999;
uint32_t lastActivityTime = 0;
uint32_t demoResetTimer = 0;
bool isDemoOver = false;

// Web-Configurable Parameters (Stored in Flash)
uint8_t current_look_ahead;
uint8_t current_brightness;
uint32_t current_idle_timeout_ms;

// Thinking Animation Helpers
uint8_t aiBrightness = 0;
bool aiFadeUp = true;

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

void updateThinkingLED(int8_t p)
{
    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;
    }
    CRGB compColor = (p == 1) ? CRGB::Yellow : CRGB::Red;
    leds[getIdx(7, 0)] = compColor.nscale8(aiBrightness);
    FastLED.show();
}

void drawStaticUI()
{
    FastLED.clear();
    CRGB borderColor = CRGB::Blue;
    if (gameState == DEMO || (gameState >= 2 && isDemoOver))
    {
        uint8_t glow = beat8(15);
        borderColor = blend(CRGB::Blue, CRGB::White, glow / 4);
    }
    for (int x = 0; x < 7; x++)
        leds[getIdx(x, 1)] = borderColor;
    for (int y = 1; y < 8; y++)
        leds[getIdx(7, y)] = borderColor;
}

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;
}

bool scanBoard(int8_t p)
{
    memset(winMask, 0, sizeof(winMask));
    bool found = false;
    auto check = [&](int c, int r, int dc, int dr)
    {
        if (board[c][r] == p && board[c + dc][r + dr] == p && board[c + 2 * dc][r + 2 * dr] == p && board[c + 3 * dc][r + 3 * dr] == p)
        {
            for (int i = 0; i < 4; i++)
                winMask[getIdx(c + i * dc, 7 - (r + i * dr))] = true;
            return true;
        }
        return false;
    };
    for (int r = 0; r < 6; r++)
        for (int c = 0; c < 4; c++)
            if (check(c, r, 1, 0))
                found = true;
    for (int r = 0; r < 3; r++)
        for (int c = 0; c < 7; c++)
            if (check(c, r, 0, 1))
                found = true;
    for (int r = 0; r < 3; r++)
        for (int c = 0; c < 4; c++)
            if (check(c, r, 1, 1))
                found = true;
    for (int r = 3; r < 6; r++)
        for (int c = 0; c < 4; c++)
            if (check(c, r, 1, -1))
                found = true;
    return found;
}

// --- AI Engine ---
int minimax(int depth, int alpha, int beta, bool isMax, int8_t aiP, int8_t huP)
{
    if (depth >= current_look_ahead - 1)
        updateThinkingLED(aiP);
    else
        yield();

    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};
    int best = isMax ? -2000 : 2000;
    for (int c : order)
    {
        int r = getFirstEmptyRow(c);
        if (r != -1)
        {
            board[c][r] = isMax ? aiP : huP;
            int val = minimax(depth - 1, alpha, beta, !isMax, aiP, huP);
            board[c][r] = 0;
            if (isMax)
            {
                if (val > best)
                    best = val;
                alpha = max(alpha, best);
            }
            else
            {
                if (val < best)
                    best = val;
                beta = min(beta, best);
            }
            if (beta <= alpha)
                break;
        }
    }
    return best;
}

void performAiMove(int8_t aiP)
{
    int8_t huP = (aiP == 1) ? 2 : 1;
    aiBrightness = 0;
    aiFadeUp = true;
    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 (current_look_ahead >= 2 && scanBoard(huP))
            {
                board[c][r] = aiP;
                leds[getIdx(7, 0)] = CRGB::Black;
                return;
            }
            board[c][r] = 0;
        }
    }
    int bestScore = -30000;
    int bestCol = 3;
    for (int c : {3, 2, 4, 1, 5, 0, 6})
    {
        int r = getFirstEmptyRow(c);
        if (r != -1)
        {
            board[c][r] = aiP;
            int score = minimax(current_look_ahead, -30000, 30000, false, aiP, huP);
            board[c][r] = 0;
            if (score > bestScore)
            {
                bestScore = score;
                bestCol = c;
            }
        }
    }
    board[bestCol][getFirstEmptyRow(bestCol)] = aiP;
    leds[getIdx(7, 0)] = CRGB::Black;
}

// --- Menu UI with Restored Serifs ---
void showMenu()
{
    isDemoOver = false;
    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;
    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(3, 3)] = CRGB::Yellow;
        leds[getIdx(1, 6)] = CRGB::Yellow;
        leds[getIdx(3, 6)] = CRGB::Yellow;
        leds[getIdx(3, 3)] = CRGB::Red;
        leds[getIdx(5, 3)] = CRGB::Red;
        leds[getIdx(3, 6)] = CRGB::Red;
        leds[getIdx(5, 6)] = CRGB::Red;
    }
    FastLED.show();
}

// --- Web Portal ---
void handleRoot()
{
    String html = "<html><head><meta name='viewport' content='width=device-width, initial-scale=1'>";
    html += "<style>body{font-family:sans-serif;background:#121212;color:white;text-align:center;} .card{background:#222;padding:25px;border-radius:15px;display:inline-block;margin-top:20px;} input{width:100%;padding:10px;margin:10px 0;border-radius:5px;border:none;}</style></head><body>";
    html += "<h1>Connect 4 Admin</h1><div class='card'><form action='/save' method='POST'>";
    html += "AI Ply (1-10):<input type='number' name='ply' value='" + String(current_look_ahead) + "'>";
    html += "Brightness (5-255):<input type='number' name='br' value='" + String(current_brightness) + "'>";
    html += "Idle Timeout (Sec):<input type='number' name='idle' value='" + String(current_idle_timeout_ms / 1000) + "'>";
    html += "<input type='submit' value='Save Settings' style='background:#28a745;color:white;font-weight:bold;'></form></div></body></html>";
    server.send(200, "text/html", html);
}

void handleSave()
{
    if (server.hasArg("ply"))
    {
        current_look_ahead = server.arg("ply").toInt();
        prefs.putUChar("ply", current_look_ahead);
    }
    if (server.hasArg("br"))
    {
        current_brightness = server.arg("br").toInt();
        FastLED.setBrightness(current_brightness);
        prefs.putUChar("br", current_brightness);
    }
    if (server.hasArg("idle"))
    {
        uint32_t s = server.arg("idle").toInt();
        current_idle_timeout_ms = s * 1000;
        prefs.putUInt("idle", s);
    }
    server.sendHeader("Location", "/");
    server.send(303);
}

void setup()
{
    Serial.begin(115200);
    prefs.begin("c4-game", false);
    current_look_ahead = prefs.getUChar("ply", 8);
    current_brightness = prefs.getUChar("br", 25);
    current_idle_timeout_ms = prefs.getUInt("idle", 60) * 1000;

    FastLED.addLeds<WS2812B, LED_PIN, GRB>(leds, NUM_LEDS);
    FastLED.setBrightness(current_brightness);
    pinMode(ENC_SW, INPUT_PULLUP);

    WiFi.disconnect(true); // Clear old settings
    WiFi.mode(WIFI_AP);    // Force Access Point mode
    delay(100);            // Give the radio a moment to reset

    // SSID, Password (MUST be 8+ chars), Channel, Hidden (0=No), Max Clients
    if (WiFi.softAP("Connect4-Config", WIFI_PASSWORD, 1, 0, 4))
    {
        Serial.println("WPA2 AP Started Successfully");
    }
    else
    {
        Serial.println("AP Failed - Check if WIFI_PASSWORD is at least 8 characters!");
    }

    server.on("/", handleRoot);
    server.on("/save", HTTP_POST, handleSave);
    server.begin();

    lastActivityTime = millis();
    showMenu();
}

void loop()
{
    server.handleClient();
    long newPos = myEnc.read() / SENSITIVITY;
    bool pressed = (digitalRead(ENC_SW) == LOW);

    // Escape Demo / Interrupt
    if (newPos != oldEncPos || (pressed && (millis() - lastActivityTime > 500)))
    {
        if (gameState == DEMO || isDemoOver)
        {
            for (int i = 0; i < 10; i++)
            {
                fadeToBlackBy(leds, NUM_LEDS, 32);
                FastLED.show();
                delay(30);
            }
            delay(2000);
            gameState = MENU;
            memset(board, 0, sizeof(board));
            showMenu();
            lastActivityTime = millis();
            oldEncPos = newPos;
            return;
        }
        lastActivityTime = millis();
    }

    if (gameState == MENU)
    {
        if (newPos != oldEncPos)
        {
            menuMode = (newPos % 3 + 3) % 3;
            oldEncPos = newPos;
            showMenu();
        }
        if (pressed)
        {
            memset(board, 0, sizeof(board));
            gameState = PLAYING;
            if (menuMode == 1)
            {
                performAiMove(1);
                currentPlayer = 2;
            }
            else
            {
                currentPlayer = 1;
            }
            delay(300);
        }
        if (millis() - lastActivityTime > current_idle_timeout_ms)
        {
            gameState = DEMO;
            memset(board, 0, sizeof(board));
            currentPlayer = 1;
        }
    }
    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();
                if (scanBoard(currentPlayer))
                    gameState = FINISHED_WIN;
                else if (isBoardFull())
                    gameState = FINISHED_DRAW;
                else
                {
                    if (menuMode < 2)
                    {
                        int8_t aiP = (menuMode == 0) ? 2 : 1;
                        performAiMove(aiP);
                        renderBoard();
                        FastLED.show();
                        if (scanBoard(aiP))
                        {
                            currentPlayer = aiP;
                            gameState = FINISHED_WIN;
                        }
                        else if (isBoardFull())
                            gameState = FINISHED_DRAW;
                    }
                    else
                    {
                        currentPlayer = (currentPlayer == 1) ? 2 : 1;
                    }
                }
                delay(300);
            }
        }
    }
    else if (gameState == DEMO)
    {
        // No idle timeout check here to prevent premature restarts
        renderBoard();
        FastLED.show();
        delay(600);
        performAiMove(currentPlayer);
        if (scanBoard(currentPlayer))
        {
            gameState = FINISHED_WIN;
            isDemoOver = true;
            demoResetTimer = millis();
        }
        else if (isBoardFull())
        {
            gameState = FINISHED_DRAW;
            isDemoOver = true;
            demoResetTimer = millis();
        }
        else
        {
            currentPlayer = (currentPlayer == 1) ? 2 : 1;
        }
    }
    else
    {
        // Monitor for Idle in Win screen to return to Demo
        if (!isDemoOver && (millis() - lastActivityTime > current_idle_timeout_ms))
        {
            memset(board, 0, sizeof(board));
            gameState = DEMO;
            currentPlayer = 1;
            return;
        }

        static uint32_t lastFlash = 0;
        static bool toggle = true;
        if (millis() - lastFlash > 300)
        {
            lastFlash = millis();
            toggle = !toggle;
            renderBoard();
            for (int i = 0; i < NUM_LEDS; i++)
            {
                if (leds[i] == CRGB::Blue)
                    continue;
                if (gameState == FINISHED_WIN)
                {
                    if (winMask[i])
                        leds[i] = toggle ? (currentPlayer == 1 ? CRGB::Yellow : CRGB::Red) : CRGB::Black;
                    else
                        leds[i].nscale8(40);
                }
                else if (gameState == FINISHED_DRAW)
                {
                    if (!toggle)
                        leds[i] = CRGB::Black;
                }
            }
            FastLED.show();
        }
        // Restart Demo loop if it was a demo game
        if (isDemoOver && (millis() - demoResetTimer > 30000))
        {
            memset(board, 0, sizeof(board));
            gameState = DEMO;
            isDemoOver = false;
        }
        if (pressed)
        {
            gameState = MENU;
            showMenu();
            delay(300);
        }
    }
}