Bug fixes and UI enhancements