(import (scheme base) (nitrate) (srfi 225))

(define gensym
  (let ((x 0))
    (lambda ()
      (set! x (+ x 1))
      (string->symbol (string-append "g" (number->string x))))))

(define (do-step~ ex var-name step-name)
  (or~ (and~ (?~ null?)
             (lambda-matcher (dict input)
               (dict-update/default! (dto) dict step-name
                                     (lambda (x) (cons
                                                  (car (dict-ref (dto)
                                                                 dict
                                                                 var-name))
                                                  x))
                                     '())))
       (must-list~ ex (b~ step-name))))

(define (do-decl~ ex var-name init-name step-name)
  (cons*~ (must-id~ ex var-name)
          (b~ init-name)
          (do-step~ ex var-name step-name)))

(define (must-id~ ex name)
  (or~ (and~ (?~ symbol?) (b~ name))
       (?~ (lambda (_) (error 'must-id "invalid identifier" ex)))))

(define (must-unique-id~ ex name)
  (and~ (lambda-matcher (dict input)
          (if (memv input (dict-ref/default (dto) dict name '()))
              (error 'must-unique-id~ "duplicated formal in let" ex)
              dict))
        (must-id~ ex name)))

(define (must-list~ ex . elements)
  (or~ (apply list~ elements)
       (lambda-matcher (dict input)
         (error 'must-list~ "invalid list" ex input))))

(define (must-list*~ ex head~ tail~)
  (or~ (list*~ head~ tail~)
       (?~ (lambda (_) (error 'must-list*~ "invalid list" ex)))))

(define (let-formals-list~ ex formal expr)
  (must-list*~ ex
               (must-list~ ex
                           (must-unique-id~ ex formal)
                           (b~ expr))
               '()))

(define (let*-formals-list~ ex formal expr)
  (must-list*~ ex
               (must-list~ ex
                           (must-id~ ex formal)
                           (b~ expr))
               '()))

(define (mustp~ ex predicate? name)
  (or~ (?~ predicate?)
       (?~ (lambda (x) (error 'mustp~ name ex)))))

(define (lambda-formals~ ex positional rest)
  (or~ (and~ (?~ symbol?) (b~ rest))
       (list*~ (and~ (?~ symbol?) (b~ positional))
               (and~ (?~ symbol?) (b~ rest)))
       (?~ (lambda (_) (error 'lambda-formals~ "invalid lambda formals" ex)))))

(define (expand ex)
  (match-pr ex (r r1)
    ;; named let
    ((cons*~ 'let (and~ (?~ symbol?) (b~ 'loop))
             (let-formals-list~ ex 'formal 'expr)
             (b~ 'body))
     (let ((formals (r 'formal))
           (exprs (r 'expr))
           (loop (r1 'loop))
           (body (r1 'body)))
       (expand `(letrec ((,loop (lambda ,formals (begin ,@body))))
                  (,loop ,@exprs)))))
    ;; let
    ((cons*~ 'let (let-formals-list~ ex 'formal 'expr)
             (b~ 'body))
     (let ((formals (r 'formal))
           (exprs (r 'expr))
           (body (r1 'body)))
       (expand `((lambda ,formals (begin ,@body)) ,@exprs))))
    ;; let*
    ((cons*~ 'let* (cons~ (must-list~ ex
                                      (must-id~ ex 'formal)
                                      (b~ 'expr))
                          (b~ 'other-bindings))
             (b~ 'body))
     (let ((formal (r1 'formal))
           (expr (r1 'expr))
           (other-bindings (r1 'other-bindings))
           (body (r1 'body)))
       (expand `(let ((,formal ,expr))
                  (let* (,other-bindings) ,@body)))))
    ((cons*~ 'let* (mustp~ ex null? "invalid let*")
             (b~ 'body))
     (expand `(begin ,@(r1 'body))))
    ;; letrec
    ((cons*~ 'letrec (let*-formals-list~ ex 'formal 'expr) (b~ 'body))
     (let ((tmps (map (lambda (formal) (list formal #f))
                      (r 'formal)))
           (sets (map (lambda (formal expr)
                        (list 'set! formal expr))
                      (r 'formal) (r 'expr)))
           (body (r1 'body)))
       (expand `(let ,tmps ,@sets ,@body))))
    ;; begin
    ((list~ 'begin) (error 'expand "invalid begin" ex))
    ((list~ 'begin (b~ 'x))
     (expand (r1 'x)))
    ((cons*~ 'begin (b~ 'x) (and~ (mustp~ ex pair?
                                          "invalid begin")
                                  (b~ 'rest)))
     (let ((tmp (gensym))
           (x (r1 'x))
           (rest (r1 'rest)))
       (expand `((lambda (,tmp) (begin ,@rest)) ,x))))
    ;; or
    ((list~ 'or (b~ 'x)) (expand (r1 'x)))
    ((cons*~ 'or (b~ 'x) (b~ 'rest))
     (let ((x (r1 'x))
           (rest (r1 'rest))
           (tmp (gensym)))
       (expand `(let ((,tmp ,x))
                  (if ,tmp ,tmp (or ,@rest))))))
    ((cons~ 'or (mustp~ ex null? "invalid or")) #f)
    ;; and
    ((list~ 'and (b~ 'x)) (expand (r1 'x)))
    ((cons*~ 'and (b~ 'x) (b~ 'rest))
     (let ((x (r1 'x)) (rest (r1 'rest)))
       (expand `(if ,x (and ,@rest) #f))))
    ((cons~ 'and (mustp~ ex null? "invalid and")) #t)
    ;; do
    ((cons*~ 'do (must-list*~ ex (do-decl~ ex 'var 'init 'step) '())
             (or~ (cons~ (b~ 'condition) (b~ 'final))
                  (?~ (lambda (_) (error 'expand "invalid do" ex))))
             (b~ 'body))
     (let ((var (r 'var))
           (init (r 'init))
           (step (r 'step))
           (condition (r1 'condition))
           (final (r1 'final))
           (body (r1 'body))
           (tmp (gensym)))
       (expand `(letrec ((,tmp (lambda ,var
                                 (if ,condition
                                     (begin ,@final)
                                     (begin
                                       ,@body
                                       (,tmp ,@step))))))
                  (,tmp ,@init)))))
    ;; cond
    ((cons*~ 'cond (cons~ 'else (b~ 'rest)))
     (expand `(begin ,@(r1 'rest))))
    ((cons*~ 'cond (list~ (b~ 'pred) '=> (b~ 'recv)) (b~ 'rest))
     (let ((tmp (gensym))
           (pred (r1 'pred))
           (recv (r1 'recv))
           (rest (r1 'rest)))
       (expand `(let ((,tmp ,pred))
                  (if ,tmp
                      (,recv ,tmp)
                      (cond ,@rest))))))
    ((cons*~ 'cond (list~ (b~ 'test)) (b~ 'rest))
     (let ((tmp (gensym))
           (test (r1 'test))
           (rest (r1 'rest)))
       (expand `(let ((,tmp ,test))
                  (if ,tmp ,tmp (cond ,@rest))))))
    ((cons*~ 'cond (cons~ (b~ 'test) (b~ 'consequent)) (b~ 'rest))
     (let ((test (r1 'test))
           (consequent (r1 'consequent))
           (rest (r1 'rest)))
       (expand `(if ,test (begin ,@consequent) (cond ,@rest)))))
    ((cons~ 'cond (mustp~ ex null? "invalid cond")) #f)
    ;; lambda
    ((list~ 'lambda (b~ 'formals) (b~ 'expr))
     `(lambda ,(r1 'formals) ,(expand (r1 'expr))))
    ((cons*~ 'lambda (b~ 'formals) (b~ 'body))
     `(lambda ,(r1 'formals) ,(expand `(begin ,@(r1 'body)))))
    ;; if
    ((cons~ 'if
            (must-list~ ex (b~ 'test) (b~ 'consequent) (b~ 'alternative)))
     `(if ,(expand (r1 'test))
          ,(expand (r1 'consequent))
          ,(expand (r1 'alternative))))
    ;; set!
    ((cons~ 'set!
            (list~ (must-id~ ex 'id)
                   (b~ 'expr)))
     `(set! ,(r1 'id) ,(expand (r1 'expr))))
    ;; quote
    ((cons~ 'quote _~) ex)
    ;; function application
    ((?~ pair?) (map expand ex))
    ;; Atoms
    (_~ ex)))