#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <readline/history.h>
#include <readline/readline.h>
#include <sys/param.h>
#include <CoreFoundation/CoreFoundation.h>
#include <MobileDevice/MobileDevice.h>

char path_to_iphone_files[PATH_MAX];
struct am_recovery_device *recovery_device;

void recovery_progress(unsigned int progress_number, unsigned int opcode)
{
    if (opcode == 9) {
        _sendCommandToDevice(recovery_device, CFSTR(
            "setenv boot-args rd=md0 -v"));
        fprintf(stderr, "!");
    } else
        fprintf(stderr, ".");
}

void recovery_connect(struct am_recovery_device *rdev)
{
    CFMutableDictionaryRef opts;
    mach_error_t err;

    fprintf(stderr, "Connected to recovery mode, switching to restore");

    opts = AMRestoreCreateDefaultOptions(kCFAllocatorDefault);
    CFDictionarySetValue(opts, CFSTR("CreateFilesystemPartitions"),
        kCFBooleanFalse);
    CFDictionarySetValue(opts, CFSTR("RestoreBundlePath"),
        CFStringCreateWithCString(kCFAllocatorDefault, path_to_iphone_files,
        kCFStringEncodingUTF8));
    CFDictionarySetValue(opts, CFSTR("RestoreBootArgs"), CFSTR("rd=md0 -v"));

    recovery_device = rdev;

    err = AMRestorePerformRecoveryModeRestore(rdev, opts, recovery_progress,
        NULL);
    if (err) {
        fprintf(stderr, "... failed.\n");
        exit(1);
    }

    fprintf(stderr, "... done.\n");
}

void recovery_disconnect(struct am_recovery_device *rdev)
{
    fprintf(stderr, "Disconnected from recovery mode.\n");
}

void restore_main(struct am_restore_device *rdev)
{
    CFMutableDictionaryRef req;
    char *line;
    char *tok;
    mach_error_t err;

    while (1) {
        line = readline("(restored) ");
        if (!line)
            return;
        if (line[0] == '\0')
            continue;
        add_history(line);

        req = CFDictionaryCreateMutable(kCFAllocatorDefault, 0,
            &kCFTypeDictionaryKeyCallBacks,
            &kCFTypeDictionaryValueCallBacks);

        tok = strsep(&line, " ");
        if (!strcmp(tok, "Mount")) {
            CFDictionarySetValue(req, CFSTR("Operation"), CFSTR("Mount"));
            if (!(tok = strsep(&line, " ")))
                continue;
            CFDictionarySetValue(req, CFSTR("DeviceName"),
                CFStringCreateWithCString(kCFAllocatorDefault, tok,
                kCFStringEncodingUTF8));
            if (!(tok = strsep(&line, " ")))
                continue;
            CFDictionarySetValue(req, CFSTR("MountPoint"),
                CFStringCreateWithCString(kCFAllocatorDefault, tok,
                kCFStringEncodingUTF8));
        } else if (!strcmp(tok, "Partition")) {
            CFDictionarySetValue(req, CFSTR("Operation"), CFSTR("Partition"));
        } else if (!strcmp(tok, "Erase")) {
            CFDictionarySetValue(req, CFSTR("Operation"), CFSTR("Erase"));
            if (!(tok = strsep(&line, " ")))
                continue;
            CFDictionarySetValue(req, CFSTR("DeviceName"),
                CFStringCreateWithCString(kCFAllocatorDefault, tok,
                kCFStringEncodingUTF8));
            if (!(tok = strsep(&line, " ")))
                continue;
            CFDictionarySetValue(req, CFSTR("VolumeName"),
                CFStringCreateWithCString(kCFAllocatorDefault, tok,
                kCFStringEncodingUTF8));
        } else if (!strcmp(tok, "Ditto")) {
            CFDictionarySetValue(req, CFSTR("Operation"), CFSTR("Ditto"));
            if (!(tok = strsep(&line, " ")))
                continue;
            CFDictionarySetValue(req, CFSTR("SourcePath"),
                CFStringCreateWithCString(kCFAllocatorDefault, tok,
                kCFStringEncodingUTF8));
            if (!(tok = strsep(&line, " ")))
                continue;
            CFDictionarySetValue(req, CFSTR("DestinationPath"),
                CFStringCreateWithCString(kCFAllocatorDefault, tok,
                kCFStringEncodingUTF8));
        } else if (!strcmp(tok, "Unmount")) {
            CFDictionarySetValue(req, CFSTR("Operation"), CFSTR("Unmount"));
            if (!(tok = strsep(&line, " ")))
                continue;
            CFDictionarySetValue(req, CFSTR("MountPoint"),
                CFStringCreateWithCString(kCFAllocatorDefault, tok,
                kCFStringEncodingUTF8));
        } else if (!strcmp(tok, "FilesystemCheck")) {
            CFDictionarySetValue(req, CFSTR("Operation"),
                CFSTR("FilesystemCheck"));
            if (!(tok = strsep(&line, " ")))
                continue;
            CFDictionarySetValue(req, CFSTR("DeviceName"),
                CFStringCreateWithCString(kCFAllocatorDefault, tok,
                kCFStringEncodingUTF8));
        } else {
            fprintf(stderr, "Unknown command '%s'.\n", tok);
            CFRelease(req);
            continue;
        }

        err = _performOperation(rdev, req);

        fprintf(stderr, "ok, result = %d\n", (int)err); 
    }
}

void restore_connect(struct am_device_notification_callback_info *info)
{
    CFDictionaryRef req, resp;
    CFTypeRef *keys, *values;
    struct am_restore_device *rdev;

    rdev = AMRestoreModeDeviceCreate(0, AMDeviceGetConnectionID(info->dev), 0);
    rdev->port = _socketForPort(rdev, 0xf27e);

    keys = (CFTypeRef *)malloc(sizeof(CFTypeRef) * 1);
    values = (CFTypeRef *)malloc(sizeof(CFTypeRef) * 1);
    keys[0] = CFSTR("Request");
    values[0] = CFSTR("QueryType");

    req = CFDictionaryCreate(kCFAllocatorDefault, (const void **)keys,
        (const void **)values, 1, &kCFTypeDictionaryKeyCallBacks,
        &kCFTypeDictionaryValueCallBacks);
    restored_send_message(rdev->port, req);
    resp = restored_receive_message(rdev->port);

    free(keys);
    free(values);

    fprintf(stderr, "Requested query type:\n");
    CFShow(resp);

    restore_main(rdev);
}

int main(int argc, char **argv)
{
    struct am_device_notification *notif;

    if (argc < 2 || !realpath(argv[1], path_to_iphone_files)) {
        fprintf(stderr, "usage: %s path-to-iphone-files\n", argv[0]);
        return 0;
    }

    AMRestoreRegisterForDeviceNotifications(NULL, recovery_connect, NULL,
        recovery_disconnect, 0, NULL);
    AMDeviceNotificationSubscribe(restore_connect, 0, 0, 0, &notif);

    fprintf(stderr, "Please connect an iPhone non-destructively put into "
        "recovery mode...\n");

    CFRunLoopRun();

    return 0;
}