Improve main.c more