#!/usr/bin/env racket
#lang racket/base

(require racket/cmdline
         racket/file
         racket/list
         racket/match
         racket/path
         racket/port
         racket/string
         racket/system)


(define version "0.4.0")

(struct disk (node whole? internal? external? removable? read-only? virtual? size-bytes name) #:transparent)
(struct autoinstall (seed-dir seed-volume standalone-dir hostname username password-hash ssh-key locale keyboard timezone) #:transparent)
(struct plan (image disk image-bytes remaining-bytes mode autoinstall) #:transparent)

(define (usage)
  (displayln "Starflash - cautious macOS SD/USB image flasher")
  (displayln "")
  (displayln "Usage:")
  (displayln "  ./starflash --guide")
  (displayln "  ./starflash --list-disks")
  (displayln "  ./starflash --image IMAGE --disk /dev/diskN --dry-run")
  (displayln "  ./starflash --image IMAGE --disk /dev/diskN --write --yes")
  (displayln "  ./starflash --ubuntu-autoinstall --seed-dir cidata --render-seed")
  (displayln "  ./starflash --ubuntu-autoinstall --seed-volume /Volumes/CIDATA --render-seed")
  (displayln "  ./starflash --ubuntu-autoinstall --iso-tree /Volumes/Ubuntu --standalone-dir ubuntu-auto --render-standalone")
  (displayln "  ./starflash --ubuntu-autoinstall --image ubuntu.iso --disk /dev/diskN --seed-dir cidata --dry-run")
  (displayln "")
  (displayln "Options:")
  (displayln "  --guide            Ask questions and build a safe plan interactively")
  (displayln "  --guide-script PATH  Read guide answers from a file for testing")
  (displayln "  --list-disks       Show disks using diskutil list")
  (displayln "  --image PATH       Image to flash (.img, .iso, .raw, or .img.gz)")
  (displayln "  --disk /dev/diskN  Whole target disk")
  (displayln "  --dry-run          Validate and print the plan without writing")
  (displayln "  --write            Actually erase/write the disk")
  (displayln "  --yes              Required with --write")
  (displayln "  --ubuntu-autoinstall  Render Ubuntu Server autoinstall NoCloud seed files")
  (displayln "  --seed-dir PATH    Seed output directory, default ./cidata")
  (displayln "  --seed-volume PATH Mounted CIDATA volume to receive seed files")
  (displayln "  --render-seed      Write user-data/meta-data into --seed-dir")
  (displayln "  --iso-tree PATH    Mounted or copied Ubuntu ISO tree")
  (displayln "  --standalone-dir PATH  Output directory for one-USB autoinstall tree")
  (displayln "  --render-standalone    Embed seed and patch boot config in --standalone-dir")
  (displayln "  --hostname NAME    Autoinstall hostname, default ubuntu")
  (displayln "  --username NAME    Autoinstall user, default ubuntu")
  (displayln "  --password-hash HASH  Autoinstall password hash, default locked password")
  (displayln "  --ssh-key KEY      SSH authorized key for the autoinstall user")
  (displayln "  --locale LOCALE    Ubuntu locale, default en_US.UTF-8")
  (displayln "  --keyboard LAYOUT  Ubuntu keyboard layout, default us")
  (displayln "  --timezone TZ      Ubuntu timezone, default UTC")
  (displayln "  --version          Print version")
  (displayln "  --help             Show this help"))

(define (fail fmt . args)
  (apply eprintf (string-append "starflash: " fmt "\n") args)
  (exit 1))

(define (blank? value)
  (string=? (string-trim value) ""))

(define (yes-answer? value)
  (member (string-downcase (string-trim value)) '("y" "yes")))

(define (no-answer? value)
  (member (string-downcase (string-trim value)) '("n" "no")))

(define (run/text . args)
  (define out (open-output-string))
  (define err (open-output-string))
  (define ok?
    (parameterize ([current-output-port out]
                   [current-error-port err])
      (apply system* args)))
  (unless ok?
    (fail "command failed: ~a\n~a" (string-join args " ") (string-trim (get-output-string err))))
  (get-output-string out))

(define (diskutil . args)
  (apply run/text (append (list "/usr/sbin/diskutil") args)))

(define (info-text node)
  (or (getenv "STARFLASH_DISKUTIL_INFO")
      (diskutil "info" node)))

(define (field-line text label)
  (define prefix (format "~a:" label))
  (for/first ([line (in-list (string-split text "\n"))]
              #:do [(define trimmed (string-trim line))]
              #:when (string-prefix? trimmed prefix))
    (string-trim (substring trimmed (string-length prefix)))))

(define (bool-field text label)
  (match (field-line text label)
    ["Yes" #t]
    ["No" #f]
    [_ #f]))

(define (string-field text label)
  (field-line text label))

(define (bytes-field text label)
  (match (regexp-match #px"\\(([0-9]+) Bytes\\)" (or (field-line text label) ""))
    [(list _ value) (string->number value)]
    [_ #f]))

(define (read-disk node)
  (define text (info-text node))
  (disk (or (string-field text "Device Node") node)
        (bool-field text "Whole")
        (bool-field text "Internal")
        (equal? (string-field text "Device Location") "External")
        (equal? (string-field text "Removable Media") "Removable")
        (bool-field text "Media Read-Only")
        (bool-field text "Virtual")
        (bytes-field text "Disk Size")
        (or (string-field text "Device / Media Name") "unknown")))

(define (allowed-image? path)
  (for/or ([suffix (in-list '(".img" ".iso" ".raw" ".img.gz"))])
    (string-suffix? (path->string path) suffix)))

(define (plain-image? path)
  (not (string-suffix? (path->string path) ".gz")))

(define (validate-image path)
  (unless path
    (fail "missing --image"))
  (unless (file-exists? path)
    (fail "image does not exist: ~a" path))
  (unless (allowed-image? path)
    (fail "unsupported image type: ~a" path))
  (file-size path))

(define (validate-disk target)
  (unless target
    (fail "missing --disk"))
  (unless (regexp-match? #px"^/dev/disk[0-9]+$" target)
    (fail "target must be a whole macOS disk like /dev/disk4, got: ~a" target))
  (define d (read-disk target))
  (unless (disk-whole? d)
    (fail "target is not a whole disk: ~a" target))
  (when (disk-internal? d)
    (fail "refusing internal disk: ~a" target))
  (unless (disk-external? d)
    (fail "refusing non-external disk: ~a" target))
  (when (disk-read-only? d)
    (fail "target is read-only: ~a" target))
  (when (disk-virtual? d)
    (fail "refusing virtual disk: ~a" target))
  (unless (disk-size-bytes d)
    (fail "could not determine disk size for: ~a" target))
  d)

(define (make-plan image target mode autoinstall)
  (define image-bytes (validate-image image))
  (define d (validate-disk target))
  (when (and (plain-image? image)
             (> image-bytes (disk-size-bytes d)))
    (fail "image is larger than target: ~a bytes > ~a bytes"
          image-bytes
          (disk-size-bytes d)))
  (plan image d image-bytes (- (disk-size-bytes d) image-bytes) mode autoinstall))

(define (default-password-hash)
  "!")

(define (make-autoinstall seed-dir seed-volume standalone-dir hostname username password-hash ssh-key locale keyboard timezone)
  (autoinstall (or seed-dir (build-path (current-directory) "cidata"))
               seed-volume
               standalone-dir
               hostname
               username
               password-hash
               ssh-key
               locale
               keyboard
               timezone))

(define (yaml-quote value)
  (format "\"~a\"" (string-replace value "\"" "\\\"")))

(define (seed-ssh-lines cfg)
  (if (and (autoinstall-ssh-key cfg)
           (not (string=? (autoinstall-ssh-key cfg) "")))
      (string-append
       "        authorized-keys:\n"
       "          - " (yaml-quote (autoinstall-ssh-key cfg)) "\n")
      "        authorized-keys: []\n"))

(define (render-user-data cfg)
  (string-append
   "#cloud-config\n"
   "autoinstall:\n"
   "  version: 1\n"
   (format "  locale: ~a\n" (autoinstall-locale cfg))
   "  keyboard:\n"
   (format "    layout: ~a\n" (autoinstall-keyboard cfg))
   (format "  timezone: ~a\n" (autoinstall-timezone cfg))
   "  identity:\n"
   (format "    hostname: ~a\n" (yaml-quote (autoinstall-hostname cfg)))
   (format "    username: ~a\n" (yaml-quote (autoinstall-username cfg)))
   (format "    password: ~a\n" (yaml-quote (autoinstall-password-hash cfg)))
   "  ssh:\n"
   "    install-server: true\n"
   "    allow-pw: false\n"
   (seed-ssh-lines cfg)
   "  storage:\n"
   "    layout:\n"
   "      name: direct\n"
   "  updates: security\n"
   "  late-commands:\n"
   "    - curtin in-target --target=/target -- sh -lc 'date -u +%Y-%m-%dT%H:%M:%SZ > /etc/mayphus-autoinstall'\n"))

(define (render-meta-data cfg)
  (format "instance-id: ~a-autoinstall\nlocal-hostname: ~a\n"
          (autoinstall-hostname cfg)
          (autoinstall-hostname cfg)))

(define (render-seed-readme cfg)
  (string-append
   "Ubuntu autoinstall NoCloud seed\n\n"
   "Copy user-data and meta-data to a small FAT/exFAT volume labeled CIDATA.\n"
   "Boot the Ubuntu Server installer with this seed attached. The installer can\n"
   "then run without the language/user/password screens.\n\n"
   "Seed directory: " (path->string (simplify-path (autoinstall-seed-dir cfg))) "\n"))

(define (render-seed-files! cfg dest)
  (make-directory* dest)
  (call-with-output-file (build-path dest "user-data")
    (lambda (out) (display (render-user-data cfg) out))
    #:exists 'truncate/replace)
  (call-with-output-file (build-path dest "meta-data")
    (lambda (out) (display (render-meta-data cfg) out))
    #:exists 'truncate/replace)
  (call-with-output-file (build-path dest "README.txt")
    (lambda (out) (display (render-seed-readme cfg) out))
    #:exists 'truncate/replace))

(define (render-seed! cfg)
  (render-seed-files! cfg (autoinstall-seed-dir cfg))
  (printf "Rendered Ubuntu autoinstall seed: ~a\n" (path->string (simplify-path (autoinstall-seed-dir cfg))))
  (when (autoinstall-seed-volume cfg)
    (unless (directory-exists? (autoinstall-seed-volume cfg))
      (fail "--seed-volume is not a mounted directory: ~a" (autoinstall-seed-volume cfg)))
    (render-seed-files! cfg (autoinstall-seed-volume cfg))
    (printf "Copied seed files to mounted CIDATA volume: ~a\n"
            (path->string (simplify-path (autoinstall-seed-volume cfg))))))

(define standalone-boot-args
  "autoinstall ds=nocloud\\;s=/cdrom/nocloud/")

(define (boot-config-file? path)
  (member (path->string (file-name-from-path path))
          '("grub.cfg" "loopback.cfg" "txt.cfg" "isolinux.cfg" "syslinux.cfg")))

(define (patch-boot-config text)
  (cond
    [(regexp-match? #rx"autoinstall[ ]+ds=nocloud" text) text]
    [(regexp-match? #rx" ---" text)
     (regexp-replace* #rx" ---" text (lambda (_) (format " ~a ---" standalone-boot-args)))]
    [else text]))

(define (all-files dir)
  (apply append
         (for/list ([path (in-list (directory-list dir #:build? #t))])
           (cond
             [(directory-exists? path) (all-files path)]
             [(file-exists? path) (list path)]
             [else '()]))))

(define (relative-to base path)
  (find-relative-path base path))

(define (patch-standalone-boot-configs! dir)
  (define patched
    (for/list ([path (in-list (all-files dir))]
               #:when (boot-config-file? path)
               #:do [(define before (file->string path))
                     (define after (patch-boot-config before))]
               #:when (not (string=? before after)))
      (call-with-output-file path
        (lambda (out) (display after out))
        #:exists 'truncate/replace)
      (relative-to dir path)))
  (unless (pair? patched)
    (fail "no boot config with a kernel '---' marker was patched in: ~a" dir))
  patched)

(define (render-standalone-readme cfg)
  (string-append
   "Standalone Ubuntu autoinstall tree\n\n"
   "This tree has /nocloud/user-data and /nocloud/meta-data embedded.\n"
   "Boot config files were patched with:\n"
   standalone-boot-args "\n\n"
   "Make this tree into a bootable ISO/USB image with an ISO remastering tool,\n"
   "then flash that image with Starflash. The installed machine should not need\n"
   "a second CIDATA volume, a MacBook, or Ethernet for installer answers.\n"))

(define (render-standalone-tree! cfg iso-tree standalone-dir)
  (unless iso-tree
    (fail "--render-standalone requires --iso-tree"))
  (unless (directory-exists? iso-tree)
    (fail "--iso-tree is not a directory: ~a" iso-tree))
  (unless standalone-dir
    (fail "--render-standalone requires --standalone-dir"))
  (when (directory-exists? standalone-dir)
    (fail "--standalone-dir already exists; move it first: ~a" standalone-dir))
  (copy-directory/files iso-tree standalone-dir)
  (define seed-target (build-path standalone-dir "nocloud"))
  (render-seed-files! cfg seed-target)
  (define patched (patch-standalone-boot-configs! standalone-dir))
  (call-with-output-file (build-path standalone-dir "STARFLASH-AUTOINSTALL.txt")
    (lambda (out) (display (render-standalone-readme cfg) out))
    #:exists 'truncate/replace)
  (printf "Rendered standalone Ubuntu autoinstall tree: ~a\n" (path->string (simplify-path standalone-dir)))
  (for ([path (in-list patched)])
    (printf "Patched boot config: ~a\n" (path->string path))))

(define (show-plan p)
  (define d (plan-disk p))
  (printf "Mode: ~a\n" (plan-mode p))
  (printf "Image: ~a (~a bytes~a)\n"
          (plan-image p)
          (plan-image-bytes p)
          (if (plain-image? (plan-image p)) "" ", compressed size"))
  (printf "Target: ~a\n" (disk-node d))
  (printf "Media: ~a, external, ~a, ~a bytes\n"
          (disk-name d)
          (if (disk-removable? d) "removable" "not-removable")
          (disk-size-bytes d))
  (when (plain-image? (plan-image p))
    (printf "Remaining after image: ~a bytes\n" (plan-remaining-bytes p)))
  (when (plan-autoinstall p)
    (printf "Ubuntu autoinstall seed: ~a\n" (path->string (simplify-path (autoinstall-seed-dir (plan-autoinstall p)))))
    (displayln "Seed handoff: NoCloud files on a volume labeled CIDATA. No live desktop setup screens."))
  (when (eq? (plan-mode p) 'dry-run)
    (displayln "Dry run only: no unmount, erase, write, sync, or eject will run.")))

(define (raw-disk node)
  (regexp-replace #px"^/dev/disk" node "/dev/rdisk"))

(define (dd-image image node)
  (system* "/bin/dd"
           (format "if=~a" (path->string image))
           (format "of=~a" node)
           "bs=4m"
           "conv=sync"))

(define (write-image p yes?)
  (unless yes?
    (fail "--write requires --yes"))
  (unless (plain-image? (plan-image p))
    (fail "writing compressed images is not implemented yet; decompress to .img first"))
  (define d (plan-disk p))
  (displayln "About to erase/write. This is destructive.")
  (show-plan p)
  (diskutil "unmountDisk" (disk-node d))
  (unless (dd-image (plan-image p) (raw-disk (disk-node d)))
    (displayln "Raw disk write failed; retrying buffered whole-disk path.")
    (unless (dd-image (plan-image p) (disk-node d))
      (fail "dd failed")))
  (diskutil "eject" (disk-node d))
  (displayln "Done."))

(define (prompt-line in prompt [default #f])
  (printf "~a~a: " prompt (if default (format " [~a]" default) ""))
  (flush-output)
  (define value (read-line in 'any))
  (when (eof-object? value)
    (fail "guide input ended at: ~a" prompt))
  (define trimmed (string-trim value))
  (cond
    [(and (blank? trimmed) default) default]
    [else trimmed]))

(define (prompt-required in prompt [default #f])
  (let loop ()
    (define value (prompt-line in prompt default))
    (if (blank? value)
        (begin
          (displayln "Please enter a value.")
          (loop))
        value)))

(define (prompt-yes/no in prompt [default-no? #t])
  (define suffix (if default-no? "y/N" "Y/n"))
  (let loop ()
    (define value (prompt-line in (format "~a (~a)" prompt suffix)))
    (cond
      [(blank? value) (not default-no?)]
      [(yes-answer? value) #t]
      [(no-answer? value) #f]
      [else
       (displayln "Please answer yes or no.")
       (loop)])))

(define (default-ssh-key-path)
  (build-path (find-system-path 'home-dir) ".ssh" "id_ed25519.pub"))

(define (read-ssh-key-value value)
  (define trimmed (string-trim value))
  (cond
    [(blank? trimmed) ""]
    [(or (string-prefix? trimmed "ssh-")
         (string-prefix? trimmed "ecdsa-"))
     trimmed]
    [(file-exists? (string->path trimmed))
     (string-trim (file->string (string->path trimmed)))]
    [else trimmed]))

(define (guided-autoinstall in)
  (and (prompt-yes/no in "Create Ubuntu cloud-init/autoinstall seed files?" #t)
       (let* ([seed-dir (string->path (prompt-line in "Seed output directory" "cidata"))]
              [hostname (prompt-line in "Hostname" "ubuntu")]
              [username (prompt-line in "Username" (or (getenv "USER") "ubuntu"))]
              [ssh-default (path->string (default-ssh-key-path))]
              [ssh-key (read-ssh-key-value (prompt-line in "SSH public key or key file path" ssh-default))]
              [locale (prompt-line in "Locale" "en_US.UTF-8")]
              [keyboard (prompt-line in "Keyboard layout" "us")]
              [timezone (prompt-line in "Timezone" "UTC")]
              [cfg (make-autoinstall seed-dir #f #f hostname username (default-password-hash) ssh-key locale keyboard timezone)])
         (render-seed! cfg)
         cfg)))

(define (run-guide #:script [script #f])
  (define in (if script (open-input-file script) (current-input-port)))
  (dynamic-wind
    void
    (lambda ()
      (displayln "Starflash guide")
      (displayln "This guide runs a dry-run first. It writes only if you type WRITE.")
      (when (prompt-yes/no in "Show disks now?" #f)
        (display (diskutil "list")))
      (define image-path (string->path (prompt-required in "Image path (.img, .iso, .raw)" #f)))
      (define target (prompt-required in "Target whole disk, for example /dev/disk4" #f))
      (define autoinstall-config (guided-autoinstall in))
      (define p (make-plan image-path target 'dry-run autoinstall-config))
      (show-plan p)
      (displayln "")
      (define write-answer (prompt-line in "Type WRITE to erase/write this disk, or press Enter to stop"))
      (if (string=? write-answer "WRITE")
          (write-image (make-plan image-path target 'write autoinstall-config) #t)
          (displayln "Stopped after dry-run. No write was performed.")))
    (lambda ()
      (when script (close-input-port in)))))

(define (main)
  (define cli-args (current-command-line-arguments))
  (when (member "--help" (vector->list cli-args))
    (usage)
    (exit 0))
  (define image #f)
  (define target #f)
  (define list-disks? #f)
  (define dry-run? #f)
  (define write? #f)
  (define yes? #f)
  (define guide? (zero? (vector-length cli-args)))
  (define guide-script #f)
  (define ubuntu-autoinstall? #f)
  (define render-seed? #f)
  (define render-standalone? #f)
  (define seed-dir #f)
  (define seed-volume #f)
  (define iso-tree #f)
  (define standalone-dir #f)
  (define hostname "ubuntu")
  (define username "ubuntu")
  (define password-hash (default-password-hash))
  (define ssh-key "")
  (define locale "en_US.UTF-8")
  (define keyboard "us")
  (define timezone "UTC")
  (command-line
   #:program "starflash"
   #:once-each
   [("--version") "Show version" (displayln version) (exit 0)]
   [("--guide") "Interactive guide" (set! guide? #t)]
   [("--guide-script") path "Read guide answers from file" (set! guide? #t) (set! guide-script (string->path path))]
   [("--list-disks") "List disks" (set! list-disks? #t)]
   [("--dry-run") "Validate without writing" (set! dry-run? #t)]
   [("--write") "Actually write image" (set! write? #t)]
   [("--yes") "Confirm destructive write" (set! yes? #t)]
   [("--image") path "Image path" (set! image (string->path path))]
   [("--disk") node "Target /dev/diskN" (set! target node)]
   [("--ubuntu-autoinstall") "Enable Ubuntu autoinstall seed planning" (set! ubuntu-autoinstall? #t)]
   [("--render-seed") "Render Ubuntu autoinstall seed files" (set! render-seed? #t)]
   [("--seed-dir") path "Seed output directory" (set! seed-dir (string->path path))]
   [("--seed-volume") path "Mounted CIDATA volume" (set! seed-volume (string->path path))]
   [("--iso-tree") path "Mounted or copied Ubuntu ISO tree" (set! iso-tree (string->path path))]
   [("--standalone-dir") path "Standalone autoinstall tree output directory" (set! standalone-dir (string->path path))]
   [("--render-standalone") "Render standalone one-USB autoinstall tree" (set! render-standalone? #t)]
   [("--hostname") value "Autoinstall hostname" (set! hostname value)]
   [("--username") value "Autoinstall username" (set! username value)]
   [("--password-hash") value "Autoinstall password hash" (set! password-hash value)]
   [("--ssh-key") value "SSH authorized key" (set! ssh-key value)]
   [("--locale") value "Ubuntu locale" (set! locale value)]
   [("--keyboard") value "Ubuntu keyboard layout" (set! keyboard value)]
   [("--timezone") value "Ubuntu timezone" (set! timezone value)])
  (define autoinstall-config
    (and ubuntu-autoinstall?
         (make-autoinstall seed-dir seed-volume standalone-dir hostname username password-hash ssh-key locale keyboard timezone)))
  (when guide?
    (run-guide #:script guide-script)
    (exit 0))
  (when render-seed?
    (unless autoinstall-config
      (fail "--render-seed requires --ubuntu-autoinstall"))
    (render-seed! autoinstall-config))
  (when render-standalone?
    (unless autoinstall-config
      (fail "--render-standalone requires --ubuntu-autoinstall"))
    (render-standalone-tree! autoinstall-config iso-tree standalone-dir))
  (cond
    [list-disks?
     (display (diskutil "list"))]
    [write?
     (write-image (make-plan image target 'write autoinstall-config) yes?)]
    [(and (or render-seed? render-standalone?) (not image) (not target))
     (void)]
    [else
     (set! dry-run? #t)
     (show-plan (make-plan image target 'dry-run autoinstall-config))]))

(module+ main
  (main))