module Compiler.AST.Expression where

import Control.Applicative
import Control.Monad
import Data.List.NonEmpty as NE
import qualified Data.Map as Map
import Data.Maybe hiding (maybe)
import Data.Text as T

import Common
import Compiler.AST.Common
import Compiler.AST.Parser.Common
import Compiler.Lexer
import Parser.Lib
import Parser.Parser
import Test.Common as C

-- This is only used to modify nested bindings in a
-- let statement, and is not actually a part of the expression.
data Subscript
 = SubscriptExpr Subscript ExpressionWithLoc
 | PropertySubscript Subscript Identifier
 | NoSubscript Identifier
 deriving (Eq, Show)

instance ToSource Subscript where
  toSource = \case
    SubscriptExpr sub i     -> T.concat [toSource sub, "[", toSource i, "]"]
    PropertySubscript sub i -> T.concat [toSource sub, ".", toSource i]
    NoSubscript i           -> toSource i

instance HasGen Subscript where
  getGen = recursive choice [NoSubscript <$> getGen]
    [ SubscriptExpr <$> getGen <*> getGen
    , PropertySubscript <$> getGen <*> getGen
    ]

data LiteralExpression
  = LAtomic Literal
  | LArray [ExpressionWithLoc]
  | LObject (Map.Map Text ExpressionWithLoc)
  deriving (Eq, Show)

instance ToSource LiteralExpression where
  toSource (LAtomic l) = toSource l
  toSource (LArray l)  = T.concat ["[", T.intercalate ", " (toSource <$> l), "]"]
  toSource (LObject l)  = T.concat ["{", T.intercalate ", " (toSource' <$> Map.toList l), "}"]
    where
      toSource' (x, y) = T.concat [toSource x, ": ", toSource y]

instance HasGen LiteralExpression where
  getGen = choice
    [ LAtomic <$> getGen
    , LArray <$> getGen
    , LObject <$> (Map.fromList <$> list (linear 1 2) (do a <- (text (linear 1 50) (enum 'a' 'z')); b <- getGen; pure ("key"<>a, b)))
    ]

data SubscriptedExpression
  = EArraySubscript ExpressionWithLoc ExpressionWithLoc
  | EPropertySubscript ExpressionWithLoc Identifier
  deriving (Show, Eq)

instance ToSource SubscriptedExpression where
  toSource = \case
    EArraySubscript sub i    -> T.concat [toSource sub, "[", toSource i, "]"]
    EPropertySubscript sub i -> T.concat [toSource sub, ".", toSource i]

instance HasGen SubscriptedExpression where
  getGen = choice
    [ EArraySubscript <$> getGen <*> getGen
    , EPropertySubscript <$> getGen <*> getGen
    ]

data ExpressionWithLoc = ExpressionWithLoc { elExpression :: Expression, elLocation ::  Location }
  deriving (Show)

instance Eq ExpressionWithLoc where
  (ExpressionWithLoc e1 _) == (ExpressionWithLoc e2 _) = e1 == e2

instance ToSource ExpressionWithLoc where
  toSource (ExpressionWithLoc e1 _) = toSource e1

data Expression
  = ELiteral LiteralExpression
  | EVar Identifier
  | ESubscripted SubscriptedExpression
  | EOperator Operator ExpressionWithLoc ExpressionWithLoc
  | ECall Identifier [ExpressionWithLoc] Bool -- Boolean is used during execution to mark tail calls
  | EConditional ExpressionWithLoc ExpressionWithLoc ExpressionWithLoc
  | EParan ExpressionWithLoc
  | EUnnamedFn (Maybe (NonEmpty Identifier)) ExpressionWithLoc
  | ENegated ExpressionWithLoc
  deriving (Show)

instance Eq Expression where -- Specialized implementation to implement equality for EParen wrapped expressions
  (ELiteral l1) == (ELiteral l2) = l1 == l2
  (EVar l1) == (EVar l2) = l1 == l2
  (ESubscripted l1) == (ESubscripted l2) = l1 == l2
  (EOperator o ex1 ex2) == (EOperator o1 ex3 ex4) = (o == o1) && (ex1 == ex3) && (ex2 == ex4)
  (ECall idef args _) == (ECall idef1 args2 _) = idef == idef1 && (args == args2)
  (EConditional ex1 ex2 ex3) == (EConditional ex4 ex5 ex6) = (ex1 == ex4) && (ex2 == ex5) && (ex3 == ex6)
  (EUnnamedFn args2 ex2) == (EUnnamedFn args3 ex3) = (args2 == args3) && (ex2 == ex3)
  (EParan ex1) == (EParan ex2) = ex1 == ex2
  (EParan ex1) == ex2 = (elExpression ex1) == ex2
  ex1 == (EParan ex2) = ex1 == (elExpression ex2)
  _ == _ = False

instance ToSource Expression where
  toSource = \case
    ELiteral l           -> toSource l
    ENegated l           -> toSource OpMinus <> toSource l
    EVar subscript               -> toSource subscript
    ESubscripted subscript               -> toSource subscript
    EOperator op exp1 exp2 -> T.concat [pOpen, toSource exp1, ws, toSource op, ws, toSource exp2, pClose]
    ECall i args _ -> T.concat $ [toSource i, toSource DlParenOpen] <> [T.intercalate ", " (toSource <$> args)] <>  [toSource DlParenClose]
    EConditional bexp exp1 exp2 -> T.concat
      [ pOpen
      , toSource KwIf
      , ws
      , toSource DlParenOpen
      , toSource bexp
      , toSource DlParenClose
      , ws
      , toSource KwThen
      , ws
      , toSource exp1
      , ws
      , toSource KwElse
      , ws
      , toSource exp2
      , pClose
      ]
    EParan exp1 -> T.concat [toSource exp1]
    EUnnamedFn args expr ->
      let
        argsSrc = case args of
          Just args' ->  T.intercalate ", " (toSource <$> (NE.toList args'))
          Nothing -> ""
      in T.concat [toSource KwFn, ws, toSource DlParenOpen, argsSrc, toSource DlParenClose, ws, toSource expr, ws, toSource KwEndFn]

    where
      ws = toSource (Space 1)
      pOpen = toSource DlParenOpen
      pClose = toSource DlParenClose

instance HasGen ExpressionWithLoc where
  getGen = ExpressionWithLoc <$> getGen <*> (pure emptyLocation)

instance HasGen Expression where
  getGen = recursive choice
    [ ELiteral <$> getGen
    , EVar <$> getGen
    , ESubscripted <$> getGen
    ]
    [ EOperator <$> getGen <*> getGen <*> getGen
    , ECall <$> getGen <*> (list (linear 1 2) getGen) <*> (pure False)
    , EConditional <$> getGen <*> getGen <*> getGen
    , EParan <$> getGen
    , EUnnamedFn <$> (C.maybe (nonEmptyGen getGen)) <*> getGen
    ]

addLRecursion :: ExpressionWithLoc -> AstParser ExpressionWithLoc
addLRecursion exp0 = do
    exp1 <- (parseAnyDots exp0) <|> (pure exp0)
    exp2 <- (parseAnySubscripts exp1) <|> (pure exp1)
    ((precedenceSort <$> (operatorParser exp2)) <|> (pure exp2))

parseAnyDots :: ExpressionWithLoc -> AstParser ExpressionWithLoc
parseAnyDots exp0 = do
  surroundWs_ (parseDelimeter DlPeriod)
  idf <- mandatory parseIdentifier
  addLRecursion (ExpressionWithLoc (ESubscripted (EPropertySubscript exp0 idf)) (elLocation exp0))

parseAnySubscripts :: ExpressionWithLoc -> AstParser ExpressionWithLoc
parseAnySubscripts exp0 = do
  surroundWs_ (parseDelimeter DlSquareParenOpen)
  indexExpr <- mandatory (astParser @ExpressionWithLoc)
  surroundWs_ (parseDelimeter DlSquareParenClose)
  addLRecursion (ExpressionWithLoc (ESubscripted (EArraySubscript exp0 indexExpr)) (elLocation exp0))

instance HasAstParser Expression where
  astParser = nameParser "Expression" $ elExpression <$> astParser

instance HasAstParser ExpressionWithLoc where
  astParser = nameParser "ExpressionWithLoc" $ do
      loc <- getParserLocation
      expr <- parserWithoutLR
      addLRecursion (ExpressionWithLoc expr loc)
    where
      parserWithoutLR
        =   literalParser
        <|> unnamedFnParser
        <|> callParser
        <|> varParser
        <|> conditionalParser
        <|> parenthesisParser
        <|> negatedExpressionParser

unnamedFnParser :: AstParser Expression
unnamedFnParser = surroundWs $ do
  surroundWs_ (parseKeyword KwFn)
  args <- parseItemListInParen parseIdentifier
  expr <- mandatory (surroundWs (astParser @ExpressionWithLoc))
  mandatory $ surroundWs_ $ parseKeyword KwEndFn
  pure $ EUnnamedFn args expr

parenthesisParser :: AstParser Expression
parenthesisParser = surroundWs $ do
  void $ parseDelimeter DlParenOpen
  expr <- surroundWs (mandatory (astParser @ExpressionWithLoc))
  void $ mandatory (parseDelimeter DlParenClose)
  pure $ EParan expr

negatedExpressionParser :: AstParser Expression
negatedExpressionParser = surroundWs $ do
  parseOperator >>= \case
    OpMinus -> do
      e <- (mandatory (astParser @ExpressionWithLoc))
      pure $ ENegated e
    _ -> cantHandle

literalParser :: AstParser Expression
literalParser = ELiteral <$> (atomicLiteralParser <|> arrayLiteralParser <|> objectLiteralParser)

atomicLiteralParser :: AstParser LiteralExpression
atomicLiteralParser = surroundWs $ do
  l <- parseToken "Atomic Literal" (\case
    TkLiteral l             -> Just l
    _                       -> Nothing)
  pure $ LAtomic l

arrayLiteralParser :: AstParser LiteralExpression
arrayLiteralParser = do
  let itemParser = astParser @ExpressionWithLoc
  surroundWs_ (parseDelimeter DlSquareParenOpen)
  args <- optional itemParser >>= \case
    Just argHead -> do
      argsTail <- many $ do
        surroundWs_ $ parseDelimeter DlComma
        mandatory itemParser
      pure (argHead : argsTail)
    Nothing -> pure []
  surroundWs_ $ mandatory (parseDelimeter DlSquareParenClose)
  pure $ LArray args

objectLiteralParser :: AstParser LiteralExpression
objectLiteralParser = do
  let mapKeyParser = (unIdentifer <$> parseIdentifier) <|> (parseToken "Map key" $ \case
        TkLiteral (LitString t) -> Just t
        _                       -> Nothing)
  let itemParser = do
        key <- mapKeyParser
        surroundWs_ $ mandatory $ nameParser "Colon" $ parseDelimeter DlColon
        expr <- mandatory (astParser @ExpressionWithLoc)
        pure (key, expr)
  surroundWs_ (parseDelimeter DlBraceParenOpen)
  args <- optional itemParser >>= \case
    Just argHead -> do
      argsTail <- many $ do
        surroundWs_ $ parseDelimeter DlComma
        mandatory itemParser
      pure (argHead : argsTail)
    Nothing -> pure []
  surroundWs_ (mandatory $ parseDelimeter DlBraceParenClose)
  pure $ LObject $ Map.fromList args

parseSubscript :: AstParser Subscript
parseSubscript = (NoSubscript <$> parseIdentifier) >>= parseSubscript'

parseSubscript' :: Subscript -> AstParser Subscript
parseSubscript' subin =
  optional (parseKeySubscript subin <|> parsePropertySubscript subin) >>= \case
    Just x  -> pure x
    Nothing -> pure subin

parsePropertySubscript :: Subscript -> AstParser Subscript
parsePropertySubscript subin = do
  whitespaceOrNl
  void $ parseDelimeter DlPeriod
  identi <- mandatory $ surroundWs parseIdentifier
  parseSubscript' (PropertySubscript subin identi)

parseKeySubscript :: Subscript -> AstParser Subscript
parseKeySubscript subin = do
  whitespaceOrNl
  void $ parseDelimeter DlSquareParenOpen
  s <- mandatory $ surroundWs (astParser @ExpressionWithLoc)
  void $ parseDelimeter DlSquareParenClose
  parseSubscript' (SubscriptExpr subin s)

varParser :: AstParser Expression
varParser = nameParser "Variable" $ surroundWs (EVar <$> parseIdentifier)

callParser :: AstParser Expression
callParser = do
  (idf, args) <- callParser_
  pure $ ECall idf args False

callParser_ :: AstParser (Identifier, [ExpressionWithLoc])
callParser_ = surroundWs $ do
  idf <- parseIdentifier
  margs <- parseItemListInParen (astParser @ExpressionWithLoc)
  pure (idf, fromMaybe [] (NE.toList <$> margs))

conditionalParser :: AstParser Expression
conditionalParser = surroundWs $ do
  _ <- parseKeyword KwIf
  whitespace
  boolExp <- mandatory (astParser @ExpressionWithLoc)
  surroundWs_ $ mandatory (parseKeyword KwThen)
  exp1 <- mandatory (astParser @ExpressionWithLoc)
  surroundWs_ $ mandatory (parseKeyword KwElse)
  exp2 <- mandatory (astParser @ExpressionWithLoc)
  pure $ EConditional boolExp exp1 exp2

operatorParser :: ExpressionWithLoc -> AstParser ExpressionWithLoc
operatorParser lexp = surroundWs $ do
  optional parseOperator >>= \case
    Just operator -> do
      rexp <- surroundWs (mandatory (astParser @ExpressionWithLoc))
      pure $ ExpressionWithLoc (EOperator operator lexp rexp) (elLocation lexp)
    Nothing -> pure lexp

precedenceSort :: ExpressionWithLoc -> ExpressionWithLoc
precedenceSort (ExpressionWithLoc (EOperator op exL (ExpressionWithLoc (EOperator op1 exRL exRR) l2)) l1) =
  if op > op1
      then ExpressionWithLoc (EOperator op1 (ExpressionWithLoc (EOperator op (precedenceSort exL) (precedenceSort exRL)) l2) (precedenceSort exRR)) l1
      else ExpressionWithLoc (EOperator op (precedenceSort exL) (ExpressionWithLoc (EOperator op1 (precedenceSort exRL) (precedenceSort exRR)) l2)) l1
precedenceSort ex = ex