Refactor search logic