some fixes & moving code to where it belongs to