Mały refactor i bugfixing