Refactor drawing logic