Improve design and fix JS bug