snake-ai-game/src/display.ts

A terminal-based snake game powered by AI.

import chalk from 'chalk';
import { Direction, GameState, ActionType } from './types.js';
import { clearScreen, formatTime } from './utils.js';
import { findShortestPathToFood, getValidMoves, calculateNewPosition, analyzeSpatialState, suggestOptimalMove } from './gameLogic.js';

const COLORS = {
  PATH: chalk.rgb(120, 255, 120),
  FOOD: chalk.rgb(255, 100, 100),
  DANGER: chalk.rgb(255, 50, 50),
  SAFE: chalk.rgb(50, 180, 255),
  SNAKE: chalk.rgb(200, 200, 200),
  HEAD: chalk.rgb(255, 255, 100),
  HIGHLIGHT: chalk.rgb(255, 215, 0),
  OBSTACLE: chalk.rgb(150, 75, 0),
  CONFIDENCE: {
    HIGH: chalk.rgb(50, 255, 50),
    MEDIUM: chalk.rgb(255, 255, 50),
    LOW: chalk.rgb(255, 100, 50)
  },
  INFO: chalk.rgb(100, 200, 255),
  AUTO: chalk.rgb(100, 255, 200),
  AI: chalk.rgb(255, 180, 100),
  DANGER_ZONE: chalk.bgRgb(100, 0, 0).rgb(255, 255, 255),
  SAFE_ZONE: chalk.bgRgb(0, 100, 0).rgb(255, 255, 255),
  NEUTRAL_ZONE: chalk.bgRgb(70, 70, 70).rgb(255, 255, 255)
};

const PATH_SYMBOLS = { UP: '↑', DOWN: '↓', LEFT: '←', RIGHT: '→' };
const DIRECTION_EMOJI = { UP: '⬆️', DOWN: '⬇️', LEFT: '⬅️', RIGHT: '➡️', SCAN: '🔍' };
const MAX_HISTORY_OCCURRENCES = 8;

const moveHistory: Array<{
  timestamp: number,
  gameTime: number,
  action: ActionType,
  tool: string,
  confidence?: number,
  safetyScore?: number,
  pathLength?: number,
  openSpaces?: number,
  isAutoPilot?: boolean,
  isAIMove?: boolean,
  moveRemaining?: number
}> = [];

const groupedHistory: Array<{
  action: ActionType,
  count: number,
  startTime: number,
  endTime: number,
  tool: string,
  confidence?: number,
  safetyScore?: number,
  pathLength?: number,
  openSpaces?: number,
  isAutoPilot?: boolean,
  isAIMove?: boolean,
  moveRemaining?: number
}> = [];

let gameStartTime = Date.now();
let tokenUsage = {
  promptTokens: 0,
  completionTokens: 0,
  totalTokens: 0,
  history: [] as Array<{ timestamp: number, prompt: number, completion: number, total: number }>
};

export const addMoveToHistory = (
  action: ActionType,
  tool: string = "move",
  metrics: {
    confidence?: number,
    safetyScore?: number,
    pathLength?: number,
    openSpaces?: number,
    isAutoPilot?: boolean,
    isAIMove?: boolean,
    moveRemaining?: number
  } = {}
): void => {
  const currentTime = Date.now();
  const gameElapsedTime = Math.floor((currentTime - gameStartTime) / 1000);
  
  moveHistory.push({ 
    timestamp: currentTime,
    gameTime: gameElapsedTime,
    action,
    tool,
    ...metrics
  });
  
  const lastGroup = groupedHistory.length > 0 ? groupedHistory[groupedHistory.length - 1] : null;
  
  if (lastGroup && lastGroup.action === action && lastGroup.tool === tool &&
      lastGroup.isAutoPilot === metrics.isAutoPilot && lastGroup.isAIMove === metrics.isAIMove) {
    lastGroup.count++;
    lastGroup.endTime = gameElapsedTime;
    if (metrics.confidence !== undefined) lastGroup.confidence = metrics.confidence;
    if (metrics.safetyScore !== undefined) lastGroup.safetyScore = metrics.safetyScore;
    if (metrics.pathLength !== undefined) lastGroup.pathLength = metrics.pathLength;
    if (metrics.openSpaces !== undefined) lastGroup.openSpaces = metrics.openSpaces;
    if (metrics.moveRemaining !== undefined) lastGroup.moveRemaining = metrics.moveRemaining;
  } else {
    groupedHistory.push({
      action,
      count: 1,
      startTime: gameElapsedTime,
      endTime: gameElapsedTime,
      tool,
      ...metrics,
    });
    
    if (groupedHistory.length > MAX_HISTORY_OCCURRENCES) groupedHistory.shift();
  }
};

export const updateTokenUsage = (
  usage: { promptTokens: number, completionTokens: number, totalTokens: number },
  response?: any
): void => {
  let promptTokens = usage.promptTokens || 0;
  let completionTokens = usage.completionTokens || 0;

  if (response?.body?.usageMetadata) {
    promptTokens = response.body.usageMetadata.promptTokenCount || promptTokens;
    completionTokens = response.body.usageMetadata.candidatesTokenCount || completionTokens;
  }

  const totalTokens = promptTokens + completionTokens;

  tokenUsage.promptTokens += promptTokens;
  tokenUsage.completionTokens += completionTokens;
  tokenUsage.totalTokens += totalTokens;

  tokenUsage.history.push({
    timestamp: Date.now(),
    prompt: promptTokens,
    completion: completionTokens,
    total: totalTokens,
  });

  if (tokenUsage.history.length > 10) {
    tokenUsage.history.shift();
  }
};

export const renderGameFrame = async (
  state: GameState,
  moveNumber: number,
  _decisionMetrics: Record<string, any> = {},
  _showThinking = false,
  timeRemaining?: number
): Promise<void> => {
  clearScreen();
  
  if (moveNumber === 0) {
    gameStartTime = Date.now();
    moveHistory.length = 0;
    groupedHistory.length = 0;
    tokenUsage = { promptTokens: 0, completionTokens: 0, totalTokens: 0, history: [] };
  }
  
  if (_decisionMetrics.scanning) {
    addMoveToHistory('SCAN', 'scan', { confidence: 1, safetyScore: 0, pathLength: 0, openSpaces: 0 });
  } else if (_decisionMetrics.direction) {
    const spatialAnalysis = analyzeSpatialState(state);
    const pathToFood = findShortestPathToFood(state);
    const safetyScore = spatialAnalysis.safetyScores?.[_decisionMetrics.direction as Direction] || 0;
    
    addMoveToHistory(
      _decisionMetrics.direction,
      _decisionMetrics.isAIMove ? 'ai_move' : 'auto_move',
      {
        confidence: _decisionMetrics.confidence,
        safetyScore,
        pathLength: pathToFood.length,
        openSpaces: spatialAnalysis.openSpaces,
        isAutoPilot: _decisionMetrics.isAutoPilot,
        isAIMove: _decisionMetrics.isAIMove,
        moveRemaining: _decisionMetrics.moveRemaining,
      }
    );
  }
  
  const layout = renderUnifiedLayout(state, moveNumber, timeRemaining, _decisionMetrics, _showThinking);
  console.log(layout);
  
  if (state.isGameOver) {
    console.log(timeRemaining === 0
      ? chalk.magenta('=== TIME\'S UP! ===')
      : chalk.red('=== GAME OVER ==='));
    console.log(chalk.yellow(`Final Score: ${state.score}`));
  }
};

const renderUnifiedLayout = (
  state: GameState,
  moveNumber: number,
  timeRemaining?: number,
  _decisionMetrics: Record<string, any> = {},
  _showThinking = false
): string => {
  const { gridSize, snake, food, score, obstacles = [], visibilityRadius } = state;
  const optimalPath = findShortestPathToFood(state);
  const head = snake[0];
  const spatialAnalysis = analyzeSpatialState(state);
  const suggestion = suggestOptimalMove(state);
  
  // const gameGridWidth = gridSize * 2 + 2; // Unused
  // const visualGridWidth = gridSize * 2 + 2; // Unused
  const spacing = 3;
  
  const lines: string[] = [];
  lines.push(chalk.cyan(`=== Snake Game ===`) + ' '.repeat(15) + chalk.cyan('=== AI Thinking Board ==='));
  const timeDisplay = timeRemaining !== undefined ? formatTime(timeRemaining) : '';
  lines.push(chalk.gray(`Move ${moveNumber} | Score: ${score}${timeRemaining !== undefined ? ` | Time: ${chalk.yellow(timeDisplay)}` : ''}`));
  lines.push('');
  
  // Build path data
  const pathData: Array<{x: number, y: number, dir: Direction}> = [];
  if (optimalPath.length > 0) {
    let curPos = {...snake[0]};
    for (let i = 0; i < optimalPath.length; i++) {
      const dir = optimalPath[i];
      curPos = calculateNewPosition(curPos, dir, gridSize);
      pathData.push({...curPos, dir});
    }
  }
  
  // Build danger zones (potential traps/dead ends)
  const dangerZones = new Set<string>();
  if (spatialAnalysis.potentialTraps) {
    spatialAnalysis.potentialTraps.forEach(trap => {
      dangerZones.add(`${trap.x},${trap.y}`);
    });
  }
  
  // Game grid rendering
  const gameGrid: string[][] = Array(gridSize).fill(0).map(() => Array(gridSize).fill('..'));
  if (snake[0].y < gridSize && snake[0].x < gridSize) gameGrid[snake[0].y][snake[0].x] = '👀';
  for (const segment of snake.slice(1)) {
    if (segment.y < gridSize && segment.x < gridSize) gameGrid[segment.y][segment.x] = '🟢';
  }
  if (food.y < gridSize && food.x < gridSize) gameGrid[food.y][food.x] = '🍎';
  for (const obstacle of obstacles) {
    if (obstacle.y < gridSize && obstacle.x < gridSize) gameGrid[obstacle.y][obstacle.x] = '🧱';
  }
  
  // Visual grid for AI thinking
  const visualGrid: string[][] = Array(gridSize).fill(0).map(() => Array(gridSize).fill('..'));
  
  // Add visibility limits and obstacles
  for (let y = 0; y < gridSize; y++) {
    for (let x = 0; x < gridSize; x++) {
      
      if (visibilityRadius !== undefined) {
        const distanceToHead = Math.max(Math.abs(x - head.x), Math.abs(y - head.y));
        if (distanceToHead > visibilityRadius) {
          visualGrid[y][x] = '??';
          continue;
        }
      }
      
      if (obstacles.some(o => o.x === x && o.y === y)) {
        visualGrid[y][x] = '##';
      }
    }
  }
  
  // Add snake, food and path to visual grid
  if (snake[0].y < gridSize && snake[0].x < gridSize) visualGrid[snake[0].y][snake[0].x] = '👀';
  for (const segment of snake.slice(1)) {
    if (segment.y < gridSize && segment.x < gridSize && visualGrid[segment.y][segment.x] === '..') 
      visualGrid[segment.y][segment.x] = '◯◯';
  }
  if (food.y < gridSize && food.x < gridSize && visualGrid[food.y][food.x] === '..') 
    visualGrid[food.y][food.x] = '🍎';
  
  // Add direction indicators and safety zones
  const validMoves = getValidMoves(state);
  for (const dir of validMoves) {
    const nextPos = calculateNewPosition(head, dir, gridSize);
    if (nextPos.y < gridSize && nextPos.x < gridSize && visualGrid[nextPos.y][nextPos.x] === '..') {
      const safetyScore = spatialAnalysis.safetyScores?.[dir] || 0;
      
      if (dangerZones.has(`${nextPos.x},${nextPos.y}`)) {
        visualGrid[nextPos.y][nextPos.x] = `${PATH_SYMBOLS[dir]}!`;
      } else if (dir === suggestion.direction) {
        visualGrid[nextPos.y][nextPos.x] = `${PATH_SYMBOLS[dir]}${PATH_SYMBOLS[dir]}`;
      } else {
        visualGrid[nextPos.y][nextPos.x] = `${PATH_SYMBOLS[dir]}${safetyScore}`;
      }
    }
  }
  
  // Add optimal path
  for (const {x, y, dir} of pathData) {
    if (y < gridSize && x < gridSize && 
        (visualGrid[y][x] === '..' || visualGrid[y][x].includes(PATH_SYMBOLS[dir]))) {
      visualGrid[y][x] = PATH_SYMBOLS[dir].repeat(2);
    }
  }
  
  // Render grids
  const gameGridTopBorder = chalk.cyan('+' + '-'.repeat(gridSize * 2) + '+');
  const visualGridTopBorder = chalk.cyan('+' + '-'.repeat(gridSize * 2) + '+');
  
  lines.push(gameGridTopBorder + ' '.repeat(spacing) + visualGridTopBorder);
  
  for (let y = 0; y < gridSize; y++) {
    let gameLine = chalk.cyan('|');
    for (let x = 0; x < gridSize; x++) gameLine += gameGrid[y][x];
    gameLine += chalk.cyan('|');
    
    let visualLine = chalk.cyan('|');
    for (let x = 0; x < gridSize; x++) {
      const cell = visualGrid[y][x];
      
      if (cell === '👀') {
        visualLine += cell;
      } else if (cell === '◯◯') {
        visualLine += COLORS.SNAKE('◯◯');
      } else if (cell === '🍎') {
        visualLine += cell;
      } else if (cell === '##') {
        visualLine += COLORS.OBSTACLE('##');
      } else if (cell === '??') {
        visualLine += chalk.gray('??');
      } else if (cell.endsWith('!')) {
        visualLine += COLORS.DANGER_ZONE(cell);
      } else if (cell === '↑↑' || cell === '↓↓' || cell === '←←' || cell === '→→') {
        visualLine += COLORS.PATH(cell);
      } else if (cell.length === 2 && (cell.startsWith('↑') || cell.startsWith('↓') || 
                                        cell.startsWith('←') || cell.startsWith('→'))) {
        const safetyScore = parseInt(cell[1]);
        if (safetyScore >= 3) {
          visualLine += COLORS.SAFE_ZONE(cell);
        } else if (safetyScore <= 1) {
          visualLine += COLORS.DANGER_ZONE(cell);
        } else {
          visualLine += COLORS.NEUTRAL_ZONE(cell);
        }
      } else {
        visualLine += '..';
      }
    }
    visualLine += chalk.cyan('|');
    
    lines.push(gameLine + ' '.repeat(spacing) + visualLine);
  }
  
  const gameGridBottomBorder = chalk.cyan('+' + '-'.repeat(gridSize * 2) + '+');
  const visualGridBottomBorder = chalk.cyan('+' + '-'.repeat(gridSize * 2) + '+');
  lines.push(gameGridBottomBorder + ' '.repeat(spacing) + visualGridBottomBorder);
  
  // Only show legend, AI thoughts, decision history, and tools when game is not over
  if (!state.isGameOver) {
    // Render legend
    lines.push('');
    lines.push(chalk.cyan('=== Legend ==='));
    lines.push(`👀 Snake head   🍎 Food   🟢/◯◯ Snake body   🧱/## Obstacle   ${COLORS.PATH('↑↑')} Optimal path`);
    lines.push(`${COLORS.SAFE_ZONE('↑3')} Safe move   ${COLORS.NEUTRAL_ZONE('↑2')} Neutral move   ${COLORS.DANGER_ZONE('↑1')} Risky move   ${COLORS.DANGER_ZONE('↑!')} Danger zone`);
    
    // Render thinking stats
    if (groupedHistory.length > 0) {
      lines.push('');
      lines.push(chalk.cyan('=== AI Thoughts ==='));
      
      lines.push(`Score: ${chalk.yellow(score.toString())}   Open Space: ${COLORS.INFO(spatialAnalysis.openSpaces.toString())}   Regions: ${spatialAnalysis.connectedRegions}`);
      
      const confidenceValue = Math.round(suggestion.confidence * 100);
      let confidenceColor = COLORS.CONFIDENCE.MEDIUM;
      if (suggestion.confidence >= 0.7) confidenceColor = COLORS.CONFIDENCE.HIGH;
      else if (suggestion.confidence < 0.4) confidenceColor = COLORS.CONFIDENCE.LOW;
      
      const directionEmoji = DIRECTION_EMOJI[suggestion.direction] || '';
      lines.push(`Suggested move: ${directionEmoji} ${suggestion.direction} ${confidenceColor(`(${confidenceValue}%)`)}   Path to food: ${spatialAnalysis.optimalPathLength ? COLORS.PATH(`${spatialAnalysis.optimalPathLength} steps`) : COLORS.DANGER('None')}`);
      
      // Render decision history
      lines.push('');
      lines.push(chalk.cyan('=== Decision History ==='));
      
      groupedHistory.forEach(group => {
        const dirEmoji = DIRECTION_EMOJI[group.action as keyof typeof DIRECTION_EMOJI] || '';
        const timeInfo = group.count > 1 ? `${group.startTime}s-${group.endTime}s` : `${group.startTime}s`;
        
        let confidenceText = '-';
        if (group.confidence !== undefined) {
          const confValue = Math.round(group.confidence * 100);
          let confColor = COLORS.CONFIDENCE.MEDIUM;
          if (group.confidence >= 0.7) confColor = COLORS.CONFIDENCE.HIGH;
          else if (group.confidence < 0.4) confColor = COLORS.CONFIDENCE.LOW;
          confidenceText = confColor(`${confValue}%`);
        }
        
        const safetyText = group.safetyScore !== undefined
          ? COLORS.SAFE(`${group.safetyScore}`)
          : '-';
        
        const pathText = group.pathLength !== undefined
          ? COLORS.PATH(`${group.pathLength}`)
          : '-';
        
        let displayText = `${timeInfo}: ${dirEmoji} ${group.action} ${confidenceText}   Safety: ${safetyText}   Path: ${pathText}`;
        
        if (group.tool === 'scan') {
          displayText = `${timeInfo}: ${dirEmoji} SCAN`;
        } else if (group.isAutoPilot) {
          const moreInfo = group.moveRemaining !== undefined ? ` +${group.moveRemaining} more` : '';
          displayText = `${timeInfo}: ${COLORS.AUTO(`🚀 AUTO ${group.action}`)} ${confidenceText}   Path: ${pathText}${moreInfo}`;
        } else if (group.isAIMove) {
          displayText = `${timeInfo}: ${COLORS.AI(`🧠 AI ${group.action}`)} ${confidenceText}   Safety: ${safetyText}   Path: ${pathText}`;
        }
        
        lines.push(displayText);
      });
      
      // Render token usage
      lines.push('');
      lines.push(chalk.cyan('=== Token Usage ==='));
      const tokenHistory = tokenUsage.history.map(h => `${h.prompt}/${h.completion}`).join(' | ');
      lines.push(`Total: ${chalk.yellow(tokenUsage.totalTokens)} (${chalk.cyan(tokenUsage.promptTokens)}P/${chalk.cyan(tokenUsage.completionTokens)}C)`);
      lines.push(`History: ${tokenHistory}`);
    }

    lines.push('');
    lines.push(chalk.cyan('=== Available Tools ==='));
    lines.push(`- scan: Analyze game state and get optimal moves.`);
    lines.push(`- move: Move the snake in a sequence of up to 5 directions.`);
    lines.push(`- giveUp: Give up the game if stuck.`);
  }
  
  return lines.join('\n');
};