Add missing files and fix build