const menuPanel = document.getElementById('menu');
const gamePanel = document.getElementById('game');
const statusDiv = document.getElementById('connection-status');
const createBtn = document.getElementById('create-btn');
const joinBtn = document.getElementById('join-btn');
const joinInput = document.getElementById('join-code');
const displayCode = document.getElementById('display-code');
const turnStatus = document.getElementById('turn-status');
const leaveBtn = document.getElementById('leave-btn');

// ... (other vars)

// ...

leaveBtn.addEventListener('click', () => {
    if (confirm('Are you sure you want to leave the lobby?')) {
        send('LEAVE');
        resetGame();
    }
});

surrenderBtn.addEventListener('click', () => {
    if (confirm('Are you sure you want to surrender?')) {
        send('SURRENDER');
    }
});

// ...

function handleMessage(msg) {
    console.log("Received:", msg);
    const parts = msg.split(' ');
    const cmd = parts[0];

    switch (cmd) {
        case 'GAME_CREATED':
            currentCode = parts[1];
            mySymbol = 'X';
            enterGame();
            updateStatus("Waiting for opponent...");
            break;
        case 'JOIN_SUCCESS':
            currentCode = joinInput.value;
            mySymbol = 'O';
            enterGame();
            updateStatus("Connected! Game starting soon...");
            break;
        case 'START':
            updateStatus("Game Started! X goes first.");
            break;
        case 'RESTART':
            updateStatus("Game Restarted! X goes first.");
            cells.forEach(c => {
                c.textContent = '';
                c.className = 'cell';
            });
            break;
        case 'BOARD':
            const boardStr = msg.substring(6);
            updateBoard(boardStr);
            break;
        case 'TURN':
            const turn = parts[1];
            if (turn === mySymbol) {
                updateStatus(`Your Turn (${mySymbol})`, true);
            } else {
                updateStatus(`Opponent's Turn (${turn})`, false);
            }
            break;
        case 'WIN':
            const winner = msg.substring(4);
            setTimeout(() => {
                if (confirm(`Game Over! Winner: ${winner}\nDo you want to play again?`)) {
                    send('RESTART');
                }
            }, 100);
            break;
        case 'DRAW':
            setTimeout(() => {
                if (confirm("Game Over! It's a draw!\nDo you want to play again?")) {
                    send('RESTART');
                }
            }, 100);
            break;
        case 'OPPONENT_LEFT':
            alert("Opponent has left the lobby.");
            resetGame();
            break;
        case 'ERROR:':
            alert(msg);
            break;
    }
}

function enterGame() {
    menuPanel.classList.add('hidden');
    gamePanel.classList.remove('hidden');
    displayCode.textContent = currentCode;
    // Clear board
    cells.forEach(c => {
        c.textContent = '';
        c.className = 'cell';
    });
}

function resetGame() {
    menuPanel.classList.remove('hidden');
    gamePanel.classList.add('hidden');
    currentCode = '';
    mySymbol = '';
    joinInput.value = '';
    statusDiv.textContent = 'Connected to Server';
    statusDiv.style.color = '#10b981';
}

function updateBoard(boardStr) {
    cells.forEach((cell, i) => {
        if (i < boardStr.length) {
            const char = boardStr.charAt(i);
            cell.textContent = char === ' ' ? '' : char;
            cell.className = 'cell';
            if (char === 'X') cell.classList.add('x');
            if (char === 'O') cell.classList.add('o');
        }
    });
}

function updateStatus(text, isAction = false) {
    turnStatus.textContent = text;
    if (isAction) {
        turnStatus.style.color = '#10b981'; // Green for 'your turn'
    } else {
        turnStatus.style.color = '#a1a1aa';
    }
}

connect();