code rework, can't test now, hope it works