#include "version.h"
#include "wallabylib.h"
#include "diag_main.h"

void v517patch();
void v622patch();

void pw_menu();
void about();
void backup_boot_to_sd();
void backup_ce_to_sd();
void backup_bootce_to_sd();
void backup_RAM_to_sd();
void restore_sd();
void pw_status();
void pw_manip(int cmd);
int get_menu_selection( int line, int num_options, char **menu_options );
int wait_for_menu_button();

void diag_main() {
    char *str_stdout_detected = "Wallaby Bootloader v%s detected\n";
    char *str_display_detected = "Detected:";
    
    _WALLABYLIB( printf, "WallabyPatchTool v1.5\n" );
    clear_display();
    switch ( wallaby_version ) {
        case WALLABY_514:
            _WALLABYLIB( printf, str_stdout_detected, REQUIRED_514_WALLABY_IDENT );
            _WALLABYLIB( display_set_cursor, 4, 0 );
            _WALLABYLIB( display_print, str_display_detected );
            _WALLABYLIB( display_set_cursor, 4, 9 );
            _WALLABYLIB( display_print, REQUIRED_514_WALLABY_IDENT );
            _WALLABYLIB( display_set_cursor, 5, 0 );
            _WALLABYLIB( display_print, "No Patch Need" );
            break;
        case WALLABY_515:
            _WALLABYLIB( printf, str_stdout_detected, REQUIRED_515_WALLABY_IDENT );
            _WALLABYLIB( display_set_cursor, 4, 0 );
            _WALLABYLIB( display_print, str_display_detected );
            _WALLABYLIB( display_set_cursor, 4, 9 );
            _WALLABYLIB( display_print, REQUIRED_515_WALLABY_IDENT );
            _WALLABYLIB( display_set_cursor, 5, 0 );
            _WALLABYLIB( display_print, "No Patch Need" );
            break;
        case WALLABY_515b:
            _WALLABYLIB( printf, str_stdout_detected, REQUIRED_515b_WALLABY_IDENT );
            _WALLABYLIB( display_set_cursor, 4, 0 );
            _WALLABYLIB( display_print, str_display_detected );
            _WALLABYLIB( display_set_cursor, 4, 9 );
            _WALLABYLIB( display_print, REQUIRED_515b_WALLABY_IDENT );
            _WALLABYLIB( display_set_cursor, 5, 0 );
            _WALLABYLIB( display_print, "No Patch Need" );
            break;

        case WALLABY_517:
            _WALLABYLIB( printf, str_stdout_detected, REQUIRED_517_WALLABY_IDENT );
            _WALLABYLIB( display_set_cursor, 4, 0 );
            _WALLABYLIB( display_print, str_display_detected );
            _WALLABYLIB( display_set_cursor, 4, 9 );
            _WALLABYLIB( display_print, REQUIRED_517_WALLABY_IDENT );
            _WALLABYLIB( display_set_inverse, 1 );
            _WALLABYLIB( display_set_cursor, 5, 0 );
            _WALLABYLIB( display_print, " Patching... " );
            _WALLABYLIB( display_set_inverse, 0 );
            v517patch();
            _WALLABYLIB( printf, "Bootloader patched\n" );
            _WALLABYLIB( display_set_cursor, 5, 0 );
            _WALLABYLIB( display_print, " Patch Done! " );
            break;
        case WALLABY_622:
            _WALLABYLIB( printf, str_stdout_detected, REQUIRED_622_WALLABY_IDENT );
            _WALLABYLIB( display_set_cursor, 4, 0 );
            _WALLABYLIB( display_print, str_display_detected );
            _WALLABYLIB( display_set_cursor, 4, 9 );
            _WALLABYLIB( display_print, REQUIRED_622_WALLABY_IDENT );
            _WALLABYLIB( display_set_inverse, 1 );
            _WALLABYLIB( display_set_cursor, 5, 0 );
            _WALLABYLIB( display_print, " Patching... " );
            _WALLABYLIB( display_set_inverse, 0 );
            v622patch();
            _WALLABYLIB( printf, "Bootloader patched\n" );
            _WALLABYLIB( display_set_cursor, 5, 0 );
            _WALLABYLIB( display_print, " Patch Done! " );
            break;

    }
    pause( 1 );
    
    _WALLABYLIB( display_set_cursor, 7, 0 );
    _WALLABYLIB( display_print, "Press ACTION " );
    
    wait_for_action_key();
    
    main_menu();
}


void pause( int cycles ) {
    int c;
    int i;
    
    for ( c = cycles; c >= 0; c-- ) {
        for ( i = 0x01000000; i >= 0; i-- ) {
        }
    }
}

int wait_for_button() {
    int bcode;
    int selected_bcode;
    
    do {
        _WALLABYLIB_FUNC( get_button_code );
        bcode = _WALLABYLIB_CALL();
    } while ( bcode != 0 );
    do {
        _WALLABYLIB_FUNC( get_button_code );
        bcode = _WALLABYLIB_CALL();
    } while ( bcode == 0 );
    selected_bcode = bcode;
    do {
        _WALLABYLIB_FUNC( get_button_code );
        bcode = _WALLABYLIB_CALL();
    } while ( bcode != 0 );
    
    return selected_bcode;
}

void wait_for_action_key() {
    int bcode;
    
    do {
        bcode = wait_for_button();
    } while ( bcode != 1 );
}

void clear_display() {
    _WALLABYLIB( set_backlight, 1 );
    _WALLABYLIB( display_clear );
    _WALLABYLIB( display_set_cursor, 0, 0 );
    _WALLABYLIB( display_print, "Wallaby Patch" );
    _WALLABYLIB( display_set_cursor, 1, 0 );
    _WALLABYLIB( display_print, "Tool 1.5/    " );
    _WALLABYLIB( display_set_cursor, 2, 0 );
    _WALLABYLIB( display_print, "-------------" );
    _WALLABYLIB( display_set_cursor, 1, 9 );
    switch ( wallaby_version ) {
        case WALLABY_514:
            _WALLABYLIB( display_print, REQUIRED_514_WALLABY_IDENT );
            break;
        case WALLABY_515:
            _WALLABYLIB( display_print, REQUIRED_515_WALLABY_IDENT );
            break;
        case WALLABY_515b:
            _WALLABYLIB( display_print, REQUIRED_515b_WALLABY_IDENT );
            break;
        case WALLABY_517:
            _WALLABYLIB( display_print, REQUIRED_517_WALLABY_IDENT );
            break;
        case WALLABY_622:
            _WALLABYLIB( display_print, REQUIRED_622_WALLABY_IDENT );
            break;
    }
}

void main_menu() {
    char *menu_options[] = { "Boot -> SD   ", "CE -> SD     ", "Boot+CE -> SD", "RAM -> SD    ", "Password Menu", "Restore SD   ", "About        ", "Exit         " };
    int menu_choice;
    
    while ( 1 ) {
        clear_display();
        menu_choice = get_menu_selection( 4, 8, menu_options );
        switch ( menu_choice ) {
            case 0:
                backup_boot_to_sd();
                break;
            case 1:
                backup_ce_to_sd();
                break;
            case 2:
                backup_bootce_to_sd();
                break;
            case 3:
                backup_RAM_to_sd();
                break;
            case 4:
                pw_menu();
                break;
            case 5:
                restore_sd();
                break;
            case 6:
                about();
                break;
            case 7:
                return;
        }
    }
}

void about() {
    clear_display();
    
    _WALLABYLIB( display_set_cursor, 4, 0 );
    _WALLABYLIB( display_print, "This patch is" );
    _WALLABYLIB( display_set_cursor, 5, 0 );
    _WALLABYLIB( display_print, "   made by   " );
    _WALLABYLIB( display_set_cursor, 6, 0 );
    _WALLABYLIB( display_print, "www.xda-devel" );
    _WALLABYLIB( display_set_cursor, 7, 0 );
    _WALLABYLIB( display_print, "  opers.com  " );
    _WALLABYLIB( display_set_cursor, 9, 0 );
    _WALLABYLIB( display_print, "Press ACTION " );
    
    wait_for_action_key();
}

void backup_rom( int image_start, int image_size ) {
    char *menu_options[] = { "Abort        ", "Continue     " };
    int menu_choice;
    
    clear_display();
    
    _WALLABYLIB( display_set_cursor, 4, 0 );
    _WALLABYLIB( display_print, "  Insert or  " );
    _WALLABYLIB( display_set_cursor, 5, 0 );
    _WALLABYLIB( display_print, " replace SD. " );
    _WALLABYLIB( display_set_cursor, 7, 0 );
    _WALLABYLIB( display_print, " SD will be  " );
    _WALLABYLIB( display_set_cursor, 8, 0 );
    _WALLABYLIB( display_print, "overwritten! " );
    
    menu_choice = get_menu_selection( 10, 2, menu_options );
    switch ( menu_choice ) {
        case 1:
            _WALLABYLIB( backup_rom_to_sd, image_start, image_size );
    }
}

void backup_boot_to_sd() {
    backup_rom( 0xa0000000, 0x40000 );
}

void backup_ce_to_sd() {
    backup_rom( 0xa0040000, 0x1fc0000 );
}

void backup_bootce_to_sd() {
    backup_rom( 0xa0000000, 0x2000000 );
}

void backup_RAM_to_sd() {
    backup_rom( 0xac000000, 0x2000000 );
}

void restore_sd() {
    char *menu_options[] = { "Abort        ", "Continue     " };
    int menu_choice;
    
    clear_display();
    
    _WALLABYLIB( display_set_cursor, 4, 0 );
    _WALLABYLIB( display_print, "Insert valid " );
    _WALLABYLIB( display_set_cursor, 5, 0 );
    _WALLABYLIB( display_print, "SD ROM image." );
    _WALLABYLIB( display_set_cursor, 7, 0 );
    _WALLABYLIB( display_print, "  WARNING:   " );
    _WALLABYLIB( display_set_cursor, 8, 0 );
    _WALLABYLIB( display_print, " Upgrade ROM " );
    _WALLABYLIB( display_set_cursor, 9, 0 );
    _WALLABYLIB( display_print, " at your own " );
    _WALLABYLIB( display_set_cursor, 10, 0 );
    _WALLABYLIB( display_print, "   risk!     " );
    
    menu_choice = get_menu_selection( 12, 2, menu_options );
    switch ( menu_choice ) {
        case 1:
            _WALLABYLIB( restore_sd_to_rom );
    }
}    

#define PW_DEACTIVATE	1
#define PW_ACTIVATE	2
#define PW_WIPE		3

void pw_menu()
{
    char *menu_options[] = { "Show PW stats ", "Deactivate PW", "Activate PW", "Wipe PW", "Return Main" };
    int menu_choice;
    int quit = 0;

    while(!quit) {
        clear_display();
        menu_choice = get_menu_selection( 4, 5, menu_options );
        switch ( menu_choice ) {
            case 0:
	        pw_status();
                break;
            case 1:
                pw_manip(PW_DEACTIVATE);
                break;
            case 2:
                pw_manip(PW_ACTIVATE);
                break;
            case 3:
                pw_manip(PW_WIPE);
                break;
    	    case 4:
	    default:
	        quit = 1;
	}
    }
}

void pw_status()
{
    unsigned char *i;
    unsigned char val;
    //char *str_stdout = "Address %x\n";

    clear_display();
    
    _WALLABYLIB( display_set_cursor, 4, 0 );
    _WALLABYLIB( display_print, "Searching ..." );

    for (i = (unsigned char *)0xAC000000; i<(unsigned char *)0xAE000000; i+=4) {
	if (*(int *)i == 0x4D494B45 && *(int *)(i+4) == 0x4D494B45) break;
    }
   
    if (i<(unsigned char *)0xAE000000) {
        _WALLABYLIB( display_set_cursor, 6, 0 );
        _WALLABYLIB( display_print, "Found Heap at" );
        _WALLABYLIB( display_set_cursor, 7, 0 );
        _WALLABYLIB( display_print_hex_word, i-4 );
        _WALLABYLIB( display_set_cursor, 9, 0 );
        _WALLABYLIB( display_print, "  Password   " );
        _WALLABYLIB( display_set_cursor, 10, 0 );

	val = *(i+0x3c);
        if (val & 1) {
	    _WALLABYLIB( display_print, "     set     " );
	} else {
	    _WALLABYLIB( display_print, "   not set   " );
	}
        _WALLABYLIB( display_set_cursor, 11, 0 );
        _WALLABYLIB( display_print, "     and     " );
        _WALLABYLIB( display_set_cursor, 12, 0 );
        if (val & 2) {
	    _WALLABYLIB( display_print, "  activated  " );
	} else {
	    _WALLABYLIB( display_print, " deactivated " );
	}

    } else {
        _WALLABYLIB( display_set_cursor, 6, 0 );
        _WALLABYLIB( display_print, "Not Found!" );
    }

    _WALLABYLIB( display_set_cursor, 14, 0 );
    _WALLABYLIB( display_print, "Press ACTION " );
    wait_for_action_key();
}

void pw_manip(int cmd)
{
    unsigned char *i;
    unsigned char val;

    clear_display();
    
    _WALLABYLIB( display_set_cursor, 4, 0 );
    _WALLABYLIB( display_print, "Searching ..." );

    for (i = (unsigned char *)0xAC000000; i<(unsigned char *)0xAE000000; i+=4) {
	if (*(int *)i == 0x4D494B45 && *(int *)(i+4) == 0x4D494B45) break;
    }
   
    if (i<(unsigned char *)0xAE000000) {
	val = *(i+0x3c) & 0x3;

	switch (cmd) {
	    case PW_DEACTIVATE:

	        if ((val & 2) == 0) {
                    _WALLABYLIB( display_set_cursor, 6, 0 );
                    _WALLABYLIB( display_print, "Not active!" );
        	} else {
        	    *(i+0x3c) = (*(i+0x3c)) & 0xFD;
                    _WALLABYLIB( display_set_cursor, 6, 0 );
                    _WALLABYLIB( display_print, "Flags set to:" );
                    _WALLABYLIB( display_set_cursor, 7, 0 );
                    _WALLABYLIB( display_print_hex_word, *(i+0x3c) );
                    _WALLABYLIB( display_set_cursor, 9, 0 );
        	    _WALLABYLIB( display_print, "  Password   " );
                    _WALLABYLIB( display_set_cursor, 10, 0 );
	    	    _WALLABYLIB( display_print, " deactivated " );
	        }
		break;

	    case PW_ACTIVATE:
	        if ((val & 2) == 2) {
                    _WALLABYLIB( display_set_cursor, 6, 0 );
                    _WALLABYLIB( display_print, "Already");
                    _WALLABYLIB( display_set_cursor, 7, 0 );
                    _WALLABYLIB( display_print, "active!" );
        	} else {
		    // Activate, but also set 'set' bit just to be sure
        	    *(i+0x3c) = (*(i+0x3c)) | 3;
                    _WALLABYLIB( display_set_cursor, 6, 0 );
                    _WALLABYLIB( display_print, "Flags set to:" );
                    _WALLABYLIB( display_set_cursor, 7, 0 );
                    _WALLABYLIB( display_print_hex_word, *(i+0x3c) );
                    _WALLABYLIB( display_set_cursor, 9, 0 );
        	    _WALLABYLIB( display_print, "  Password   " );
                    _WALLABYLIB( display_set_cursor, 10, 0 );
	    	    _WALLABYLIB( display_print, "  activated  " );
		}
		break;

	    case PW_WIPE:
	        if ((val & 3) == 0) {
                    _WALLABYLIB( display_set_cursor, 6, 0 );
                    _WALLABYLIB( display_print, "Not active!" );
        	} else {
        	    *(i+0x3c) = (*(i+0x3c)) & 0xFC;
                    _WALLABYLIB( display_set_cursor, 6, 0 );
                    _WALLABYLIB( display_print, "Flags set to:" );
                    _WALLABYLIB( display_set_cursor, 7, 0 );
                    _WALLABYLIB( display_print_hex_word, *(i+0x3c) );
                    _WALLABYLIB( display_set_cursor, 9, 0 );
        	    _WALLABYLIB( display_print, "  Password   " );
                    _WALLABYLIB( display_set_cursor, 10, 0 );
	    	    _WALLABYLIB( display_print, "    wiped    " );
		}
		break;

	    default:
                _WALLABYLIB( display_set_cursor, 6, 0 );
                _WALLABYLIB( display_print, "Unkown CMD" );
	}
    } else {
        _WALLABYLIB( display_set_cursor, 6, 0 );
        _WALLABYLIB( display_print, "Not Found!" );
    }

    _WALLABYLIB( display_set_cursor, 12, 0 );
    _WALLABYLIB( display_print, "Press ACTION " );
    wait_for_action_key();
}

int get_menu_selection( int line, int num_options, char **menu_options ) {
    int current_choice = 0;
    int button_code;
    int i;
    
    while ( 1 ) {
        for ( i = 0; i < num_options; i++ ) {
            _WALLABYLIB( display_set_cursor, line + i, 0 );
            _WALLABYLIB( display_set_inverse, ( i == current_choice ) ? 1 : 0 );
            _WALLABYLIB( display_print, menu_options[ i ] );
        }
        _WALLABYLIB( display_set_inverse, 0 );
        
        button_code = wait_for_menu_button();
        switch ( button_code ) {
            case 2:
                current_choice--;
                if ( current_choice < 0 ) current_choice = num_options - 1;
                break;
            case 8:
                current_choice++;
                if ( current_choice >= num_options ) current_choice = 0;
                break;
            case 1:
                return current_choice;
        }
    }
}
    

int wait_for_menu_button() {
    int bcode;
    
    do {
        bcode = wait_for_button();
    } while ( ( bcode != 2 ) && ( bcode != 8 ) && ( bcode != 1 ) );
    
    return bcode;
}