module Parse.Expression (term, typeAnnotation, definition, expr) where

import Text.Parsec hiding (newline, spaces)
import Text.Parsec.Indent (block, withPos, checkIndent)

import qualified Parse.Binop as Binop
import Parse.Helpers
import qualified Parse.Helpers as Help
import qualified Parse.Literal as Literal
import qualified Parse.Pattern as Pattern
import qualified Parse.State as State
import qualified Parse.Type as Type

import AST.V0_16
import qualified AST.Expression as E
import qualified AST.Pattern as P
import qualified AST.Variable as Var
import qualified Reporting.Annotation as A


--------  Basic Terms  --------

varTerm :: IParser E.Expr'
varTerm =
  boolean <|> ((\x -> E.Var $ Var.VarRef x) <$> var)


boolean :: IParser E.Expr'
boolean =
  let t = const (Boolean True) <$> string "True"
      f = const (Boolean False) <$> string "False"
  in
    E.Literal <$> try (t <|> f)


accessor :: IParser E.Expr'
accessor =
  do  lbl <- try (string "." >> rLabel)
      return $ E.AccessFunction lbl


negative :: IParser E.Expr'
negative =
  do  nTerm <-
          try $
            do  char '-'
                notFollowedBy (char '.' <|> char '-')
                term

      return $ E.Unary E.Negative nTerm


--------  Complex Terms  --------

listTerm :: IParser E.Expr'
listTerm =
    shader' <|> try (braces range) <|> commaSeparated
  where
    range =
      do
          lo <- expr
          (loPost, _, hiPre) <- padded (string "..")
          hi <- expr
          return $ \loPre hiPost multiline ->
              E.Range
                  (Commented loPre lo loPost)
                  (Commented hiPre hi hiPost)
                  multiline

    shader' =
      do  pos <- getPosition
          rawSrc <- Help.shader
          return $ E.GLShader (filter (/='\r') rawSrc)

    commaSeparated =
      do
        pushNewlineContext
        result <- braces'' expr
        sawNewline <- popNewlineContext
        return $
          case result of
            Left comments ->
              E.EmptyList comments
            Right terms ->
              E.ExplicitList terms sawNewline


parensTerm :: IParser E.Expr
parensTerm =
  choice
    [ try (addLocation $ parens' opFn)
    , try (addLocation $ parens' tupleFn)
    , do
          (start, e, end) <- located $ parens (parened <|> unit)
          return $ A.at start end e
    ]
  where
    opFn =
      E.Var <$> anyOp

    tupleFn =
      do  commas <- many1 comma
          return $ E.TupleFunction (length commas + 1)

    parened =
      do  expressions <- commaSep1 ((\e a b -> Commented a e b) <$> expr)
          return $ \pre post multiline ->
            case expressions pre post of
              [single] ->
                  E.Parens single

              expressions' ->
                  E.Tuple expressions' multiline

    unit =
        return $ \pre post _ -> E.Unit (pre ++ post)


recordTerm :: IParser E.Expr
recordTerm =
  addLocation $ brackets $ choice
    [ do  starter <- try (addLocation rLabel)
          (_, postStarter) <- whitespace
          choice
            [ update starter postStarter
            , literal starter postStarter
            ]
    , return $ \pre post _ -> E.EmptyRecord (pre ++ post)
    ]
  where
    update (A.A ann starter) postStarter =
      do  try (string "|")
          (_, postBar) <- whitespace
          fields <- commaSep1 field
          return $ \pre post multiline -> (E.RecordUpdate (Commented pre (A.A ann $ E.Var $ Var.VarRef starter) postStarter) (fields postBar post) multiline)

    literal (A.A _ starter) postStarter =
      do  pushNewlineContext
          try equals -- TODO: can the try break newline tracking?
          (_, preExpr) <- whitespace
          value <- expr
          multiline' <- popNewlineContext
          (_, postExpr) <- whitespace
          choice
            [ do  try comma
                  (_, preNext) <- whitespace
                  fields <- commaSep field
                  return $ \pre post multiline -> (E.Record ((Commented pre starter postStarter, Commented preExpr value postExpr, multiline') : (fields preNext post)) multiline)
            , return $ \pre post multiline -> (E.Record [(Commented pre starter postStarter, Commented preExpr value (postExpr ++ post), multiline')] multiline)
            ]

    field =
      do  pushNewlineContext
          key <- rLabel
          (postKey, _, preExpr) <- padded equals
          value <- expr
          multiline <- popNewlineContext
          return $ \pre post -> (Commented pre key postKey, Commented preExpr value post, multiline)


term :: IParser E.Expr
term =
  addLocation (choice [ E.Literal <$> Literal.literal, listTerm, accessor, negative ])
    <|> accessible (addLocation varTerm <|> parensTerm <|> recordTerm)
    <?> "an expression"


--------  Applications  --------

appExpr :: IParser E.Expr
appExpr =
  expecting "an expression" $
  do  pushNewlineContext
      start <- getMyPosition
      t <- term
      ts <- constrainedSpacePrefix term
      end <- getMyPosition
      sawNewline <- popNewlineContext
      return $
          case ts of
            [] ->
              t
            _  ->
              A.at start end $ E.App t ts sawNewline


--------  Normal Expressions  --------

expr :: IParser E.Expr
expr =
  addLocation (choice [ letExpr, caseExpr, ifExpr ])
    <|> lambdaExpr
    <|> binaryExpr
    <?> "an expression"


binaryExpr :: IParser E.Expr
binaryExpr =
    Binop.binops appExpr lastExpr anyOp
  where
    lastExpr =
        addLocation (choice [ letExpr, caseExpr, ifExpr ])
        <|> lambdaExpr
        <?> "an expression"


ifExpr :: IParser E.Expr'
ifExpr =
  let
    elseKeyword =
      (reserved "else" <?> "an 'else' branch")
        >> (snd <$> whitespace)
  in
    do
      first <- ifClause
      rest <- many (try $ (,) <$> elseKeyword <*> ifClause)
      final <- (,) <$> elseKeyword <*> expr

      return $ E.If first rest final


ifClause :: IParser E.IfClause
ifClause =
  do
    try (reserved "if")
    (_ , preCondition) <- whitespace
    condition <- expr
    (postCondition, _, bodyComments) <- padded (reserved "then")
    thenBranch <- expr
    (_, preElse) <- whitespace <?> "an 'else' branch"
    return
      ( Commented preCondition condition postCondition
      , Commented bodyComments thenBranch preElse
      )


lambdaExpr :: IParser E.Expr
lambdaExpr =
  addLocation $
  do  pushNewlineContext
      char '\\' <|> char '\x03BB' <?> "an anonymous function"
      args <- spacePrefix Pattern.term
      (preArrowComments, _, bodyComments) <- padded rightArrow
      body <- expr
      multiline <- popNewlineContext
      return $ E.Lambda args (preArrowComments ++ bodyComments) body multiline


caseExpr :: IParser E.Expr'
caseExpr =
  do  try (reserved "case")
      pushNewlineContext
      e <- (\(pre, e, post) -> Commented pre e post) <$> padded expr
      reserved "of"
      multilineSubject <- popNewlineContext
      (_, firstPatternComments) <- whitespace
      updateState $ State.setNewline -- because if statements are always formatted as multiline, we pretend we saw a newline here to avoid problems with the Box rendering model
      result <- without firstPatternComments -- <|> with
      return $ E.Case (e, multilineSubject) result
  where
    case_ preComments =
      do
          ((_, patternComments), p, (preArrowComments, _, bodyComments)) <-
              try ((,,)
                  <$> whitespace
                  <*> (checkIndent >> Pattern.expr)
                  <*> padded rightArrow
                  )
          result <- expr
          return
            ( Commented (preComments ++ patternComments) p preArrowComments
            , (bodyComments, result)
            )


    -- bracketed case statements are removed in 0.16
    -- with =
    --   brackets (fmap (const . const . const) $ semiSep1 (case_ <?> "cases { x -> ... }")) -- TODO: use comments

    without preComments =
        withPos $
            do
                r1 <- case_ preComments
                r <- many $ case_ []
                return $ r1:r


-- LET

letExpr :: IParser E.Expr'
letExpr =
  do  try (reserved "let")
      commentsAfterLet <- map E.LetComment <$> snd <$> whitespace
      defs <-
        block $
          do  def <- typeAnnotation E.LetAnnotation <|> definition E.LetDefinition
              (_, commentsAfterDef) <- whitespace
              return $ def : (map E.LetComment commentsAfterDef)
      (_, _, bodyComments) <- padded (reserved "in") -- TODO: pre comments are always empty because any whitespace was consumed before padded?
      E.Let (commentsAfterLet ++ concat defs) bodyComments <$> expr


-- TYPE ANNOTATION

typeAnnotation :: ((Var.Ref, Comments) -> (Comments, Type) -> a) -> IParser a
typeAnnotation fn =
    (\(v, pre, post) e -> fn (v, pre) (post, e)) <$> try start <*> Type.expr
  where
    start =
      do  v <- (Var.VarRef <$> lowVar) <|> parens' symOp
          (preColon, _, postColon) <- padded hasType
          return (v, preColon, postColon)


-- DEFINITION

definition :: (P.Pattern -> [(Comments, P.Pattern)] -> Comments -> E.Expr -> a) -> IParser a
definition fn =
  withPos $
    do
        (name, args) <- defStart
        (preEqualsComments, _, postEqualsComments) <- padded equals
        body <- expr
        return $ fn name args (preEqualsComments ++ postEqualsComments) body


defStart :: IParser (P.Pattern, [(Comments, P.Pattern)])
defStart =
    choice
      [ do  pattern <- try Pattern.term
            infics pattern <|> func pattern
      , do  opPattern <- addLocation (P.Var <$> parens' symOp)
            func opPattern
      ]
      <?> "the definition of a variable (x = ...)"
  where
    func pattern =
        case pattern of
          A.A _ (P.Var _) ->
              ((,) pattern) <$> spacePrefix Pattern.term

          _ ->
              return (pattern, [])

    infics p1 =
      do  (c1, (start, op, end)) <- try ((,) <$> (snd <$> whitespace) <*> located anyOp)
          p2 <- (,) <$> (snd <$> whitespace) <*> Pattern.term
          return (A.at start end (P.Var op), [(c1, p1), p2 ])
