(import (rename scheme
                (display display*)
                (integer->char integer->char*)
                (char->integer char->integer*)
                (list->string list->string*)
                (string->list string->list*)))
(import (chicken base))
(import (chicken bitwise))
(import (chicken condition))
(import (chicken file))
(import (chicken format))
(import (chicken io))
(import (chicken irregex))
(import (chicken pathname))
(import (chicken platform))
(import (chicken process-context))
(import (chicken string))
(import (srfi 1))
(import (srfi 69))
(import getopt-long)
(import utf8)

;;; common

(define (bytes->string bytes) (list->string* (map integer->char* bytes)))
(define (string->bytes string) (map char->integer* (string->list* string)))

(define (strip-key-prefix key-name)
  (string-translate* key-name '(("KEY_" . ""))))

(define (fill-keymap sexps)
  (let ((keymap (make-hash-table eqv? eqv?-hash)))
    (for-each
     (lambda (sexp)
       (hash-table-set! keymap (car sexp) (cdr sexp)))
     sexps)
    keymap))

(define (built-in-path name)
  (make-pathname (chicken-home) (make-pathname "duckyscript" name)))

(define built-in-layout-directory (built-in-path "layouts"))
(define built-in-templates-directory (built-in-path "templates"))

(define built-in-key-bytes-file (built-in-path "usb-hid-keys"))
(define built-in-modifier-bytes-file (built-in-path "usb-hid-modifiers"))
(define built-in-key-aliases-file (built-in-path "key-aliases"))
(define built-in-modifier-aliases-file (built-in-path "modifier-aliases"))

(define (built-in-layouts) (directory built-in-layout-directory))
(define (built-in-templates) (directory built-in-templates-directory))

(define (read-all in)
  (read-string #f in))

(define (try-read-file path reader)
  (condition-case
   (call-with-input-file path reader)
   (e (exn i/o file) #f)))

(define (try-read-layout name-or-path)
  (or (try-read-file (make-pathname built-in-layout-directory name-or-path)
                     read-list)
      (try-read-file name-or-path read-list)
      (abort
       (condition '(exn location try-read-layout message "Layout not found")
                  '(i/o)
                  '(file)))))

(define (try-read-template name-or-path)
  (or (try-read-file (make-pathname built-in-templates-directory name-or-path)
                     read-all)
      (try-read-file name-or-path read-all)
      (abort
       (condition '(exn location try-read-template message "Template not found")
                  '(i/o)
                  '(file)))))

(define (built-in-file name)
  (alist->hash-table (call-with-input-file name read-list)))

(define (inverse-hash-table hash-table)
  (let ((ht (make-hash-table)))
    (hash-table-for-each hash-table
      (lambda (key value)
        (hash-table-set! ht value key)))
    ht))

(define modifier-aliases (built-in-file built-in-modifier-aliases-file))
(define modifier-aliases* (inverse-hash-table modifier-aliases))
(define key-aliases (built-in-file built-in-key-aliases-file))
(define key-bytes (built-in-file built-in-key-bytes-file))
(define modifier-bytes (built-in-file built-in-modifier-bytes-file))
(define key-bytes* (inverse-hash-table key-bytes))
(define modifier-bytes* (inverse-hash-table modifier-bytes))

(define (dump-aliases-symbols)
  (for-each
   (lambda (label hash-table)
     (print label)
     (hash-table-for-each
      hash-table
      (lambda (key alias)
        (printf " ~a (~a)\n" key alias))))
   (list "Modifier aliases:" "Key aliases:")
   (list modifier-aliases key-aliases))
  (print "Key symbols:")
  (hash-table-for-each key-bytes
    (lambda (key _alias)
      (printf " ~a\n" (strip-key-prefix key))))
  (print "Modifier symbols:")
  (hash-table-for-each modifier-bytes
    (lambda (modifier _alias)
      (printf " ~a\n" (strip-key-prefix modifier)))))

;;; compiler

(define dashed-combination-irx
  '(: bol (or "CTRL-ALT" "CTRL-SHIFT"
              "ALT-SHIFT" "ALT-TAB"
              "COMMAND-OPTION")))

(define (replace-dashed-combinations line)
  (irregex-replace dashed-combination-irx line
                   (lambda (match)
                     (let ((dashed (irregex-match-substring match)))
                       (string-translate dashed "-" " ")))))

(define (irregex-split-once irx string)
  (let ((match (irregex-search irx string)))
    (if match
        (list (substring string 0 (irregex-match-start-index match))
              (substring string (irregex-match-end-index match)))
        (list string))))

(define (tokenize-duckyscript input)
  (define blank "[\t ]+")
  (define newline "\r?\n")
  (let ((lines (remove (lambda (line)
                         (or (irregex-match blank line)
                             (irregex-search "^(//|REM)" line)))
                       (irregex-split newline input))))
    (filter-map (lambda (line)
                  (let* ((line (replace-dashed-combinations line))
                         (command+args (irregex-split-once blank line))
                         (command (car command+args))
                         (args (cdr command+args)))
                    (case (string->symbol command)
                      ((STRING) command+args)
                      ((STRING_DELAY)
                       (if (null? args)
                           (list command)
                           (cons command (irregex-split-once blank (car args)))))
                      (else
                       (if (null? args)
                           (list command)
                           (cons command (string-split (car args))))))))
                lines)))

(define (parse-duckyscript tokens)
  (define known-modifiers
    '("WINDOWS" "GUI" "SHIFT" "ALT" "CONTROL" "CTRL" "COMMAND" "OPTION"))
  (define key-sym-irx '(: upper (+ (or upper numeric #\_))))
  (define key-sym-or-letter-irx `(or (or lower upper numeric) ,key-sym-irx))
  (define (check-args caller args . converters)
    (when (not (= (length args) (length converters)))
      (error "unexpected number of arguments"
             caller (length args) (length converters)))
    (map (lambda (arg converter)
           (or (converter arg) (error "invalid argument" caller arg)))
         args converters))
  (map
   (lambda (line)
     (let ((command (car line))
           (args (cdr line)))
       (case (string->symbol command)
         ((DEFAULTDELAY DEFAULT_DELAY)
          `(default-delay ,@(check-args 'default-delay args string->number)))
         ((DELAY)
          `(delay ,@(check-args 'delay args string->number)))
         ((STRING_DELAY)
          `(string-delay ,@(check-args 'string-delay args string->number identity)))
         ((STRING)
          `(string ,@(check-args 'string args identity)))
         ((REPEAT)
          `(repeat ,@(check-args 'repeat args string->number)))
         ((WINDOWS GUI SHIFT ALT CONTROL CTRL COMMAND OPTION)
          (receive (modifiers key-syms)
              (span (lambda (arg) (member arg known-modifiers)) line)
            (when (> (length key-syms) 1)
              (error "trailing keys after modifiers" modifiers key-syms))
            (if (null? key-syms)
                `(combination ,(drop-right modifiers 1) ,(last modifiers))
                (let ((key-sym (car key-syms)))
                  (when (not (irregex-match key-sym-or-letter-irx key-sym))
                    (error "invalid key" key-sym))
                  `(combination ,modifiers ,key-sym)))))
         (else ; all-caps key
          (when (not (null? args))
            (error "trailing keys" args))
          (when (not (irregex-match key-sym-irx command))
            (error "expected key" command))
          `(combination () ,command)))))
   tokens))

(define (ast->sequence commands keymap)
  (define stack (list))
  (define (compile-delay delay)
    (let loop ((delay delay))
      (cond
       ((> delay 255)
        (set! stack (append '(255 0) stack))
        (loop (- delay 255)))
       ((> delay 0) ; skip delay of 0
        (set! stack (append (list delay 0) stack))))))
  (define (compile-char char)
    (let ((bindings (hash-table-ref/default keymap char #f)))
      (when (not bindings)
        (error (format "Unknown char: ~s (~a)" char (char->integer char))))
      (let loop ((bindings bindings))
        (when (pair? bindings)
          (let* ((binding (car bindings))
                 (key-sym (if (pair? binding) (car binding) binding))
                 (modifiers (if (pair? binding) (cdr binding) '())))
            (set! stack (cons key-sym stack))
            (set! stack (cons modifiers stack))
            (loop (cdr bindings)))))))
  (define (normalize key-sym hash-table)
    (format "KEY_~a"
            (if (= (string-length key-sym) 1)
                (string (char-upcase (string-ref key-sym 0)))
                (hash-table-ref/default hash-table key-sym key-sym))))
  (define (check name type table)
    (when (not (hash-table-ref/default table name #f))
      (error (format "unknown ~a name" type) name))
    name)
  (define (check-key key-name)
    (check key-name "key" key-bytes))
  (define (check-modifier modifier-name)
    (check modifier-name "modifier" modifier-bytes))
  (define (compile-combination modifiers key-sym-or-char)
    (let ((key-sym (check-key (normalize key-sym-or-char key-aliases))))
      (set! stack (cons key-sym stack)))
    (let* ((modifiers (map (lambda (key-sym)
                             (check-modifier
                              (normalize key-sym modifier-aliases)))
                           modifiers))
           (modifiers (if (pair? modifiers) modifiers '())))
      (set! stack (cons modifiers stack))))
  (let loop ((default-delay 0)
             (last-command #f)
             (commands commands))
    (if (pair? commands)
        (let* ((command (car commands))
               (type (car command))
               (args (cdr command)))
          (case type
            ((comment)
             (loop default-delay last-command (cdr commands)))
            ((default-delay)
             (loop (car args) last-command (cdr commands)))
            ((delay)
             (compile-delay (car args))
             (loop default-delay command (cdr commands)))
            ((string)
             (for-each compile-char (string->list (car args)))
             (compile-delay default-delay)
             (loop default-delay command (cdr commands)))
            ((string-delay)
             (let ((delay (car args))
                   (string (cadr args)))
               (for-each (lambda (char)
                           (compile-char char)
                           (compile-delay delay))
                         (string->list string)))
             (compile-delay default-delay)
             (loop default-delay command (cdr commands)))
            ((repeat)
             (when (not last-command)
               (error "Last command not found"))
             (let* ((repeated (make-list (car args) last-command))
                    (commands (append repeated (cdr commands))))
               (loop default-delay last-command commands)))
            ((combination)
             (compile-combination (car args) (cadr args))
             (compile-delay default-delay)
             (loop default-delay command (cdr commands)))))
        (reverse stack))))

(define (sequence->bytes sequence)
  (define (lookup-key sym)
    (hash-table-ref key-bytes sym))
  (define (lookup-modifier sym)
    (hash-table-ref modifier-bytes sym))
  (append-map (lambda (combination)
                (let ((key-sym (car combination))
                      (modifiers (cadr combination)))
                  (if (number? key-sym)
                      combination
                      (let* ((key-byte (lookup-key key-sym))
                             (modifiers (map lookup-modifier modifiers))
                             (modifiers (if (null? modifiers) '(0) modifiers))
                             (modifier-byte (apply bitwise-ior modifiers)))
                        (list key-byte modifier-byte)))))
              (chop sequence 2)))

(define (apply-template-options template options)
  (irregex-replace/all "@([-a-z0-9]+)@" template
                       (lambda (match)
                         (let ((name (irregex-match-substring match 1)))
                           (alist-ref (string->symbol name) options)))))

(define (sequence->template sequence opts)
  (define (number->hexstring number)
    (if (< number 16)
        (format "0x0~x" number)
        (format "0x~x" number)))
  (let* ((sequence (append-map (lambda (combination)
                                 (let ((key-sym (car combination))
                                       (modifiers (cadr combination)))
                                   (if (number? key-sym)
                                       (map number->hexstring combination)
                                       (let ((modifiers (if (null? modifiers)
                                                            '("0")
                                                            modifiers)))
                                         (list key-sym
                                               (string-intersperse modifiers
                                                                   "|"))))))
                               (chop sequence 2)))
         (loop-count (or (alist-ref 'loop-count opts) 1))
         (loop-delay (or (alist-ref 'loop-delay opts) 1000))
         (template-name (or (alist-ref 'template opts)
                            "digispark-digistump"))
         (template-options `((seq-array . ,(string-intersperse sequence ", "))
                             (seq-length . ,(number->string (length sequence)))
                             (loop-count . ,(number->string loop-count))
                             (loop-delay . ,(number->string loop-delay))))
         (template (try-read-template template-name)))
    (apply-template-options template template-options)))

;;; decompiler

(define (bytes->sequence bytes)
  (define (byte->modifiers byte)
    (filter-map
     (lambda (item)
       (let ((modifier-byte (car item))
             (modifier-sym (cdr item)))
         (if (zero? (bitwise-and modifier-byte byte))
             #f
             modifier-sym)))
     (hash-table->alist modifier-bytes*)))
  (map
   (lambda (byte-pair)
     (let ((key-byte (car byte-pair))
           (modifiers-byte (cadr byte-pair)))
       (if (zero? key-byte)
           `(delay ,modifiers-byte)
           `(combination ,(byte->modifiers modifiers-byte)
                         ,(hash-table-ref key-bytes* key-byte)))))
   (chop bytes 2)))

(define (decorate pred sequence)
  (let loop ((acc '())
             (sublist '())
             (last #f)
             (sequence sequence))
    (if (pair? sequence)
        (let* ((item (car sequence))
               (type (pred item)))
          (if (and last (not (eqv? last type)))
              (loop (cons (cons last (reverse sublist)) acc)
                    (list item)
                    type
                    (cdr sequence))
              (loop acc
                    (cons item sublist)
                    type
                    (cdr sequence))))
        (if (null? sublist)
            (reverse acc)
            (let ((type (pred (car sublist))))
              (reverse (cons (cons type (reverse sublist)) acc)))))))

(define (undecorate sequence)
  (append-map cdr sequence))

(define (reconstruct-char-keys decorated keymap)
  (define unprintables '(#\backspace #\tab #\return #\escape))
  (define (combination-match? combination prefix map)
    (let ((modifiers (cadr combination))
          (key-sym (list-ref combination 2)))
      (cond
       ((null? modifiers)
        (hash-table-ref/default map (append prefix (list key-sym)) #f))
       ((null? (cdr modifiers))
        (let ((subsequence (list (list key-sym (car modifiers)))))
          (hash-table-ref/default map (append prefix subsequence) #f)))
       (else
        #f))))
  ;; TODO: consider special-casing RSHIFT as LSHIFT alternative
  (let ((keymap* (inverse-hash-table keymap)))
    (map
     (lambda (item)
       (if (eqv? (car item) 'delay)
           item
           (let loop ((acc '())
                      (combinations (cdr item)))
             (if (pair? combinations)
                 (let* ((combination (car combinations))
                        (match (combination-match? combination '() keymap*)))
                   (if (and match (not (memv match unprintables)))
                       (loop (cons `(combination () ,(string match)) acc)
                             (cdr combinations))
                       ;; TODO: consider deeper match for dead keys
                       (loop (cons combination acc)
                             (cdr combinations))))
                 `(combination ,@(reverse acc))))))
     decorated)))

(define (reconstruct-delays decorated)
  (map
   (lambda (item)
     (if (eqv? (car item) 'delay)
         (let loop ((acc '())
                    (last-delays '())
                    (delays (cdr item)))
           (if (pair? delays)
               (let ((delay (car delays)))
                 (if (equal? delay '(delay 255))
                     (loop acc (cons 255 last-delays) (cdr delays))
                     (let ((delay (+ (cadr delay) (fold + 0 last-delays))))
                       (loop (cons `(delay ,delay) acc) '() (cdr delays)))))
               `(delay ,@(reverse acc))))
         item))
   decorated))

(define (reconstruct-strings decorated)
  (define (string-command? item)
    (and (eqv? (car item) 'combination)
         (null? (cadr item))
         (= (string-length (list-ref item 2)) 1)))
  (let ((decorated (decorate (lambda (item)
                               (if (string-command? item)
                                   'string
                                   (car item)))
                             (undecorate decorated))))
    (map
     (lambda (item)
       (if (eqv? (car item) 'string)
           (let ((strings (map (lambda (command) (list-ref command 2))
                               (cdr item))))
             `(string (string ,(string-intersperse strings ""))))
           item))
     decorated)))

(define (recognize-delay-strings commands)
  (let loop ((acc '())
             (commands commands))
    (if (pair? commands)
        (let ((command (car commands)))
          (if (and (eqv? (car command) 'string)
                   (= (string-length (cadr command)) 1)
                   (not (null? (cdr commands))))
              (let ((command2 (cadr commands)))
                (if (eqv? (car command2) 'delay)
                    (let* ((delay (cadr command2))
                           (string (cadr command))
                           (string-delay `(string-delay ,delay ,string)))
                      (loop (cons string-delay acc)
                            (cddr commands)))
                    (loop (cons command2 (cons command acc))
                          (cddr commands))))
              (loop (cons command acc)
                    (cdr commands))))
        (reverse acc))))

(define (reconstruct-delay-strings commands)
  (define (string-delay-command? command)
    (eqv? (car command) 'string-delay))
  (define (same-string-delay-command? command1 command2)
    (= (cadr command1) (cadr command2)))
  (reverse
   (fold
    (lambda (item acc)
      (if (and (string-delay-command? item)
               (pair? acc)
               (string-delay-command? (car acc))
               (same-string-delay-command? item (car acc)))
          (let ((delay (cadr item))
                (string (string-append (list-ref (car acc) 2)
                                       (list-ref item 2))))
            (cons `(string-delay ,delay ,string) (cdr acc)))
          (cons item acc)))
    '()
    commands)))

(define (unparse commands)
  (define (lookup-modifier modifier)
    (hash-table-ref modifier-aliases* modifier))
  (string-intersperse
   (map
    (lambda (command)
      (let ((type (car command))
            (args (cdr command)))
        (case type
          ((combination)
           (let* ((modifiers (map strip-key-prefix (car args)))
                  (modifiers (map lookup-modifier modifiers))
                  (key-sym (strip-key-prefix (cadr args))))
             (if (null? modifiers)
                 key-sym
                 (format "~a ~a" (string-intersperse modifiers " ") key-sym))))
          ((string-delay)
           (apply format "STRING_DELAY ~a ~a" args))
          ((string)
           (apply format "STRING ~a" args))
          ((delay)
           (apply format "DELAY ~a" args))
          (else
           (error "Unknown command" type)))))
    commands)
   "\n"))

(define (decompile input keymap)
  (unparse
   (reconstruct-delay-strings
    (recognize-delay-strings
     (undecorate
      (reconstruct-strings
       (reconstruct-delays
        (reconstruct-char-keys
         (decorate car (bytes->sequence (string->bytes input)))
         keymap))))))))

;;; CLI

(define (die #!rest messages)
  (for-each
   (lambda (message)
     (display message (current-error-port))
     (newline (current-error-port)))
   messages)
  (exit 1))

(define options
  `((input
     "Input file"
     (value (required "FILE"))
     (single-char #\i))
    (output
     "Output file"
     (value (required "FILE"))
     (single-char #\o))
    (layout
     "Keyboard layout (us/de or a path to a file)"
     (value (required "FILE"))
     (single-char #\l))
    (template
     "Template to use (name of a built-in template or path to a file)"
     (value (required "FILE"))
     (single-char #\t))
    (loop-count
     "Amount of loop iterations to perform (default: 1, infinite: -1)"
     (value (required "COUNT")
            (transformer ,string->number)))
    (loop-delay
     "Amount of ms to wait after each loop iteration (default: 1000ms)"
     (value (required "COUNT")
            (transformer ,string->number)))
    (keep
     "Compilation stage to stop at (tokens, ast, sequence)"
     (value (required "STAGE")
            (transformer ,string->symbol)))
    (compact
     "Generate binary instead of C source file")
    (decompile
     "Interpret input file as binary and output file as source file")
    (list-layouts
     "Lists available layouts")
    (list-templates
     "Lists available templates")
    (list-symbols
     "Lists available key/modifier aliases and key symbols")
    (help
     "Prints this help"
     (single-char #\h))))

(define usage-hint (format "Usage: ~a [options]" (program-name)))

(define (main)
  (define (call-with-output-path path proc)
    (if path
        (call-with-output-file path proc)
        (proc (current-output-port))))
  (define (dump path thing #!optional (proc write))
    (call-with-output-path path
      (lambda (out) (proc thing out) (newline out)))
    (exit 0))
  (let* ((opts (condition-case
                (getopt-long (command-line-arguments) options)
                (e (exn)
                   (die (format "Error: ~a: ~a"
                                (get-condition-property e 'exn 'message)
                                (get-condition-property e 'exn 'arguments))
                        usage-hint
                        (usage options)))))
         (keep (alist-ref 'keep opts))
         (compact? (alist-ref 'compact opts))
         (input-path (alist-ref 'input opts)))
    (when (and keep (not (memv keep '(tokens ast sequence))))
      (die (format "Unknown value for keep: ~a" keep)))
    (when (alist-ref 'help opts)
      (print usage-hint)
      (display (usage options))
      (exit 0))
    (when (alist-ref 'list-layouts opts)
      (for-each print (built-in-layouts))
      (exit 0))
    (when (alist-ref 'list-templates opts)
      (for-each print (built-in-templates))
      (exit 0))
    (when (alist-ref 'list-symbols opts)
      (dump-aliases-symbols)
      (exit 0))
    (let ((input (if input-path
                     (call-with-input-file input-path read-all)
                     (read-all (current-input-port))))
          (output-path (alist-ref 'output opts))
          (keymap (fill-keymap (try-read-layout (or (alist-ref 'layout opts)
                                                    "us")))))
      (when (alist-ref 'decompile opts)
        (dump output-path (decompile input keymap) display))
      (let ((tokens (tokenize-duckyscript input)))
        (when (equal? keep 'tokens)
          (dump output-path tokens))
        (let ((ast (parse-duckyscript tokens)))
          (when (equal? keep 'ast)
            (dump output-path ast))
          (let ((sequence (ast->sequence ast keymap)))
            (when (equal? keep 'sequence)
              (dump output-path sequence))
            (call-with-output-path output-path
              (lambda (out)
                (if compact?
                    (display* (bytes->string (sequence->bytes sequence)) out)
                    (display (sequence->template sequence opts) out))))))))))

(when (not (get-environment-variable "TEST_MODE"))
  (main))