Restructure + fix imports