JS code refactor + styling fixes