(module ast mzscheme
  (require (planet "contract-utils.ss" ("cobbe" "contract-utils.plt" 1 0))
           (lib "contract.ss")
           "../config.ss"
           "token.ss")

  ;; TODO: save location information for operators?

  (define (Property? x)
    (or (Identifier? x)
        (StringLiteral? x)
        (NumericLiteral? x)))

  (define (SourceElement? x)
    (or (Statement? x)
        (Declaration? x)))

  (define (SubStatement? x)
    (or (Statement? x)
        (and (Declaration? x)
             (not (FunctionDeclaration? x)))
        (and (allow-nested-function-declarations?)
             (FunctionDeclaration? x))))

  (define (Identifier=? id1 id2)
    (eq? (Identifier-name id1)
         (Identifier-name id2)))

  (define postfix-operators '(++ --))
  (define prefix-operators '(delete void typeof ++ -- + - ~ !))
  (define infix-operators '(* / % + -
                            << >> >>> < > <= >=
                            instanceof in
                            == != === !==
                            & ^ \|
                            && \|\|))
  (define assignment-operators '(= *= /= %= += -= <<= >>= >>>= &= ^= \|=))

  ;; assignment-operator->infix-operator : assignment-operator -> (optional infix-operator)
  (define (assignment-operator->infix-operator aop)
    (and (not (eq? aop '=))
         (let* ([aop-str (symbol->string aop)]
                [op-str (substring aop-str 0 (sub1 (string-length aop-str)))])
           (string->symbol op-str))))

  (define (postfix-operator? x) (and (memq x postfix-operators) #t))
  (define (prefix-operator? x) (and (memq x prefix-operators) #t))
  (define (infix-operator? x) (and (memq x infix-operators) #t))
  (define (assignment-operator? x) (and (memq x assignment-operators) #t))

  (define PostfixOperator/c (apply symbols postfix-operators))
  (define PrefixOperator/c (apply symbols prefix-operators))
  (define InfixOperator/c (apply symbols infix-operators))
  (define AssignmentOperator/c (apply symbols assignment-operators))

  ;; ===========================================================================
  ;; TERMS
  ;; ===========================================================================

  (define-struct Term (location))

  (define-struct (Declaration Term) ())
  (define-struct (Statement Term) ())
  (define-struct (Expression Term) ())

  ;; ===========================================================================
  ;; DECLARATIONS
  ;; ===========================================================================

  ;; Identifier * (listof Identifier) * (listof SourceElement)
  (define-struct (FunctionDeclaration Declaration) (name args body) #f)

  ;; (nelistof VariableInitializer)
  (define-struct (VariableDeclaration Declaration) (bindings) #f)

  ;; Identifier * (optional Expression)
  (define-struct (VariableInitializer Term) (id init) #f)

  ;; TODO: should `let-ness' be a proprety of FunctionDeclarations, rather than
  ;;       wrapping FunctionDeclarations inside LetDeclarations?
  ;; (union (nelistof VariableInitializer) (nelistof FunctionDeclaration))
  (define-struct (LetDeclaration Declaration) (bindings) #f)

  ;; ===========================================================================
  ;; EXPRESSIONS
  ;; ===========================================================================

  ;; string
  (define-struct (StringLiteral Expression) (value) #f)

  ;; number
  (define-struct (NumericLiteral Expression) (value) #f)

  ;; boolean
  (define-struct (BooleanLiteral Expression) (value) #f)

  ;;
  (define-struct (NullLiteral Expression) () #f)

  ;; string * boolean * boolean
  (define-struct (RegexpLiteral Expression) (pattern global? case-insensitive?) #f)

  ;; (listof (optional Expression))
  (define-struct (ArrayLiteral Expression) (elements) #f)

  ;; (listof (cons Property Expression))
  (define-struct (ObjectLiteral Expression) (properties) #f)

  ;;
  (define-struct (ThisReference Expression) () #f)

  ;; Identifier
  (define-struct (VarReference Expression) (id) #f)

  ;; Expression * Expression
  (define-struct (BracketReference Expression) (container key) #f)

  ;; Expression * Identifier
  (define-struct (DotReference Expression) (container id) #f)

  ;; Expression * (listof Expression)
  (define-struct (NewExpression Expression) (constructor arguments) #f)

  ;; Expression * PostfixOperator
  (define-struct (PostfixExpression Expression) (expression operator) #f)

  ;; PrefixOperator * Expression
  (define-struct (PrefixExpression Expression) (operator expression) #f)

  ;; Expression * InfixOperator * Expression
  (define-struct (InfixExpression Expression) (left operator right) #f)

  ;; Expression * Expression * Expression
  (define-struct (ConditionalExpression Expression) (test consequent alternate) #f)

  ;; Expression * AssignmentOperator * Expression
  (define-struct (AssignmentExpression Expression) (lhs operator rhs) #f)

  ;; (optional Identifier) * (listof Identifier) * (listof SourceElement)
  (define-struct (FunctionExpression Expression) (name args body) #f)

  ;; (listof VariableInitializer) Expression
  (define-struct (LetExpression Expression) (bindings body) #f)

  ;; Expression * (listof Expression)
  (define-struct (CallExpression Expression) (method args) #f)

  ;; Expression
  (define-struct (ParenExpression Expression) (expression) #f)

  ;; (listof Expression)
  (define-struct (ListExpression Expression) (expressions) #f)

  ;; ===========================================================================
  ;; STATEMENTS
  ;; ===========================================================================

  ;; (listof SubStatement)
  (define-struct (BlockStatement Statement) (statements) #f)

  ;;
  (define-struct (EmptyStatement Statement) () #f)

  ;; Expression
  (define-struct (ExpressionStatement Statement) (expression) #f)

  ;; Expression * SubStatement * (optional SubStatement)
  (define-struct (IfStatement Statement) (test consequent alternate) #f)

  ;; SubStatement * Expression
  (define-struct (DoWhileStatement Statement) (body test) #f)

  ;; Expression * SubStatement
  (define-struct (WhileStatement Statement) (test body) #f)

  ;; (union (optional Expression) VariableDeclaration) * (optional Expression) * (optional Expression) * SubStatement
  (define-struct (ForStatement Statement) (init test incr body) #f)

  ;; (union Expression VariableDeclaration) * Expression * SubStatement
  (define-struct (ForInStatement Statement) (lhs container body) #f)

  ;; (optional Identifier)
  (define-struct (ContinueStatement Statement) (label) #f)

  ;; (optional Identifier)
  (define-struct (BreakStatement Statement) (label) #f)

  ;; (optional Expression)
  (define-struct (ReturnStatement Statement) (value) #f)

  ;; (listof VariableInitializer) SubStatement
  (define-struct (LetStatement Statement) (bindings body) #f)

  ;; Expression * SubStatement
  (define-struct (WithStatement Statement) (context body) #f)

  ;; Expression * (listof CaseClause)
  (define-struct (SwitchStatement Statement) (expression cases) #f)

  ;; Identifier * SubStatement
  (define-struct (LabelledStatement Statement) (label statement) #f)

  ;; Expression
  (define-struct (ThrowStatement Statement) (value) #f)

  ;; Statement * (listof CatchClause) * (optional Statement)
  (define-struct (TryStatement Statement) (body catch finally) #f)

  ;; symbol
  (define-struct (Identifier Term) (name) #f)

  ;; (optional Expression) * (listof SubStatement)
  (define-struct (CaseClause Term) (question answer) #f)

  ;; Identifier * Statement
  (define-struct (CatchClause Term) (id body) #f)

  ;; ===========================================================================
  ;; LOCATION INFORMATION
  ;; ===========================================================================

  ;; has-location? : any -> boolean
  (define (has-location? x)
    (or (token? x) (Term? x)))

  ;; ast-location : has-location -> (optional region)
  (define (ast-location ast)
    (cond
      [(token? ast) (token-location ast)]
      [(Term? ast) (Term-location ast)]
      [else (error 'ast-location "not an ast node")]))

  ;; ast-source : has-location -> (optional any)
  (define (ast-source t)
    (cond
      [(ast-location t) => region-source]
      [else #f]))

  ;; ast-start : has-location -> (optional position)
  (define (ast-start t)
    (cond
      [(ast-location t) => region-start]
      [else #f]))

  ;; ast-end : has-location -> (optional position)
  (define (ast-end t)
    (cond
      [(ast-location t) => region-end]
      [else #f]))

  ;; @ : (optional has-location) (optional has-location) -> (optional region)
  (define (@ start end)
    (and start
         end
         (let ([source (ast-source start)]
               [start (ast-start start)]
               [end (ast-end end)])
           (and start end (make-region source start end)))))

  (provide/contract
    [has-location? predicate/c]
    [ast-location (has-location? . -> . (optional/c region?))]
    [ast-source (has-location? . -> . (optional/c any/c))]
    [ast-start (has-location? . -> . (optional/c position?))]
    [ast-end (has-location? . -> . (optional/c position?))]
    [@ ((optional/c has-location?) (optional/c has-location?) . -> . (optional/c region?))])

  (provide/contract
    [Property? predicate/c]
    [SourceElement? predicate/c]
    [Identifier=? (Identifier? Identifier? . -> . boolean?)])

  (provide/contract
    [postfix-operators (listof symbol?)]
    [prefix-operators (listof symbol?)]
    [infix-operators (listof symbol?)]
    [assignment-operators (listof symbol?)]
    [assignment-operator->infix-operator (assignment-operator? . -> . (optional/c infix-operator?))]
    [postfix-operator? predicate/c]
    [prefix-operator? predicate/c]
    [infix-operator? predicate/c]
    [assignment-operator? predicate/c]
    [PostfixOperator/c flat-contract?]
    [PrefixOperator/c flat-contract?]
    [InfixOperator/c flat-contract?]
    [AssignmentOperator/c flat-contract?])

  (provide/contract
    [Term? predicate/c]
    [Term-location (Term? . -> . (optional/c region?))]
    [SubStatement? predicate/c])

  (provide/contract
    (struct (Declaration Term) ([location (optional/c region?)]))
    (struct (Expression Term) ([location (optional/c region?)]))
    (struct (Statement Term) ([location (optional/c region?)]))
    (struct (Identifier Term) ([location (optional/c region?)]
                               [name symbol?]))
    (struct (CaseClause Term) ([location (optional/c region?)]
                               [question (optional/c Expression?)]
                               [answer (listof SubStatement?)]))
    (struct (CatchClause Term) ([location (optional/c region?)]
                                [id Identifier?]
                                [body Statement?]))
    (struct (VariableInitializer Term) ([location (optional/c region?)]
                                        [id Identifier?]
                                        [init (optional/c Expression?)])))

  (provide/contract
    (struct (FunctionDeclaration Declaration) ([location (optional/c region?)]
                                               [name Identifier?]
                                               [args (listof Identifier?)]
                                               [body (listof SourceElement?)]))
    (struct (VariableDeclaration Declaration) ([location (optional/c region?)]
                                               [bindings (nelistof/c VariableInitializer?)]))
    (struct (LetDeclaration Declaration) ([location (optional/c region?)]
                                          [bindings (union (nelistof/c VariableInitializer?)
                                                           (nelistof/c FunctionDeclaration?))])))

  (provide/contract
    (struct (StringLiteral Expression) ([location (optional/c region?)]
                                        [value string?]))
    (struct (RegexpLiteral Expression) ([location (optional/c region?)]
                                        [pattern string?]
                                        [global? boolean?]
                                        [case-insensitive? boolean?]))
    (struct (NumericLiteral Expression) ([location (optional/c region?)]
                                         [value number?]))
    (struct (BooleanLiteral Expression) ([location (optional/c region?)]
                                         [value boolean?]))
    (struct (NullLiteral Expression) ([location (optional/c region?)]))
    (struct (ArrayLiteral Expression) ([location (optional/c region?)]
                                       [elements (listof (optional/c Expression?))]))
    (struct (ObjectLiteral Expression) ([location (optional/c region?)]
                                        [properties (listof (cons/c Property? Expression?))]))
    (struct (ThisReference Expression) ([location (optional/c region?)]))
    (struct (VarReference Expression) ([location (optional/c region?)]
                                       [id Identifier?]))
    (struct (BracketReference Expression) ([location (optional/c region?)]
                                           [container Expression?]
                                           [key Expression?]))
    (struct (DotReference Expression) ([location (optional/c region?)]
                                       [container Expression?]
                                       [id Identifier?]))
    (struct (NewExpression Expression) ([location (optional/c region?)]
                                        [constructor Expression?]
                                        [arguments (listof Expression?)]))
    (struct (PostfixExpression Expression) ([location (optional/c region?)]
                                            [expression Expression?]
                                            [operator PostfixOperator/c]))
    (struct (PrefixExpression Expression) ([location (optional/c region?)]
                                           [operator PrefixOperator/c]
                                           [expression Expression?]))
    (struct (InfixExpression Expression) ([location (optional/c region?)]
                                          [left Expression?]
                                          [operator InfixOperator/c]
                                          [right Expression?]))
    (struct (ConditionalExpression Expression) ([location (optional/c region?)]
                                                [test Expression?]
                                                [consequent Expression?]
                                                [alternate Expression?]))
    (struct (AssignmentExpression Expression) ([location (optional/c region?)]
                                               [lhs Expression?]
                                               [operator AssignmentOperator/c]
                                               [rhs Expression?]))
    (struct (FunctionExpression Expression) ([location (optional/c region?)]
                                             [name (optional/c Identifier?)]
                                             [args (listof Identifier?)]
                                             [body (listof SourceElement?)]))
    (struct (LetExpression Expression) ([location (optional/c region?)]
                                        [bindings (listof VariableInitializer?)]
                                        [body Expression?]))
    (struct (CallExpression Expression) ([location (optional/c region?)]
                                         [method Expression?]
                                         [args (listof Expression?)]))
    (struct (ParenExpression Expression) ([location (optional/c region?)]
                                          [expression Expression?]))
    (struct (ListExpression Expression) ([location (optional/c region?)]
                                         [expressions (listof Expression?)])))

  (provide/contract
    (struct (BlockStatement Statement) ([location (optional/c region?)]
                                        [statements (listof SubStatement?)]))
    (struct (EmptyStatement Statement) ([location (optional/c region?)]))
    (struct (ExpressionStatement Statement) ([location (optional/c region?)]
                                             [expression Expression?]))
    (struct (IfStatement Statement) ([location (optional/c region?)]
                                     [test Expression?]
                                     [consequent SubStatement?]
                                     [alternate (optional/c SubStatement?)]))
    (struct (DoWhileStatement Statement) ([location (optional/c region?)]
                                          [body SubStatement?]
                                          [test Expression?]))
    (struct (WhileStatement Statement) ([location (optional/c region?)]
                                        [test Expression?]
                                        [body SubStatement?]))
    (struct (ForStatement Statement) ([location (optional/c region?)]
                                      [init (union (optional/c Expression?) VariableDeclaration?)]
                                      [test (optional/c Expression?)]
                                      [incr (optional/c Expression?)]
                                      [body SubStatement?]))
    (struct (ForInStatement Statement) ([location (optional/c region?)]
                                        [lhs (union Expression? VariableDeclaration?)]
                                        [container Expression?]
                                        [body SubStatement?]))
    (struct (ContinueStatement Statement) ([location (optional/c region?)]
                                           [label (optional/c Identifier?)]))
    (struct (BreakStatement Statement) ([location (optional/c region?)]
                                        [label (optional/c Identifier?)]))
    (struct (ReturnStatement Statement) ([location (optional/c region?)]
                                         [value (optional/c Expression?)]))
    (struct (LetStatement Statement) ([location (optional/c region?)]
                                      [bindings (listof VariableInitializer?)]
                                      [body SubStatement?]))
    (struct (WithStatement Statement) ([location (optional/c region?)]
                                       [context Expression?]
                                       [body SubStatement?]))
    (struct (SwitchStatement Statement) ([location (optional/c region?)]
                                         [expression Expression?]
                                         [cases (listof CaseClause?)]))
    (struct (LabelledStatement Statement) ([location (optional/c region?)]
                                           [label Identifier?]
                                           [statement SubStatement?]))
    (struct (ThrowStatement Statement) ([location (optional/c region?)]
                                        [value Expression?]))
    (struct (TryStatement Statement) ([location (optional/c region?)]
                                      [body Statement?]
                                      [catch (listof CatchClause?)]
                                      [finally (optional/c Statement?)]))))