(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)
        (FunctionDeclaration? x)))

  (define (SubStatement? x)
    (or (Statement? 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)

  ;; Identifier * (optional Expression)
  (define-struct (VariableDeclaration Declaration) (id init) #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 VariableDeclaration) SubStatement
  (define-struct (LetExpression Expression) (bindings body) #f)

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

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

  ;; TODO: add BeginExpression e1, ..., en
  ;;   - evaluates to last value
  ;;   - change ForExpression to use expression? instead of expression*

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

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

  ;; (listof VariableDeclaration)
  (define-struct (VariableStatement Statement) (declarations) #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)

  ;; (optional (union (listof Expression) (nelistof VariableDeclaration))) * (listof 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)

  ;; 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]))

  (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?))])

  (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?])))

  (provide/contract
    (struct (FunctionDeclaration Declaration) ([location (optional/c region?)]
                                               [name Identifier?]
                                               [args (listof Identifier?)]
                                               [body (listof SourceElement?)]))
    (struct (VariableDeclaration Declaration) ([location (optional/c region?)]
                                               [id Identifier?]
                                               [init (optional/c Expression?)])))

  (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?]
                                                [right 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 VariableDeclaration?)]
                                        [body SubStatement?]))
    (struct (CallExpression Expression) ([location (optional/c region?)]
                                         [method Expression?]
                                         [args (listof Expression?)]))
    (struct (ParenExpression Expression) ([location (optional/c region?)]
                                          [expression Expression?])))

  (provide/contract
    (struct (BlockStatement Statement) ([location (optional/c region?)]
                                        [statements (listof SubStatement?)]))
    (struct (VariableStatement Statement) ([location (optional/c region?)]
                                           [declarations (nelistof/c VariableDeclaration?)]))
    (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 (listof Expression?) (nelistof/c VariableDeclaration?))]
                                      [test (optional/c Expression?)]
                                      [incr (listof 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 (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?)]))))