Fix AiState and cleanup the code