Refactored the AI logic.