WIP: Add controller, refactor a bit.