/* ============================================================
 *  Connect Four — Browser Edition
 *  A single-file game: AI (minimax + alpha-beta + heuristic), demo mode,
 *  game log (localStorage), blunder mode, idle timeout.
 *
 *  Include this script in an HTML page that has:
 *    <canvas id="gameCanvas"></canvas>
 *
 *  Works in Firefox, Chrome, Edge, Safari, and Brave.
 * ============================================================ */

// --- Configurable Parameters --------------------------------
const COLS             = 7;        // board columns
const ROWS             = 6;        // board rows
const LOOK_AHEAD       = 8;        // AI search depth (plies)
const BLUNDER_ENABLED  = false;    // allow random AI mistakes
const BLUNDER_CHANCE   = 20;       // percent chance of blunder (0-100)
const DEMO_RESET_PAUSE = 5;       // seconds before auto-demo after game end
const IDLE_TIMEOUT     = 60;       // seconds of inactivity before demo starts
const MAX_GAME_LOG     = 100;      // max stored game entries (localStorage)

// --- Visual Parameters --------------------------------------
const CELL_SIZE        = 70;       // pixel size of each board cell
const DISC_RADIUS      = 28;       // radius of a disc
const BOARD_PAD_TOP    = 100;      // space above the board (cursor + col numbers)
const BOARD_PAD_X      = 40;       // horizontal padding
const BOARD_PAD_BOTTOM = 40;       // space below the board
const ANIM_DROP_SPEED  = 1200;     // pixels per second for drop animation
const FONT_FAMILY      = "system-ui, -apple-system, sans-serif";

// --- Colors -------------------------------------------------
const COLOR_BG         = "#1a1a2e";
const COLOR_BOARD      = "#16213e";
const COLOR_GRID_LINE  = "#0f3460";
const COLOR_EMPTY      = "#0a1628";
const COLOR_P1         = "#ffd700"; // Yellow (player 1)
const COLOR_P2         = "#e63946"; // Red (player 2)
const COLOR_P1_DIM     = "#8b7500";
const COLOR_P2_DIM     = "#7a1f26";
const COLOR_HIGHLIGHT  = "#ffffff";
const COLOR_TEXT        = "#e0e0e0";
const COLOR_TEXT_DIM    = "#666680";
const COLOR_MENU_BG    = "#1a1a2e";
const COLOR_MENU_SEL   = "#0f3460";
// ------------------------------------------------------------

const COL_ORDER = [3, 2, 4, 1, 5, 0, 6];

const State = Object.freeze({
    MENU:          0,
    PLAYING:       1,
    AI_TURN:       2,
    FINISHED_WIN:  3,
    FINISHED_DRAW: 4,
    DEMO:          5,
});

// --- Canvas setup -------------------------------------------
const canvas = document.getElementById("gameCanvas");
const ctx = canvas.getContext("2d");

const BOARD_W = COLS * CELL_SIZE;
const BOARD_H = ROWS * CELL_SIZE;
const CANVAS_W = BOARD_W + BOARD_PAD_X * 2;
const CANVAS_H = BOARD_PAD_TOP + BOARD_H + BOARD_PAD_BOTTOM;

canvas.width = CANVAS_W;
canvas.height = CANVAS_H;
canvas.style.display = "block";
canvas.style.margin = "0 auto";
canvas.tabIndex = 0;
canvas.focus();

// --- Game state ---------------------------------------------
let board = makeBoard();
let gameState = State.MENU;
let menuMode = 0;
let currentPlayer = 1;
let activeCol = 3;
let winnerPlayer = 0;
let winPositions = [];
let currentMoves = "";
let gameMenuMode = 0;
let gameLevel = LOOK_AHEAD;
let games = loadGameLog();
let demoPly = [4, 4];
let lastActivity = performance.now() / 1000;
let demoResetTimer = 0;
let flashToggle = true;
let lastFlash = 0;
let hoverCol = -1;

// Drop animation state
let dropping = false;
let dropCol = -1;
let dropPlayer = 0;
let dropTargetRow = -1;
let dropY = 0;
let dropTargetY = 0;

// --- Board helpers ------------------------------------------
function makeBoard() {
    const b = [];
    for (let c = 0; c < COLS; c++) {
        b[c] = new Array(ROWS).fill(0);
    }
    return b;
}

function resetGame() {
    board = makeBoard();
    winnerPlayer = 0;
    winPositions = [];
    currentMoves = "";
}

function getFirstEmptyRow(b, col) {
    for (let r = 0; r < ROWS; r++) {
        if (b[col][r] === 0) return r;
    }
    return -1;
}

function isBoardFull(b) {
    for (let c = 0; c < COLS; c++) {
        if (b[c][ROWS - 1] === 0) return false;
    }
    return true;
}

function scanBoard(b) {
    function check(c, r, dc, dr) {
        const p = b[c][r];
        if (p === 0) return [0, []];
        const pos = [];
        for (let i = 0; i < 4; i++) {
            const cc = c + i * dc;
            const rr = r + i * dr;
            if (cc < 0 || cc >= COLS || rr < 0 || rr >= ROWS) return [0, []];
            if (b[cc][rr] !== p) return [0, []];
            pos.push([cc, rr]);
        }
        return [p, pos];
    }

    for (let r = 0; r < ROWS; r++)
        for (let c = 0; c <= COLS - 4; c++) {
            const [w, pos] = check(c, r, 1, 0);
            if (w) return [w, pos];
        }
    for (let r = 0; r <= ROWS - 4; r++)
        for (let c = 0; c < COLS; c++) {
            const [w, pos] = check(c, r, 0, 1);
            if (w) return [w, pos];
        }
    for (let r = 0; r <= ROWS - 4; r++)
        for (let c = 0; c <= COLS - 4; c++) {
            const [w, pos] = check(c, r, 1, 1);
            if (w) return [w, pos];
        }
    for (let r = 3; r < ROWS; r++)
        for (let c = 0; c <= COLS - 4; c++) {
            const [w, pos] = check(c, r, 1, -1);
            if (w) return [w, pos];
        }
    return [0, []];
}

function evaluateBoard(b, aiP, huP) {
    let score = 0;
    let aiThreats = 0, huThreats = 0;

    // Center column bonus
    for (let r = 0; r < ROWS; r++) {
        if (b[3][r] === aiP) score += 3;
        else if (b[3][r] === huP) score -= 3;
    }

    // Score a window of 4 cells by piece counts
    function scoreWindow(c, r, dc, dr) {
        let ai = 0, hu = 0, emptyC = -1, emptyR = -1;
        for (let i = 0; i < 4; i++) {
            const cc = c + i * dc;
            const rr = r + i * dr;
            const v = b[cc][rr];
            if (v === aiP) ai++;
            else if (v === huP) hu++;
            else { emptyC = cc; emptyR = rr; }
        }
        if (ai > 0 && hu > 0) return 0;
        if (ai === 3) {
            aiThreats++;
            const playable = emptyR === 0 || b[emptyC][emptyR - 1] !== 0;
            return playable ? 100 : 40;
        }
        if (ai === 2) return 5;
        if (hu === 3) {
            huThreats++;
            const playable = emptyR === 0 || b[emptyC][emptyR - 1] !== 0;
            return playable ? -100 : -40;
        }
        if (hu === 2) return -5;
        return 0;
    }

    // Horizontal
    for (let r = 0; r < ROWS; r++)
        for (let c = 0; c <= COLS - 4; c++)
            score += scoreWindow(c, r, 1, 0);
    // Vertical
    for (let r = 0; r <= ROWS - 4; r++)
        for (let c = 0; c < COLS; c++)
            score += scoreWindow(c, r, 0, 1);
    // Diagonal up-right
    for (let r = 0; r <= ROWS - 4; r++)
        for (let c = 0; c <= COLS - 4; c++)
            score += scoreWindow(c, r, 1, 1);
    // Diagonal down-right
    for (let r = 3; r < ROWS; r++)
        for (let c = 0; c <= COLS - 4; c++)
            score += scoreWindow(c, r, 1, -1);

    // Fork bonus: multiple threats are disproportionately dangerous
    if (aiThreats >= 2) score += 200;
    if (huThreats >= 2) score -= 200;

    return score;
}

// --- AI -----------------------------------------------------
function minimax(b, depth, alpha, beta, isMax, aiP, huP) {
    const [winner] = scanBoard(b);
    if (winner === aiP) return 1000 + depth;
    if (winner === huP) return -1000 - depth;
    if (depth === 0 || isBoardFull(b)) return evaluateBoard(b, aiP, huP);

    let best = isMax ? -10000 : 10000;
    for (const c of COL_ORDER) {
        const r = getFirstEmptyRow(b, c);
        if (r === -1) continue;
        b[c][r] = isMax ? aiP : huP;
        const score = minimax(b, depth - 1, alpha, beta, !isMax, aiP, huP);
        b[c][r] = 0;
        if (isMax) {
            if (score > best) best = score;
            if (best > alpha) alpha = best;
        } else {
            if (score < best) best = score;
            if (best < beta) beta = best;
        }
        if (beta <= alpha) break;
    }
    return best;
}

function performAiMove(b, aiP, lookAhead, isDemo = false, dPly = 4) {
    const huP = aiP === 1 ? 2 : 1;
    const ply = isDemo ? dPly : lookAhead;

    // Phase 1a: check ALL columns for instant AI win
    for (let c = 0; c < COLS; c++) {
        const r = getFirstEmptyRow(b, c);
        if (r === -1) continue;
        b[c][r] = aiP;
        if (scanBoard(b)[0] === aiP) { b[c][r] = 0; return c; }
        b[c][r] = 0;
    }

    // Phase 1b: check ALL columns for opponent block
    for (let c = 0; c < COLS; c++) {
        const r = getFirstEmptyRow(b, c);
        if (r === -1) continue;
        b[c][r] = huP;
        if (scanBoard(b)[0] === huP) { b[c][r] = 0; return c; }
        b[c][r] = 0;
    }

    // Phase 2: blunder
    if (!isDemo && BLUNDER_ENABLED && Math.random() * 100 < BLUNDER_CHANCE) {
        const valid = [];
        for (let c = 0; c < COLS; c++) if (getFirstEmptyRow(b, c) !== -1) valid.push(c);
        return valid[Math.floor(Math.random() * valid.length)];
    }

    // Phase 3: minimax
    let bestScore = -30000;
    let bestCol = 3;
    for (const c of COL_ORDER) {
        const r = getFirstEmptyRow(b, c);
        if (r === -1) continue;
        b[c][r] = aiP;
        const score = minimax(b, ply, -30000, 30000, false, aiP, huP);
        b[c][r] = 0;
        if (score > bestScore) {
            bestScore = score;
            bestCol = c;
        }
    }
    return bestCol;
}

function randomizeDemoPlies() {
    const strong = 4 + Math.floor(Math.random() * 2);
    const weak = 2 + Math.floor(Math.random() * 2);
    return Math.random() < 0.5 ? [strong, weak] : [weak, strong];
}

// --- Game log (localStorage) --------------------------------
function loadGameLog() {
    try {
        const raw = localStorage.getItem("connectFourLog");
        if (!raw) return [];
        return JSON.parse(raw).slice(-MAX_GAME_LOG);
    } catch { return []; }
}

function saveGameLog(g) {
    try {
        localStorage.setItem("connectFourLog", JSON.stringify(g.slice(-MAX_GAME_LOG)));
    } catch { /* storage full or unavailable */ }
}

function logGame(g, gMenuMode, level, winner, moves) {
    const type = gMenuMode === 0 ? "Y" : gMenuMode === 1 ? "R" : "2";
    const winChar = winner === 1 ? "Y" : winner === 2 ? "R" : "D";
    g.push({ type, level: String(level), winner: winChar, moves });
    g = g.slice(-MAX_GAME_LOG);
    saveGameLog(g);
    return g;
}

// --- Check game end -----------------------------------------
function checkGameEnd() {
    const [w, pos] = scanBoard(board);
    winnerPlayer = w;
    winPositions = pos;
    const won = w !== 0;
    const draw = !won && isBoardFull(board);
    if (!won && !draw) return false;

    if (gameState !== State.DEMO) {
        games = logGame(games, gameMenuMode, gameLevel, won ? w : 0, currentMoves);
        console.log(`Game: ${currentMoves} → ${won ? playerName(w) + " wins" : "Draw"}`);
    }
    gameState = won ? State.FINISHED_WIN : State.FINISHED_DRAW;
    demoResetTimer = performance.now() / 1000;
    lastActivity = performance.now() / 1000;
    return true;
}

// --- Drawing ------------------------------------------------
function playerColor(p) { return p === 1 ? COLOR_P1 : COLOR_P2; }
function playerColorDim(p) { return p === 1 ? COLOR_P1_DIM : COLOR_P2_DIM; }
function playerName(p) { return p === 1 ? "Yellow" : "Red"; }

function cellX(c) { return BOARD_PAD_X + c * CELL_SIZE + CELL_SIZE / 2; }
function cellY(r) { return BOARD_PAD_TOP + (ROWS - 1 - r) * CELL_SIZE + CELL_SIZE / 2; }

function isWinPos(c, r) {
    for (const [wc, wr] of winPositions) {
        if (wc === c && wr === r) return true;
    }
    return false;
}

function drawDisc(x, y, radius, color) {
    ctx.beginPath();
    ctx.arc(x, y, radius, 0, Math.PI * 2);
    ctx.fillStyle = color;
    ctx.fill();
}

function drawBoard() {
    // Board background
    ctx.fillStyle = COLOR_BOARD;
    const bx = BOARD_PAD_X;
    const by = BOARD_PAD_TOP;
    ctx.beginPath();
    ctx.roundRect(bx - 5, by - 5, BOARD_W + 10, BOARD_H + 10, 12);
    ctx.fill();

    // Grid lines
    ctx.strokeStyle = COLOR_GRID_LINE;
    ctx.lineWidth = 1;
    for (let c = 1; c < COLS; c++) {
        const x = BOARD_PAD_X + c * CELL_SIZE;
        ctx.beginPath();
        ctx.moveTo(x, BOARD_PAD_TOP);
        ctx.lineTo(x, BOARD_PAD_TOP + BOARD_H);
        ctx.stroke();
    }
    for (let r = 1; r < ROWS; r++) {
        const y = BOARD_PAD_TOP + r * CELL_SIZE;
        ctx.beginPath();
        ctx.moveTo(BOARD_PAD_X, y);
        ctx.lineTo(BOARD_PAD_X + BOARD_W, y);
        ctx.stroke();
    }

    // Cells
    for (let c = 0; c < COLS; c++) {
        for (let r = 0; r < ROWS; r++) {
            const x = cellX(c);
            const y = cellY(r);
            const val = board[c][r];

            // Skip drawing in cell if we're animating a drop into it
            if (dropping && c === dropCol && r === dropTargetRow) continue;

            if (val === 0) {
                drawDisc(x, y, DISC_RADIUS, COLOR_EMPTY);
            } else {
                const isWin = isWinPos(c, r);
                if (gameState === State.FINISHED_WIN) {
                    if (isWin && flashToggle) {
                        drawDisc(x, y, DISC_RADIUS, COLOR_EMPTY);
                    } else if (!isWin) {
                        drawDisc(x, y, DISC_RADIUS, playerColorDim(val));
                    } else {
                        drawDisc(x, y, DISC_RADIUS, playerColor(val));
                    }
                } else if (gameState === State.FINISHED_DRAW && flashToggle) {
                    drawDisc(x, y, DISC_RADIUS, COLOR_EMPTY);
                } else {
                    drawDisc(x, y, DISC_RADIUS, playerColor(val));
                }
            }
        }
    }

    // Drop animation disc
    if (dropping) {
        drawDisc(cellX(dropCol), dropY, DISC_RADIUS, playerColor(dropPlayer));
    }
}

function drawCursor() {
    if (gameState === State.PLAYING && !dropping) {
        const x = cellX(activeCol);
        const y = BOARD_PAD_TOP - 45;
        drawDisc(x, y, DISC_RADIUS * 0.8, playerColor(currentPlayer));
    }
    if (gameState === State.PLAYING && hoverCol >= 0 && hoverCol !== activeCol && !dropping) {
        const x = cellX(hoverCol);
        const y = BOARD_PAD_TOP - 45;
        drawDisc(x, y, DISC_RADIUS * 0.5, playerColorDim(currentPlayer));
    }
}

function drawColNumbers() {
    ctx.font = `14px ${FONT_FAMILY}`;
    ctx.textAlign = "center";
    ctx.textBaseline = "middle";
    for (let c = 0; c < COLS; c++) {
        ctx.fillStyle = (c === activeCol && gameState === State.PLAYING) ? COLOR_TEXT : COLOR_TEXT_DIM;
        ctx.fillText(String(c + 1), cellX(c), BOARD_PAD_TOP - 12);
    }
}

function drawStatus() {
    const y = BOARD_PAD_TOP + BOARD_H + 25;
    ctx.font = `bold 18px ${FONT_FAMILY}`;
    ctx.textAlign = "center";
    ctx.textBaseline = "middle";

    if (gameState === State.PLAYING) {
        ctx.fillStyle = playerColor(currentPlayer);
        const label = gameMenuMode === 2 ? `${playerName(currentPlayer)}'s turn`
            : currentPlayer === (gameMenuMode === 0 ? 1 : 2) ? "Your turn" : "AI thinking...";
        ctx.fillText(label, CANVAS_W / 2, y);
    } else if (gameState === State.AI_TURN) {
        const aiP = gameMenuMode === 0 ? 2 : 1;
        ctx.fillStyle = playerColor(aiP);
        ctx.fillText("AI thinking...", CANVAS_W / 2, y);
    } else if (gameState === State.FINISHED_WIN) {
        ctx.fillStyle = playerColor(winnerPlayer);
        ctx.fillText(`${playerName(winnerPlayer)} wins!`, CANVAS_W / 2, y);
        ctx.font = `14px ${FONT_FAMILY}`;
        ctx.fillStyle = COLOR_TEXT_DIM;
        ctx.fillText("Click or press any key for menu", CANVAS_W / 2, y + 24);
    } else if (gameState === State.FINISHED_DRAW) {
        ctx.fillStyle = COLOR_TEXT;
        ctx.fillText("Draw!", CANVAS_W / 2, y);
        ctx.font = `14px ${FONT_FAMILY}`;
        ctx.fillStyle = COLOR_TEXT_DIM;
        ctx.fillText("Click or press any key for menu", CANVAS_W / 2, y + 24);
    } else if (gameState === State.DEMO) {
        ctx.fillStyle = COLOR_TEXT_DIM;
        ctx.font = `14px ${FONT_FAMILY}`;
        ctx.fillText("Demo mode - click or press any key for menu", CANVAS_W / 2, y);
    }
}

function drawMenu() {
    ctx.fillStyle = COLOR_MENU_BG;
    ctx.fillRect(0, 0, CANVAS_W, CANVAS_H);

    ctx.font = `bold 36px ${FONT_FAMILY}`;
    ctx.textAlign = "center";
    ctx.textBaseline = "middle";
    ctx.fillStyle = COLOR_P1;
    ctx.fillText("Connect", CANVAS_W / 2 - 60, 80);
    ctx.fillStyle = COLOR_P2;
    ctx.fillText("Four", CANVAS_W / 2 + 70, 80);

    const items = [
        { label: "1P Yellow (you start)", color: COLOR_P1 },
        { label: "1P Red (AI starts)",    color: COLOR_P2 },
        { label: "Multiplayer",           color: "#5dade2" },
    ];

    const startY = 160;
    const itemH = 60;
    const itemW = 340;

    for (let i = 0; i < items.length; i++) {
        const y = startY + i * itemH;
        const x = (CANVAS_W - itemW) / 2;
        const selected = i === menuMode;

        // Background
        ctx.fillStyle = selected ? COLOR_MENU_SEL : "transparent";
        ctx.beginPath();
        ctx.roundRect(x, y, itemW, 48, 8);
        ctx.fill();

        // Border for selected
        if (selected) {
            ctx.strokeStyle = items[i].color;
            ctx.lineWidth = 2;
            ctx.beginPath();
            ctx.roundRect(x, y, itemW, 48, 8);
            ctx.stroke();
        }

        // Arrow
        ctx.font = `bold 20px ${FONT_FAMILY}`;
        ctx.textAlign = "left";
        ctx.textBaseline = "middle";
        ctx.fillStyle = selected ? items[i].color : COLOR_TEXT_DIM;
        ctx.fillText(selected ? "\u25b6 " : "  ", x + 16, y + 24);

        // Label
        ctx.font = `${selected ? "bold " : ""}18px ${FONT_FAMILY}`;
        ctx.fillText(items[i].label, x + 50, y + 24);
    }

    // Instructions
    ctx.font = `14px ${FONT_FAMILY}`;
    ctx.textAlign = "center";
    ctx.fillStyle = COLOR_TEXT_DIM;
    ctx.fillText("Up/Down or hover to select, click or Enter to start", CANVAS_W / 2, startY + items.length * itemH + 20);
    ctx.fillText("During game: Arrow keys or click columns, 1-7 for direct drop", CANVAS_W / 2, startY + items.length * itemH + 44);
}

function render() {
    ctx.fillStyle = COLOR_BG;
    ctx.fillRect(0, 0, CANVAS_W, CANVAS_H);

    if (gameState === State.MENU) {
        drawMenu();
    } else {
        drawBoard();
        drawCursor();
        drawColNumbers();
        drawStatus();
    }
}

// --- Drop animation -----------------------------------------
function animateDrop(col, row, player) {
    return new Promise(resolve => {
        dropping = true;
        dropCol = col;
        dropPlayer = player;
        dropTargetRow = row;
        dropY = BOARD_PAD_TOP - 45;
        dropTargetY = cellY(row);

        function step(timestamp) {
            dropY += ANIM_DROP_SPEED * (1 / 60);
            if (dropY >= dropTargetY) {
                dropY = dropTargetY;
                dropping = false;
                resolve();
                return;
            }
            requestAnimationFrame(step);
        }
        requestAnimationFrame(step);
    });
}

// --- Input: column from mouse / touch -----------------------
function colFromEvent(e) {
    const rect = canvas.getBoundingClientRect();
    const scaleX = CANVAS_W / rect.width;
    const x = (e.clientX - rect.left) * scaleX;
    const col = Math.floor((x - BOARD_PAD_X) / CELL_SIZE);
    return (col >= 0 && col < COLS) ? col : -1;
}

function menuItemFromEvent(e) {
    const rect = canvas.getBoundingClientRect();
    const scaleY = CANVAS_H / rect.height;
    const scaleX = CANVAS_W / rect.width;
    const y = (e.clientY - rect.top) * scaleY;
    const x = (e.clientX - rect.left) * scaleX;
    const startY = 160;
    const itemH = 60;
    const itemW = 340;
    const mx = (CANVAS_W - itemW) / 2;
    for (let i = 0; i < 3; i++) {
        const iy = startY + i * itemH;
        if (x >= mx && x <= mx + itemW && y >= iy && y <= iy + 48) return i;
    }
    return -1;
}

// --- Place a disc (with animation) --------------------------
let busy = false; // prevents input during animation / AI

async function placeDisk(col, player) {
    const r = getFirstEmptyRow(board, col);
    if (r === -1) return false;
    currentMoves += String(col);
    await animateDrop(col, r, player);
    board[col][r] = player;
    return true;
}

// --- AI turn (async to not block UI) ------------------------
async function doAiTurn() {
    busy = true;
    const aiP = gameMenuMode === 0 ? 2 : 1;
    gameState = State.AI_TURN;

    // Yield a frame so "AI thinking" shows
    await new Promise(r => setTimeout(r, 50));

    const bestCol = performAiMove(board, aiP, LOOK_AHEAD);
    await placeDisk(bestCol, aiP);
    activeCol = bestCol;

    if (!checkGameEnd()) {
        gameState = State.PLAYING;
        currentPlayer = aiP === 1 ? 2 : 1;
    }
    lastActivity = performance.now() / 1000;
    busy = false;
}

// --- Demo turn ----------------------------------------------
let demoTimer = null;

function stopDemo() {
    if (demoTimer !== null) {
        clearTimeout(demoTimer);
        demoTimer = null;
    }
}

async function demoStep() {
    if (gameState !== State.DEMO) return;
    busy = true;
    const ply = demoPly[currentPlayer - 1];
    const bestCol = performAiMove(board, currentPlayer, LOOK_AHEAD, true, ply);
    await placeDisk(bestCol, currentPlayer);

    if (!checkGameEnd()) {
        currentPlayer = currentPlayer === 1 ? 2 : 1;
        demoTimer = setTimeout(demoStep, 400);
    }
    busy = false;
}

function startDemo() {
    resetGame();
    demoPly = randomizeDemoPlies();
    gameState = State.DEMO;
    currentPlayer = 1;
    lastActivity = performance.now() / 1000;
    demoTimer = setTimeout(demoStep, 400);
}

// --- Start game from menu -----------------------------------
function startGame(mode) {
    resetGame();
    gameMenuMode = mode;
    gameLevel = LOOK_AHEAD;
    currentPlayer = 1;
    activeCol = 3;
    hoverCol = -1;

    if (mode === 1) {
        gameState = State.PLAYING; // briefly, then AI
        doAiTurn();
    } else {
        gameState = State.PLAYING;
    }
    lastActivity = performance.now() / 1000;
}

function returnToMenu() {
    stopDemo();
    resetGame();
    gameState = State.MENU;
    menuMode = 0;
    lastActivity = performance.now() / 1000;
}

// --- Mouse events -------------------------------------------
canvas.addEventListener("mousemove", (e) => {
    if (gameState === State.MENU) {
        const mi = menuItemFromEvent(e);
        if (mi >= 0) menuMode = mi;
    } else if (gameState === State.PLAYING && !busy) {
        hoverCol = colFromEvent(e);
    }
});

canvas.addEventListener("click", async (e) => {
    if (busy) return;
    lastActivity = performance.now() / 1000;

    if (gameState === State.MENU) {
        const mi = menuItemFromEvent(e);
        if (mi >= 0) {
            menuMode = mi;
            startGame(mi);
        }
        return;
    }

    if (gameState === State.FINISHED_WIN || gameState === State.FINISHED_DRAW || gameState === State.DEMO) {
        returnToMenu();
        return;
    }

    if (gameState === State.PLAYING) {
        const col = colFromEvent(e);
        if (col < 0) return;
        const r = getFirstEmptyRow(board, col);
        if (r === -1) return;

        busy = true;
        activeCol = col;
        await placeDisk(col, currentPlayer);

        if (!checkGameEnd()) {
            if (gameMenuMode < 2) {
                await doAiTurn();
            } else {
                currentPlayer = currentPlayer === 1 ? 2 : 1;
            }
        }
        busy = false;
    }
});

// --- Touch support (mobile) ---------------------------------
canvas.addEventListener("touchend", (e) => {
    if (e.changedTouches.length > 0) {
        const touch = e.changedTouches[0];
        const click = new MouseEvent("click", {
            clientX: touch.clientX,
            clientY: touch.clientY,
        });
        canvas.dispatchEvent(click);
    }
    e.preventDefault();
}, { passive: false });

// --- Keyboard events ----------------------------------------
document.addEventListener("keydown", async (e) => {
    if (busy) return;
    lastActivity = performance.now() / 1000;

    if (e.key === "q" || e.key === "Q") {
        if (gameState !== State.MENU) {
            returnToMenu();
        }
        return;
    }

    if (gameState === State.MENU) {
        if (e.key === "ArrowUp") {
            menuMode = (menuMode - 1 + 3) % 3;
        } else if (e.key === "ArrowDown") {
            menuMode = (menuMode + 1) % 3;
        } else if (e.key === "Enter" || e.key === " ") {
            startGame(menuMode);
        }
        return;
    }

    if (gameState === State.FINISHED_WIN || gameState === State.FINISHED_DRAW || gameState === State.DEMO) {
        returnToMenu();
        return;
    }

    if (gameState === State.PLAYING) {
        if (e.key === "ArrowLeft") {
            activeCol = Math.max(0, activeCol - 1);
        } else if (e.key === "ArrowRight") {
            activeCol = Math.min(COLS - 1, activeCol + 1);
        } else if (e.key >= "1" && e.key <= "7") {
            const col = parseInt(e.key) - 1;
            const r = getFirstEmptyRow(board, col);
            if (r === -1) return;

            busy = true;
            activeCol = col;
            await placeDisk(col, currentPlayer);
            if (!checkGameEnd()) {
                if (gameMenuMode < 2) {
                    await doAiTurn();
                } else {
                    currentPlayer = currentPlayer === 1 ? 2 : 1;
                }
            }
            busy = false;
        } else if (e.key === "Enter" || e.key === " ") {
            const r = getFirstEmptyRow(board, activeCol);
            if (r === -1) return;

            busy = true;
            await placeDisk(activeCol, currentPlayer);
            if (!checkGameEnd()) {
                if (gameMenuMode < 2) {
                    await doAiTurn();
                } else {
                    currentPlayer = currentPlayer === 1 ? 2 : 1;
                }
            }
            busy = false;
        }
    }
});

// --- Main loop ----------------------------------------------
let lastTime = 0;

function gameLoop(timestamp) {
    const now = timestamp / 1000;

    // Flash toggle for win/draw
    if (gameState === State.FINISHED_WIN || gameState === State.FINISHED_DRAW) {
        if (now - lastFlash > 0.4) {
            lastFlash = now;
            flashToggle = !flashToggle;
        }

        // Auto-restart to demo
        if (now - demoResetTimer > DEMO_RESET_PAUSE) {
            startDemo();
        }
    }

    // Idle timeout -> demo
    if (gameState !== State.DEMO && gameState !== State.FINISHED_WIN && gameState !== State.FINISHED_DRAW) {
        if (now - lastActivity > IDLE_TIMEOUT) {
            startDemo();
        }
    }

    render();
    requestAnimationFrame(gameLoop);
}

requestAnimationFrame(gameLoop);