#lang racket/base

(provide batch-compile load-plugins)

(require "../compiler.rkt")
(require "../reader.rkt")
(require racket/match)
(require racket/pretty)
(require (only-in racket/list index-of append-map remove-duplicates last drop-right))
(require (only-in racket/string string-split string-prefix?))
(require (only-in racket/file make-parent-directory*))
(require file/glob)

(struct schema (paths
                module-path
                value)
  #:transparent)

(define (clean-input p)
  (path->string (simplify-path (path->complete-path (expand-user-path p)) #f)))

(define (compute-base paths)
  (match paths
    ['() "."]
    [(list p) (path->string (simplify-path (build-path p 'up) #f))]
    [_ (let try-index ((i 0))
         (let scan-paths ((paths paths) (ch #f))
           (match paths
             ['() (try-index (+ i 1))]
             [(cons p more-paths)
              (cond [(= i (string-length p)) (substring p 0 i)]
                    [(not ch) (scan-paths more-paths (string-ref p i))]
                    [(eqv? ch (string-ref p i)) (scan-paths more-paths ch)]
                    [else (substring p 0 i)])])))]))

(define (expand-globs globs base0 output-directory0 k)
  (define base (or base0 (compute-base globs)))
  (define output-directory (or (and output-directory0 (clean-input output-directory0))
                               base))
  (k base
     output-directory
     (for/list [(full-input-path (in-list (map path->string
                                               (remove-duplicates
                                                (append-map (lambda (g)
                                                              (define results (glob g))
                                                              (when (null? results)
                                                                (error 'preserves-schema-rkt
                                                                       "Input not found: ~v" g))
                                                              results)
                                                            globs)))))]
       (when (not (string-prefix? full-input-path base))
         (error 'preserves-schema-rkt "Input filename ~v falls outside base ~v"
                full-input-path
                base))
       (define relative-input-path (substring full-input-path (string-length base)))
       (define module-path (for/list [(p (explode-path relative-input-path))]
                             (string->symbol (path->string (path-replace-extension p "")))))
       (define relative-output-path
         (path->string (path-replace-extension relative-input-path ".rkt")))
       (schema (schema-translation-paths full-input-path
                                         relative-input-path
                                         (path->string (build-path output-directory relative-output-path))
                                         relative-output-path)
               module-path
               (file->schema full-input-path)))))

(define (batch-compile #:inputs inputs
                       #:additional-modules [additional-modules (hash)]
                       #:base-directory [base-directory #f]
                       #:output-directory [output-directory #f]
                       #:write-files? [write-files? #t]
                       #:plugins [plugins '()])
  (expand-globs inputs
                base-directory
                output-directory
                (lambda (base-directory output-directory schemas)

                  (define index
                    (for/fold [(index additional-modules)]
                              [(s (in-list schemas))]
                      (hash-set index
                                (schema-module-path s)
                                (schema-translation-paths-relative-output-path (schema-paths s)))))

                  (define outputs
                    (for/hash [(s (in-list schemas))]
                      (match-define (schema (and tps (schema-translation-paths fi ri fo ro))
                                            mp
                                            v)
                        s)
                      (values fo
                              (schema->module-stx
                               (last mp)
                               (lambda (module-path) (hash-ref index module-path #f))
                               v
                               #:translation-paths tps
                               #:plugins plugins))))

                  (when write-files?
                    (for [((output-path stx) (in-hash outputs))]
                      (make-parent-directory* output-path)
                      (with-output-to-file output-path #:exists 'replace
                        (lambda () (pretty-write stx)))))

                  outputs)))

(define (load-plugins mods)
  (for/list [(mod mods)]
    (dynamic-require mod 'schema-compiler-plugin)))

(module+ main
  (require racket/cmdline)

  (define output-directory #f)
  (define write-files? #t)
  (define stdout? #f)
  (define base-directory #f)
  (define additional-modules '())
  (define inputs '())
  (define plugin-mods '())

  (command-line #:once-each
                ["--output" directory "Output directory for modules (default: next to sources)"
                 (set! output-directory directory)]
                ["--stdout" "Prints each module to stdout one after the other instead of writing them to files in the `--output` directory"
                 (set! stdout? #t)]
                ["--no-write-files" "Disables generation of output to the filesystem"
                 (set! write-files? #f)]
                ["--base" directory "Base directory for sources (default: common prefix)"
                 (set! base-directory directory)]
                #:multi
                ["--module" namespace=path "Additional Namespace=path import"
                 (let ((i (index-of (string->list namespace=path) #\=)))
                   (when (not i)
                     (error '--module "Argument must be Namespace=path: ~v" namespace=path))
                   (let* ((namespace-str (substring namespace=path 0 i))
                          (path-str (substring namespace=path (+ i 1))))
                     (set! additional-modules
                           (cons (list (map string->symbol (string-split namespace-str "."))
                                       path-str)
                                 additional-modules))))]
                [("--plugin-lib" "-l") lib-path "Load compiler plugin library"
                 (set! plugin-mods (cons (string->symbol lib-path) plugin-mods))]
                [("--plugin-file" "-f") rkt-file-path "Load compiler plugin source file"
                 (set! plugin-mods (cons rkt-file-path plugin-mods))]
                #:args input-glob
                (set! inputs (map clean-input input-glob)))

  (when (null? inputs)
    (error 'preserves-schema-rkt "No inputs specified."))

  (set! plugin-mods (reverse plugin-mods))

  (define outputs
    (batch-compile #:inputs inputs
                   #:additional-modules (for/hash [(e (in-list additional-modules))]
                                          (values (car e) (cadr e)))
                   #:base-directory base-directory
                   #:output-directory output-directory
                   #:write-files? write-files?
                   #:plugins (load-plugins plugin-mods)))

  (when stdout?
    (for [((output-path stx) (in-hash outputs))]
      (printf ";;-----------------------------------------------------------------\n")
      (printf ";; ~a\n" output-path)
      (newline)
      (pretty-write stx)
      (newline))))