#lang scheme/base

(require (planet cobbe/contract-utils:1/contract-utils)
         scheme/list
         scheme/contract
         scheme/match
         "ast-core.ss"
         "ast-utils.ss"
         "../../private/config.ss")

;; TODO: deal with keyword/variable conflicts (array, object, field-ref, etc.)

;; sexp? : any -> boolean
(define (sexp? x)
  (or (symbol? x)
      (number? x)
      (string? x)
      (boolean? x)
      (null? x)
      (and (pair? x) (andmap sexp? x))))

;; ===========================================================================
;; SEXP PARSING
;; ===========================================================================

;; TODO: allow multiple arity infix operators in sexps

;; TODO: lots more parse error checking
;;   - check syntax of subforms
;;   - better error messages when falling off end of match

;; sexp->Expression : sexp -> Expression
(define (sexp->Expression sexp)
  (match sexp
    [(? string?) (make-StringLiteral #f sexp)]
    [(list 'regexp pattern global? case-insensitive?)
     (make-RegexpLiteral #f pattern global? case-insensitive?)]
    [(? number?) (make-NumericLiteral #f sexp)]
    [(? boolean?) (make-BooleanLiteral #f sexp)]
    ['null (make-NullLiteral #f)]
    [(list 'array elts ...)
     (make-ArrayLiteral #f (map sexp->ArrayElement elts))]
    [(list 'object [list props vals] ...)
     (make-ObjectLiteral #f (map (lambda (prop val)
                                   (cons (sexp->Property prop) (sexp->Expression val)))
                                 props
                                 vals))]
    ['this (make-ThisReference #f)]
    [(? symbol?) (make-VarReference #f (make-Identifier #f sexp))]
    [(list 'field-ref container key)
     (make-BracketReference #f (sexp->Expression container)
                            (sexp->Expression key))]
    [(list 'field container id)
     (make-DotReference #f (sexp->Expression container) (make-Identifier #f id))]
    [(list 'new constructor args ...)
     (make-NewExpression #f (sexp->Expression constructor) (map sexp->Expression args))]
    [(list 'prefix op expr)
     (make-PrefixExpression #f op (sexp->Expression expr))]
    [(list 'postfix expr op)
     (make-PostfixExpression #f (sexp->Expression expr) op)]
    [(list (? infix-operator? op) left right)
     (make-InfixExpression #f (sexp->Expression left) op (sexp->Expression right))]
    [(list '? test consequent alternate)
     (make-ConditionalExpression #f (sexp->Expression test)
                                 (sexp->Expression consequent)
                                 (sexp->Expression alternate))]
    [(list (? assignment-operator? op) left right)
     (make-AssignmentExpression #f (sexp->Expression left) op (sexp->Expression right))]
    [(list 'function (list (? symbol? name) (? symbol? args) ...) body ...)
     (make-FunctionExpression #f (make-Identifier #f name)
                              (map (lambda (arg)
                                     (make-Identifier #f arg))
                                   args)
                              (map sexp->SourceElement body))]
    [(list 'function (list (? symbol? args) ...) body ...)
     (make-FunctionExpression #f #f
                              (map (lambda (arg)
                                     (make-Identifier #f arg))
                                   args)
                              (map sexp->SourceElement body))]
    [(list 'begin e ...)
     (make-ListExpression #f (map sexp->Expression e))]
    [(list method args ...)
     (make-CallExpression #f (sexp->Expression method) (map sexp->Expression args))]
    [_ (error 'sexp->Expression "invalid sexp")]))

;; sexp->Property : sexp -> Property
(define (sexp->Property sexp)
  (cond
    [(symbol? sexp) (make-Identifier #f sexp)]
    [(string? sexp) (make-StringLiteral #f sexp)]
    [(number? sexp) (make-NumericLiteral #f sexp)]))

;; sexp->ArrayElement : sexp -> (optional Expression)
(define (sexp->ArrayElement sexp)
  (and (not (null? sexp))
       (sexp->Expression sexp)))

;; sexp->SourceElement : sexp -> SourceElement
(define (sexp->SourceElement sexp)
  (match sexp
    [(list 'function (? symbol? name) (list (? symbol? args) ...) body ...)
     (make-FunctionDeclaration #f (make-Identifier #f name)
                               (map (lambda (arg)
                                      (make-Identifier #f arg))
                                    args)
                               (map sexp->SourceElement body))]
    [_ (sexp->SubStatement sexp)]))

;; sexp->SubStatement : sexp -> SubStatement
(define (sexp->SubStatement sexp)
  (match sexp
    [(list 'function (? symbol? name) (list (? symbol? args) ...) body ...)
     ;; TODO: better error message
     (when (not (allow-nested-function-declarations?))
       (error 'sexp->SubStatement "illegally nested function definition"))
     (make-FunctionDeclaration #f (make-Identifier #f name)
                               (map (lambda (arg)
                                      (make-Identifier #f arg))
                                    args)
                               (map sexp->SourceElement body))]
    [(list 'var decls ...)
     (make-VariableDeclaration #f (map sexp->VariableInitializer decls))]
    ;; TODO: [('let decls ...) _]
    [_ (sexp->Statement sexp)]))

;; sexp->VariableInitializer : sexp -> VariableInitializer
(define (sexp->VariableInitializer sexp)
  (match sexp
    [(? symbol?)
     (make-VariableInitializer #f (make-Identifier #f sexp) #f)]
    [[list id value]
     (make-VariableInitializer #f (make-Identifier #f id) (sexp->Expression value))]
    [_ (error 'sexp->VariableInitializer "invalid sexp")]))

;; sexp->CaseClause : sexp -> CaseClause
(define (sexp->CaseClause sexp)
  (match sexp
    [(list 'default stmts ...)
     (make-CaseClause #f #f (map sexp->SubStatement stmts))]
    [(list 'case value stmts ...)
     (make-CaseClause #f (sexp->Expression value) (map sexp->SubStatement stmts))]
    [_ (error 'sexp->Expression "invalid sexp")]))

;; sexp->CatchClause : sexp -> CatchClause
(define (sexp->CatchClause sexp)
  (match sexp
    [(list 'catch id body)
     (make-CatchClause #f (make-Identifier #f id) (sexp->BlockStatement body))]
    [_ (error 'sexp->Expression "invalid sexp")]))

;; sexp->BlockStatement : sexp -> BlockStatement
(define (sexp->BlockStatement sexp)
  (match sexp
    [(list 'block elts ...)
     (make-BlockStatement #f (map sexp->SubStatement elts))]
    [_ (error 'sexp->Expression "invalid sexp")]))

;; sexp->Statement : sexp -> Statement
(define (sexp->Statement sexp)
  (match sexp
    [(list 'block elts ...)
     (sexp->BlockStatement sexp)]
    [(list)
     (make-EmptyStatement #f)]
    [(list 'if test consequent alternate)
     (make-IfStatement #f (sexp->Expression test)
                       (sexp->SubStatement consequent)
                       (sexp->SubStatement alternate))]
    [(list 'if test consequent)
     (make-IfStatement #f (sexp->Expression test)
                       (sexp->SubStatement consequent)
                       #f)]
    [(list 'do body test)
     (make-DoWhileStatement #f (sexp->SubStatement body)
                            (sexp->Expression test))]
    [(list 'while test body)
     (make-WhileStatement #f (sexp->Expression test)
                          (sexp->SubStatement body))]
    [(list 'for (list 'var inits ...) test incr body)
     (make-ForStatement #f (make-VariableDeclaration #f (map sexp->VariableInitializer inits))
                        (sexp->Expression test)
                        (sexp->Expression incr)
                        (sexp->SubStatement body))]
    [(list 'for init test incr body)
     (make-ForStatement #f (sexp->Expression init)
                        (sexp->Expression test)
                        (sexp->Expression incr)
                        (sexp->SubStatement body))]
    [(list 'for-in (list (list 'var var) container) body)
     (make-ForInStatement #f (make-VariableDeclaration #f (list (make-VariableInitializer #f (make-Identifier #f var) #f)))
                          (sexp->Expression container)
                          (sexp->SubStatement body))]
    [(list 'for-in (list lhs container) body)
     (make-ForInStatement #f (sexp->Expression lhs)
                          (sexp->Expression container)
                          (sexp->SubStatement body))]
    [(list 'continue label)
     (make-ContinueStatement #f (make-Identifier #f label))]
    [(list 'continue)
     (make-ContinueStatement #f #f)]
    [(list 'break label)
     (make-BreakStatement #f (make-Identifier #f label))]
    [(list 'break)
     (make-BreakStatement #f #f)]
    [(list 'return value)
     (make-ReturnStatement #f (sexp->Expression value))]
    [(list 'return)
     (make-ReturnStatement #f #f)]
    [(list 'with context body)
     (make-WithStatement #f (sexp->Expression context)
                         (sexp->SubStatement body))]
    [(list 'switch test cases ...)
     (make-SwitchStatement #f (sexp->Expression test)
                           (map sexp->CaseClause cases))]
    [(list 'label label stmt)
     (make-LabelledStatement #f (make-Identifier #f label) (sexp->SubStatement stmt))]
    [(list 'throw value)
     (make-ThrowStatement #f (sexp->Expression value))]
    [(list 'try body clauses ...)
     (match (last clauses)
       [(list 'finally finally)
        (make-TryStatement #f (sexp->BlockStatement body)
                           (map sexp->CatchClause (drop-right clauses 1))
                           (sexp->BlockStatement finally))]
       [_ (make-TryStatement #f (sexp->BlockStatement body)
                             (map sexp->CatchClause clauses)
                             #f)])]
    [_ (make-ExpressionStatement #f (sexp->Expression sexp))]))

;; ===========================================================================
;; SEXP GENERATION
;; ===========================================================================

;; Expression->sexp : Expression -> sexp
(define (Expression->sexp expr)
  (match expr
    [(struct StringLiteral (_ str)) str]
    [(struct RegexpLiteral (_ pattern global? case-insensitive?))
     `(regexp ,pattern ,global? ,case-insensitive?)]
    [(struct NumericLiteral (_ n)) n]
    [(struct BooleanLiteral (_ b))
     (if b #t #f)]
    [(struct NullLiteral (_)) 'null]
    [(struct ArrayLiteral (_ elts))
     `(array ,@(map ArrayElement->sexp elts))]
    [(struct ObjectLiteral (_ (list (cons props values) ...)))
     `(object ,@(map (lambda (prop val)
                       (list (Property->sexp prop) (Expression->sexp val)))
                     props
                     values))]
    [(struct ThisReference (_)) 'this]
    [(struct VarReference (_ (struct Identifier (_ id)))) id]
    [(struct BracketReference (_ container key))
     `(field-ref ,(Expression->sexp container)
                 ,(Expression->sexp key))]
    [(struct DotReference (_ container (struct Identifier (_ id))))
     `(field ,(Expression->sexp container) ,id)]
    [(struct NewExpression (_ constructor args))
     `(new ,(Expression->sexp constructor) ,@(map Expression->sexp args))]
    [(struct PrefixExpression (_ op expr))
     `(prefix ,op ,(Expression->sexp expr))]
    [(struct PostfixExpression (_ expr op))
     `(postfix ,(Expression->sexp expr) ,op)]
    [(struct InfixExpression (_ left op right))
     `(,op ,(Expression->sexp left) ,(Expression->sexp right))]
    [(struct ConditionalExpression (_ test consequent alternate))
     `(? ,(Expression->sexp test)
         ,(Expression->sexp consequent)
         ,(Expression->sexp alternate))]
    [(struct AssignmentExpression (_ lhs op rhs))
     `(,op ,(Expression->sexp lhs) ,(Expression->sexp rhs))]
    [(struct FunctionExpression (_ #f (list (struct Identifier (_ args)) ...) body))
     `(function ,args ,@(map SourceElement->sexp body))]
    [(struct FunctionExpression (_ (struct Identifier (_ name)) (list (struct Identifier (_ args)) ...) body))
     `(function ,name ,args ,@(map SourceElement->sexp body))]
    [(struct ListExpression (_ exprs))
     `(begin ,@(map Expression->sexp exprs))]
    [(struct CallExpression (_ method args))
     `(,(Expression->sexp method) ,@(map Expression->sexp args))]
    [(struct ParenExpression (_ expression))
     (Expression->sexp expression)]))

;; SourceElement->sexp : SourceElement -> sexp
(define (SourceElement->sexp elt)
  (match elt
    [(? FunctionDeclaration?)
     (FunctionDeclaration->sexp elt)]
    [(struct VariableDeclaration (_ inits))
     `(var ,@(map VariableInitializer->sexp inits))]
    [_ (Statement->sexp elt)]))

;; Property->sexp : Property -> sexp
(define (Property->sexp elt)
  (match elt
    [(struct Identifier (_ name)) name]
    [(struct StringLiteral (_ value)) value]
    [(struct NumericLiteral (_ value)) value]))

;; ArrayElement->sexp : (optional Expression) -> sexp
(define (ArrayElement->sexp elt)
  (if elt (Expression->sexp elt) '()))

;; Statement->sexp : Statement -> sexp
(define (Statement->sexp stmt)
  (match stmt
    [(struct BlockStatement (_ elts))
     `(block ,@(map SubStatement->sexp elts))]
    [(struct EmptyStatement (_))
     '()]
    [(struct ExpressionStatement (_ expr))
     (Expression->sexp expr)]
    [(struct IfStatement (_ test consequent alternate))
     (if alternate
         `(if ,(Expression->sexp test)
              ,(SubStatement->sexp consequent)
              ,(SubStatement->sexp alternate))
         `(if ,(Expression->sexp test)
              ,(SubStatement->sexp consequent)))]
    [(struct DoWhileStatement (_ body test))
     `(do ,(SubStatement->sexp body)
        ,(Expression->sexp test))]
    [(struct WhileStatement (_ test body))
     `(while ,(Expression->sexp test)
             ,(SubStatement->sexp body))]
    [(struct ForStatement (_ init test incr body))
     `(for ,(cond
              [(not init) #f]
              [(VariableDeclaration? init) (SourceElement->sexp init)]
              [else (Expression->sexp init)])
           ,(if test (Expression->sexp test) #t)
           ,(if incr (Expression->sexp incr) #f)
           ,(SubStatement->sexp body))]
    [(struct ForInStatement (_
                             (struct VariableDeclaration (_ (list (struct VariableInitializer (_ (struct Identifier (_ var)) #f)))))
                             container
                             body))
     `(for-in ((var ,var) ,(Expression->sexp container))
              ,(SubStatement->sexp body))]
    [(struct ForInStatement (_ (? Expression? var) container body))
     `(for-in (,(Expression->sexp var) ,(Expression->sexp container))
              ,(SubStatement->sexp body))]
    [(struct ContinueStatement (_ #f))
     '(continue)]
    [(struct ContinueStatement (_ (struct Identifier (_ id))))
     `(continue ,id)]
    [(struct BreakStatement (_ #f))
     '(break)]
    [(struct BreakStatement (_ (struct Identifier (_ id))))
     `(break ,id)]
    [(struct ReturnStatement (_ value))
     (if value `(return ,(Expression->sexp value)) '(return))]
    [(struct WithStatement (_ context body))
     `(with ,(Expression->sexp context) ,(SubStatement->sexp body))]
    [(struct SwitchStatement (_ test cases))
     `(switch ,(Expression->sexp test)
              ,@(map CaseClause->sexp cases))]
    [(struct LabelledStatement (_ (struct Identifier (_ label)) stmt))
     `(label ,label ,(SubStatement->sexp stmt))]
    [(struct ThrowStatement (_ value))
     `(throw ,(Expression->sexp value))]
    [(struct TryStatement (_ body catch finally))
     `(try ,(Statement->sexp body)
           ,@(map CatchClause->sexp catch)
           ,@(if finally (list `(finally ,(Statement->sexp finally))) null))]))

;; SubStatement->sexp : SubStatement -> sexp
(define (SubStatement->sexp elt)
  (SourceElement->sexp elt))

;; FunctionDeclaration->sexp : FunctionDeclaration -> sexp
(define (FunctionDeclaration->sexp elt)
  (match elt
    [(struct FunctionDeclaration (_ (struct Identifier (_ name)) (list (struct Identifier (_ args)) ...) body))
     `(function ,name ,args ,@(map SourceElement->sexp body))]))

;; VariableInitializer->sexp : VariableInitializer -> sexp
(define (VariableInitializer->sexp decl)
  (match decl
    [(struct VariableInitializer (_ (struct Identifier (_ id)) #f))
     id]
    [(struct VariableInitializer (_ (struct Identifier (_ id)) init))
     `[,id ,(Expression->sexp init)]]))

;; CaseClause->sexp : CaseClause -> sexp
(define (CaseClause->sexp clause)
  (match clause
    [(struct CaseClause (_ #f stmts))
     `(default ,@(map SubStatement->sexp stmts))]
    [(struct CaseClause (_ expr stmts))
     `(case ,(Expression->sexp expr) ,@(map SubStatement->sexp stmts))]))

;; CatchClause->sexp : CatchClause -> sexp
(define (CatchClause->sexp clause)
  (match clause
    [(struct CatchClause (_ (struct Identifier (_ id)) body))
     `(catch ,id ,(Statement->sexp body))]))

(provide/contract
 [sexp? predicate/c]
 [Expression->sexp (Expression? . -> . sexp?)]
 [Statement->sexp (Statement? . -> . sexp?)]
 [SourceElement->sexp (SourceElement? . -> . sexp?)]
 [sexp->Expression (sexp? . -> . Expression?)]
 [sexp->Statement (sexp? . -> . Statement?)]
 [sexp->SourceElement (sexp? . -> . SourceElement?)])