#include <ctype.h>
#include <errno.h>
#include <fcntl.h>
#include <getopt.h>
#include <limits.h>
#include <linux/input.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/reboot.h>
#include <sys/types.h>
#include <time.h>
#include <unistd.h>

#include <sys/wait.h>
#include <sys/limits.h>
#include <dirent.h>
#include <sys/stat.h>

#include <signal.h>
#include <sys/wait.h>

#include "bootloader.h"
#include "common.h"
#include "cutils/properties.h"
#include "firmware.h"
#include "install.h"
#include "minui/minui.h"
#include "minzip/DirUtil.h"
#include "roots.h"
#include "recovery_ui.h"

#include "commands.h"
#include "amend/amend.h"

#include "mtdutils/dump_image.h"
#include "../../external/yaffs2/yaffs2/utils/mkyaffs2image.h"
#include "../../external/yaffs2/yaffs2/utils/unyaffs.h"

int signature_check_enabled = 1;
int script_assert_enabled = 1;
static const char *SDCARD_PACKAGE_FILE = "SDCARD:update.zip";

void
toggle_signature_check()
{
    signature_check_enabled = !signature_check_enabled;
    ui_print("Signature Check: %s\n", signature_check_enabled ? "Enabled" : "Disabled");
}

void toggle_script_asserts()
{
    script_assert_enabled = !script_assert_enabled;
    ui_print("Script Asserts: %s\n", script_assert_enabled ? "Enabled" : "Disabled");
}

int install_zip(const char* packagefilepath)
{
    ui_print("\n-- Installing: %s\n", packagefilepath);
    set_sdcard_update_bootloader_message();
    int status = install_package(packagefilepath);
	ui_reset_progress();
    if (status != INSTALL_SUCCESS) {
        ui_set_background(BACKGROUND_ICON_ERROR);
        ui_print("Installation aborted.\n");
        return 1;
    } 
    if (firmware_update_pending()) {
        ui_print("\nReboot via menu to complete\ninstallation.\n");
    }
    ui_set_background(BACKGROUND_ICON_NONE);
    ui_print("\nInstall from sdcard complete.\n");
    return 0;
}

char* INSTALL_MENU_ITEMS[] = {  "apply sdcard:update.zip",
                                "choose zip from sdcard",
                                "toggle signature verification",
                                "toggle script asserts",
                                NULL };
#define ITEM_APPLY_SDCARD     0
#define ITEM_CHOOSE_ZIP       1
#define ITEM_SIG_CHECK        2
#define ITEM_ASSERTS          3

void show_install_update_menu()
{
    static char* headers[] = {  "Apply update from .zip file on SD card",
                                "",
                                NULL 
    };
    for (;;)
    {
        int chosen_item = get_menu_selection(headers, INSTALL_MENU_ITEMS, 0);
        switch (chosen_item)
        {
            case ITEM_ASSERTS:
                toggle_script_asserts();
                break;
            case ITEM_SIG_CHECK:
                toggle_signature_check();
                break;
            case ITEM_APPLY_SDCARD:
                install_zip(SDCARD_PACKAGE_FILE);
                break;
            case ITEM_CHOOSE_ZIP:
                show_choose_zip_menu();
                break;
            default:
                return;
        }
        
    }
}

char** gather_files(const char* directory, const char* fileExtensionOrDirectory, int* numFiles)
{
    char path[PATH_MAX] = "";
    DIR *dir;
    struct dirent *de;
    int total = 0;
    int i;
    char** files;
    int pass;
    *numFiles = 0;
    int dirLen = strlen(directory);

    dir = opendir(directory);
    if (dir == NULL) {
        ui_print("Couldn't open directory.\n");
        return NULL;
    }
  
    int extension_length;
    if (fileExtensionOrDirectory != NULL)
        extension_length = strlen(fileExtensionOrDirectory);
  
    int isCounting = 1;
    i = 0;
    for (pass = 0; pass < 2; pass++) {
        while ((de=readdir(dir)) != NULL) {
            // skip hidden files
            if (de->d_name[0] == '.')
                continue;
            
            // NULL means that we are gathering directories, so skip this
            if (fileExtensionOrDirectory != NULL)
            {
                // make sure that we can have the desired extension (prevent seg fault)
                if (strlen(de->d_name) < extension_length)
                    continue;
                // compare the extension
                if (strcmp(de->d_name + strlen(de->d_name) - extension_length, fileExtensionOrDirectory) != 0)
                    continue;
            }
            else
            {
                struct stat info;
                char* fullFileName = (char*)malloc(strlen(de->d_name) + dirLen + 1);
                strcpy(fullFileName, directory);
                strcat(fullFileName, de->d_name);
                stat(fullFileName, &info);
                free(fullFileName);
                // make sure it is a directory
                if (!(S_ISDIR(info.st_mode)))
                    continue;
            }
            
            if (pass == 0)
            {
                total++;
                continue;
            }
            
            files[i] = (char*) malloc(dirLen + strlen(de->d_name) + 2);
            strcpy(files[i], directory);
            strcat(files[i], de->d_name);
            if (fileExtensionOrDirectory == NULL)
                strcat(files[i], "/");
            i++;
        }
        if (pass == 1)
            break;
        if (total == 0)
            break;
        rewinddir(dir);
        *numFiles = total;
        files = (char**) malloc((total+1)*sizeof(char*));
        files[total]=NULL;
    }

    if(closedir(dir) < 0) {
        LOGE("Failed to close directory.");
    }

    if (total==0) {
        return NULL;
    }

    return files;
}

void free_string_array(char** array)
{
    char* cursor = array[0];
    int i = 0;
    while (cursor != NULL)
    {
        free(cursor);
        cursor = array[++i];
    }
    free(array);
}

// pass in NULL for fileExtensionOrDirectory and you will get a directory chooser
char* choose_file_menu(const char* directory, const char* fileExtensionOrDirectory, const char* headers[])
{
    char path[PATH_MAX] = "";
    DIR *dir;
    struct dirent *de;
    int numFiles = 0;
    int numDirs = 0;
    int i;

    int dir_len = strlen(directory);

    char** files = gather_files(directory, fileExtensionOrDirectory, &numFiles);
    char** dirs;
    if (fileExtensionOrDirectory != NULL)
        dirs = gather_files(directory, NULL, &numDirs);
    int total = numDirs + numFiles;
    if (total == 0)
    {
        ui_print("No files found.\n");
        return NULL;
    }
    char** list = (char**) malloc((total + 1) * sizeof(char*));
    list[total] = NULL;


    for (i = 0 ; i < numDirs; i++)
    {
        list[i] = strdup(dirs[i] + dir_len);
    }

    for (i = 0 ; i < numFiles; i++)
    {
        list[numDirs + i] = strdup(files[i] + dir_len);
    }

    for (;;)
    {
        int chosen_item = get_menu_selection(headers, list, 0);
        if (chosen_item == GO_BACK)
            break;
        if (chosen_item < numDirs)
        {
            char* subret = choose_file_menu(dirs[chosen_item], fileExtensionOrDirectory, headers);
            if (subret != NULL)
                return subret;
            continue;
        } 
        static char ret[PATH_MAX];
        strcpy(ret, files[chosen_item - numDirs]);
        return ret;
    }
    return NULL;
}

void show_choose_zip_menu()
{
    if (ensure_root_path_mounted("SDCARD:") != 0) {
        LOGE ("Can't mount /sdcard\n");
        return;
    }

    static char* headers[] = {  "Choose a zip to apply",
                                "",
                                NULL 
    };
    
    char* file = choose_file_menu("/sdcard/", ".zip", headers);
    if (file == NULL)
        return;
    char sdcard_package_file[1024];
    strcpy(sdcard_package_file, "SDCARD:");
    strcat(sdcard_package_file,  file + strlen("/sdcard/"));
    install_zip(sdcard_package_file);
}

// This was pulled from bionic: The default system command always looks
// for shell in /system/bin/sh. This is bad.
#define _PATH_BSHELL "/sbin/sh"
#define system recovery_system
extern char **environ;
int
system(const char *command)
{
  pid_t pid;
    sig_t intsave, quitsave;
    sigset_t mask, omask;
    int pstat;
    char *argp[] = {"sh", "-c", NULL, NULL};

    if (!command)        /* just checking... */
        return(1);

    argp[2] = (char *)command;

    sigemptyset(&mask);
    sigaddset(&mask, SIGCHLD);
    sigprocmask(SIG_BLOCK, &mask, &omask);
    switch (pid = vfork()) {
    case -1:            /* error */
        sigprocmask(SIG_SETMASK, &omask, NULL);
        return(-1);
    case 0:                /* child */
        sigprocmask(SIG_SETMASK, &omask, NULL);
        execve(_PATH_BSHELL, argp, environ);
    _exit(127);
  }

    intsave = (sig_t)  bsd_signal(SIGINT, SIG_IGN);
    quitsave = (sig_t) bsd_signal(SIGQUIT, SIG_IGN);
    pid = waitpid(pid, (int *)&pstat, 0);
    sigprocmask(SIG_SETMASK, &omask, NULL);
    (void)bsd_signal(SIGINT, intsave);
    (void)bsd_signal(SIGQUIT, quitsave);
    return (pid == -1 ? -1 : pstat);
}

int print_and_error(char* message)
{
    ui_print(message);
    return 1;
}

// TODO : Separate file for Nandroid?
int do_nandroid_backup(char* backup_name)
{
    if (ensure_root_path_mounted("SDCARD:") != 0) {
        LOGE ("Can't mount /sdcard\n");
        return 1;
    }

    struct timeval tp;
    gettimeofday(&tp, NULL);
    
    char backupdir[PATH_MAX];
    char tmp[PATH_MAX];
    if (NULL != backup_name)
		sprintf(backupdir, "/sdcard/clockworkmod/backup/%s", backup_name);
	else
		sprintf(backupdir, "/sdcard/clockworkmod/backup/%d", tp.tv_sec);
		
    sprintf(tmp, "mkdir -p %s", backupdir);
    system(tmp);

    int ret;
    ui_print("Backing up boot...\n");
    sprintf(tmp, "%s/%s", backupdir, "boot.img");
    ret = dump_image("boot", tmp, NULL);
    if (0 != ret)
        return print_and_error("Error while dumping boot image!\n");
    
    // TODO: Wrap this up in a loop?

    ui_print("Backing up system...\n");
    sprintf(tmp, "%s/%s", backupdir, "system.img");
    if (ensure_root_path_mounted("SYSTEM:") != 0)
        return print_and_error("Can't mount /system!\n");
    ret = mkyaffs2image("/system", tmp, 0, NULL);
    ensure_root_path_unmounted("SYSTEM:");
    if (0 != ret)
        return print_and_error("Error while making a yaffs2 image of system!\n");
    
    ui_print("Backing up data...\n");
    sprintf(tmp, "%s/%s", backupdir, "data.img");
    if (ensure_root_path_mounted("DATA:") != 0)
        return print_and_error("Can't mount /data!\n");
    ret = mkyaffs2image("/data", tmp, 0, NULL);
    ensure_root_path_unmounted("DATA:");
    if (0 != ret)
        return print_and_error("Error while making a yaffs2 image of data!\n");
    
    ui_print("Backing up cache...\n");
    sprintf(tmp, "%s/%s", backupdir, "cache.img");
    if (ensure_root_path_mounted("CACHE:") != 0)
        return print_and_error("Can't mount /cache!\n");
    ret = mkyaffs2image("/cache", tmp, 0, NULL);
    ensure_root_path_unmounted("CACHE:");
    if (0 != ret)
        return print_and_error("Error while making a yaffs2 image of cache!\n");
    
    sprintf(tmp, "md5sum %s/*img > %s/nandroid.md5", backupdir, backupdir);
    system(tmp);
    
    return 0;
}

int do_nandroid_restore(char* backup_path)
{
    if (ensure_root_path_mounted("SDCARD:") != 0) {
        LOGE ("Can't mount /sdcard\n");
        return 1;
    }
    
    char tmp[PATH_MAX];

    ui_print("Checking MD5 sums...\n");
    sprintf(tmp, "md5sum -c %s/nandroid.md5", backup_path);
    if (0 != system(tmp))
        return print_and_error("MD5 mismatch!\n");
    
    ui_print("Restoring system...\n");
    if (0 != ensure_root_path_unmounted("SYSTEM:")) 
        return print_and_error("Can't unmount /system!\n");
    if (0 != format_root_device("SYSTEM:"))
        return print_and_error("Error while formatting /system!\n");
    if (ensure_root_path_mounted("SYSTEM:") != 0)
        return print_and_error("Can't mount /system!\n");
    sprintf(tmp, "%s/system.img", backup_path);
    if (0 != unyaffs(tmp, "/system", NULL))
        return print_and_error("Error while restoring /system!\n");

    ui_print("Restoring data...\n");
    if (0 != ensure_root_path_unmounted("DATA:")) 
        return print_and_error("Can't unmount /data!\n");
    if (0 != format_root_device("DATA:"))
        return print_and_error("Error while formatting /data!\n");
    if (ensure_root_path_mounted("DATA:") != 0)
        return print_and_error("Can't mount /data!\n");
    sprintf(tmp, "%s/data.img", backup_path);
    if (0 != unyaffs(tmp, "/data", NULL))
        return print_and_error("Error while restoring /data!\n");

    ui_print("Restoring cache...\n");
    if (0 != ensure_root_path_unmounted("CACHE:")) 
        return print_and_error("Can't unmount /cache!\n");
    if (0 != format_root_device("CACHE:"))
        return print_and_error("Error while formatting /cache!\n");
    if (ensure_root_path_mounted("CACHE:") != 0)
        return print_and_error("Can't mount /cache!\n");
    sprintf(tmp, "%s/cache.img", backup_path);
    if (0 != unyaffs(tmp, "/cache", NULL))
        return print_and_error("Error while restoring /cache!\n");
        
    return 0;
}

void show_nandroid_restore_menu()
{
    if (ensure_root_path_mounted("SDCARD:") != 0) {
        LOGE ("Can't mount /sdcard\n");
        return 1;
    }
    
    static char* headers[] = {  "Choose an image to restore",
                                "",
                                NULL 
    };

    char* file = choose_file_menu("/sdcard/clockworkmod/backup/", NULL, headers);
    if (file == NULL)
        return;
    do_nandroid_restore(file);
}

void do_mount_usb_storage()
{
    system("echo /dev/block/mmcblk0 > /sys/devices/platform/usb_mass_storage/lun0/file");
    static char* headers[] = {  "USB Mass Storage device",
                                "Leaving this menu unmount",
                                "your SD card from your PC.",
                                "",
                                NULL 
    };
    
    static char* list[] = { "Unmount", NULL };
    
    for (;;)
    {
        int chosen_item = get_menu_selection(headers, list, 0);
        if (chosen_item == GO_BACK || chosen_item == 0)
            break;
    }
    
    system("echo '' > /sys/devices/platform/usb_mass_storage/lun0/file");
    system("echo 0 > /sys/devices/platform/usb_mass_storage/lun0/enable");
}

#define EXTENDEDCOMMAND_SCRIPT "/cache/recovery/extendedcommand"

int extendedcommand_file_exists()
{
    struct stat file_info;
    return 0 == stat(EXTENDEDCOMMAND_SCRIPT, &file_info);
}

int run_script_from_buffer(char* script_data, int script_len, char* filename)
{
    /* Parse the script.  Note that the script and parse tree are never freed.
     */
    const AmCommandList *commands = parseAmendScript(script_data, script_len);
    if (commands == NULL) {
        printf("Syntax error in update script\n");
        return 1;
    } else {
        printf("Parsed %.*s\n", script_len, filename);
    }

    /* Execute the script.
     */
    int ret = execCommandList((ExecContext *)1, commands);
    if (ret != 0) {
        int num = ret;
        char *line, *next = script_data;
        while (next != NULL && ret-- > 0) {
            line = next;
            next = memchr(line, '\n', script_data + script_len - line);
            if (next != NULL) *next++ = '\0';
        }
        printf("Failure at line %d:\n%s\n", num, next ? line : "(not found)");
        return 1;
    }    
    
    return 0;
}

int run_script(char* filename, int delete_file)
{
    struct stat file_info;
    if (0 != stat(filename, &file_info)) {
        printf("Error executing stat on file: %s\n", filename);
        return 1;
    }
    
    int script_len = file_info.st_size;
    char* script_data = (char*)malloc(script_len);
    FILE *file = fopen(filename, "rb");
    fread(script_data, script_len, 1, file);
    fclose(file);
	if (delete_file)
    	remove(filename);

	return run_script_from_buffer(script_data, script_len, filename);
}

int run_and_remove_extendedcommand()
{
    int i = 0;
    for (i = 20; i > 0; i--) {
        ui_print("Waiting for SD Card to mount (%ds)\n", i);
        if (ensure_root_path_mounted("SDCARD:") == 0) {
            ui_print("SD Card mounted...\n");
            break;
        }
        sleep(1);
    }
    if (i == 0) {
        ui_print("Timed out waiting for SD card... continuing anyways.");
    }
    
    return run_script(EXTENDEDCOMMAND_SCRIPT, 1);
}

int amend_main(int argc, char** argv)
{
    if (argc != 2) 
    {
        printf("Usage: amend <script>\n");
        return 0;
    }

    RecoveryCommandContext ctx = { NULL };
    if (register_update_commands(&ctx)) {
        LOGE("Can't install update commands\n");
    }
    return run_script(argv[1], 0);
}