fixes and refactoring
fixed issue with context, refactored