#import <AppKit/AppKit.h>
#import <CommonCrypto/CommonDigest.h>
#include <errno.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <unistd.h>

// Generated by build.py from this invocation's paths and payload hashes.
#include "writer-config.h"

static NSDictionary *Plan;

static BOOL CheckHash(NSString *name, NSString *expected) {
    NSData *data = [NSData dataWithContentsOfFile:[Base stringByAppendingPathComponent:name]];
    if (!data || data.length > UINT32_MAX) return NO;
    unsigned char digest[CC_SHA256_DIGEST_LENGTH];
    CC_SHA256(data.bytes, (CC_LONG)data.length, digest);
    NSMutableString *actual = [NSMutableString string];
    for (unsigned i = 0; i < sizeof(digest); ++i) [actual appendFormat:@"%02x", digest[i]];
    return [actual isEqualToString:expected];
}

static BOOL CheckConfiguration(void) {
    if (!CheckHash(@"write_usb.py", WriterHash)
        || !CheckHash(@"usb_device.py", DeviceHelperHash)
        || !CheckHash(@"write-plan.json", PlanHash)) {
        fprintf(stderr, "The prepared writer or disk plan changed. Refusing to write.\n");
        return NO;
    }
    NSData *data = [NSData dataWithContentsOfFile:[Base stringByAppendingPathComponent:@"write-plan.json"]];
    id plan = [NSJSONSerialization JSONObjectWithData:data options:0 error:NULL];
    if (![plan isKindOfClass:NSDictionary.class]) return NO;
    id device = plan[@"device"];
    NSRegularExpression *pattern = [NSRegularExpression regularExpressionWithPattern:@"^disk[0-9]+$" options:0 error:NULL];
    if (![device isKindOfClass:NSString.class]
        || [pattern numberOfMatchesInString:device options:0 range:NSMakeRange(0, [device length])] != 1
        || ![plan[@"registry_name"] isKindOfClass:NSString.class]
        || ![plan[@"size"] isKindOfClass:NSNumber.class] || [plan[@"size"] unsignedLongLongValue] == 0
        || ![plan[@"media_entry_id"] isKindOfClass:NSNumber.class]
        || ![plan[@"boot_session_uuid"] isKindOfClass:NSString.class]) {
        fprintf(stderr, "The disk plan is missing an explicit target attachment. Prepare a new plan.\n");
        return NO;
    }
    if (![Python isAbsolutePath] || access(Python.fileSystemRepresentation, X_OK) != 0) {
        fprintf(stderr, "The Python used to prepare this writer is no longer available. Rebuild the writer.\n");
        return NO;
    }
    Plan = plan;
    return YES;
}

static int RunPython(BOOL write, int diskFD) {
    char descriptor[32];
    snprintf(descriptor, sizeof(descriptor), "%d", diskFD);
    const char *script = [[Base stringByAppendingPathComponent:@"write_usb.py"] fileSystemRepresentation];
    pid_t child = fork();
    if (child < 0) {
        perror("fork");
        if (diskFD >= 0) close(diskFD);
        return 1;
    }
    if (child == 0) {
        // Ignore Python environment overrides and the user's site packages.
        char *const args[] = {(char *)Python.fileSystemRepresentation, "-E", "-s", "-B",
            (char *)script, write ? "--write" : "--check-device",
            write ? "--device-fd" : NULL, write ? descriptor : NULL, NULL};
        execv(args[0], args);
        perror("Starting the reviewed writer");
        _exit(1);
    }
    // The child owns raw access from here. Keeping the parent's duplicate
    // open would prevent diskutil from ejecting after read-back verification.
    if (diskFD >= 0) close(diskFD);
    int status = 0;
    while (waitpid(child, &status, 0) < 0) {
        if (errno == EINTR) continue;
        perror("waitpid");
        return 1;
    }
    return WIFEXITED(status) ? WEXITSTATUS(status) : 1;
}

static int AuthorizedWrite(void) {
    if (geteuid() != 0) {
        fprintf(stderr, "Open the prepared Makepad USB Writer app to authenticate as an administrator.\n");
        return 1;
    }
    if (!CheckConfiguration() || RunPython(NO, -1) != 0) return 1;

    // Raw access belongs to this app's Full Disk Access grant. Python receives
    // the descriptor and rechecks its identity and attachment before writing.
    NSString *rawDevice = [NSString stringWithFormat:@"/dev/r%@", Plan[@"device"]];
    int diskFD = open(rawDevice.fileSystemRepresentation, O_RDWR);
    if (diskFD < 0) {
        perror("Opening the selected USB raw device");
        fprintf(stderr, "Enable the prepared writer app in Full Disk Access, then reopen it. No bytes were written.\n");
        return 1;
    }
    struct stat opened;
    if (fstat(diskFD, &opened) != 0 || !S_ISCHR(opened.st_mode)) {
        fprintf(stderr, "The target is not a raw character device. Refusing to write.\n");
        close(diskFD);
        return 1;
    }
    return RunPython(YES, diskFD);
}

static NSString *ShellQuote(NSString *value) {
    return [NSString stringWithFormat:@"'%@'", [value stringByReplacingOccurrencesOfString:@"'" withString:@"'\\''"]];
}

static NSString *AppleScriptQuote(NSString *value) {
    NSString *escaped = [value stringByReplacingOccurrencesOfString:@"\\" withString:@"\\\\"];
    escaped = [escaped stringByReplacingOccurrencesOfString:@"\"" withString:@"\\\""];
    return [NSString stringWithFormat:@"\"%@\"", escaped];
}

@interface WriterDelegate : NSObject <NSApplicationDelegate>
@property(strong) NSWindow *window;
@end

@implementation WriterDelegate
- (BOOL)applicationShouldTerminateAfterLastWindowClosed:(NSApplication *)sender {
    (void)sender;
    return YES;
}
- (void)applicationDidFinishLaunching:(NSNotification *)notification {
    (void)notification;
    if (!CheckConfiguration()) {
        NSAlert *alert = [[NSAlert alloc] init];
        alert.messageText = @"The prepared USB writer is incomplete or has changed";
        alert.informativeText = @"Build a new writer from the current image output before writing a disk.";
        [alert runModal];
        [NSApp terminate:nil];
        return;
    }
    self.window = [[NSWindow alloc] initWithContentRect:NSMakeRect(0, 0, 560, 220)
        styleMask:NSWindowStyleMaskTitled | NSWindowStyleMaskClosable
        backing:NSBackingStoreBuffered defer:NO];
    self.window.title = @"Makepad USB Writer";
    NSString *text = [NSString stringWithFormat:
        @"Target: %@ — %@ (%.1f GB)\n\nAll data on this USB will be replaced with Arch Linux. The app verifies the selected drive, writes the prepared image, reads it back, and ejects it. Keep the drive connected until completion.",
        Plan[@"device"], Plan[@"registry_name"], [Plan[@"size"] doubleValue] / 1e9];
    NSTextField *label = [NSTextField wrappingLabelWithString:text];
    label.frame = NSMakeRect(22, 65, 516, 135);
    [self.window.contentView addSubview:label];
    NSButton *button = [NSButton buttonWithTitle:@"Erase and write" target:self action:@selector(writeImage:)];
    button.frame = NSMakeRect(375, 18, 160, 32);
    [self.window.contentView addSubview:button];
    [self.window center];
    [self.window orderFront:nil];
}
- (void)writeImage:(id)sender {
    [sender setEnabled:NO];
    NSString *command = [NSString stringWithFormat:@"%@ --authorized-write > %@ 2>&1",
        ShellQuote(NSBundle.mainBundle.executablePath),
        ShellQuote([Base stringByAppendingPathComponent:@"writer-app.log"])];
    NSAppleScript *script = [[NSAppleScript alloc] initWithSource:
        [NSString stringWithFormat:@"do shell script %@ with administrator privileges", AppleScriptQuote(command)]];
    NSDictionary *error = nil;
    NSAppleEventDescriptor *result = [script executeAndReturnError:&error];
    [self.window orderOut:nil];
    NSAlert *alert = [[NSAlert alloc] init];
    if (result && !error) {
        alert.messageText = @"Arch USB verified and ejected";
        alert.informativeText = @"Move the USB to the target PC, disable Secure Boot, and choose its UEFI USB boot entry. Connect Ethernet. Login: arch. A memorable SSH and sudo password is generated on first boot and displayed on the local console. Run sudo makepad-ssh show for connection details. Initial setup continues automatically; run makepad-status to check progress.";
    } else {
        alert.alertStyle = NSAlertStyleWarning;
        alert.messageText = @"Arch USB write did not complete";
        alert.informativeText = [NSString stringWithFormat:@"%@\n\nDetails: %@/writer-app.log\nIf raw access was denied, enable this writer app in Full Disk Access and reopen it.",
            error[NSAppleScriptErrorMessage] ?: @"The administrator action did not complete.", Base];
    }
    [alert addButtonWithTitle:@"Close"];
    [alert runModal];
    [NSApp terminate:nil];
}
@end

int main(int argc, const char *argv[]) {
    @autoreleasepool {
        if (argc == 2 && strcmp(argv[1], "--authorized-write") == 0) return AuthorizedWrite();
        if (argc == 2 && strcmp(argv[1], "--check-configuration") == 0) {
            if (!CheckConfiguration()) return 1;
            puts("Dedicated writer: pinned scripts and disk plan hashes match. No device opened.");
            return 0;
        }
        if (argc != 1) return 2;
        [NSApplication sharedApplication];
        [NSApp setActivationPolicy:NSApplicationActivationPolicyRegular];
        WriterDelegate *delegate = [[WriterDelegate alloc] init];
        NSApp.delegate = delegate;
        [NSApp run];
    }
    return 0;
}