{- This Source Code Form is subject to the terms of the Mozilla Public License,
   v. 2.0. If a copy of the MPL was not distributed with this file, You can
   obtain one at https://mozilla.org/MPL/2.0/. -}

module Language.Elna.Frontend.NameAnalysis
    ( nameAnalysis
    , Error(..)
    ) where

import qualified Data.List.NonEmpty as NonEmpty
import qualified Data.Vector as Vector
import qualified Language.Elna.Frontend.AST as AST
import qualified Language.Elna.Frontend.SymbolTable as SymbolTable
import Language.Elna.Frontend.SymbolTable
    ( SymbolTable
    , Info(..)
    , ParameterInfo(..)
    )
import Control.Monad.Trans.Except (Except, runExcept, throwE)
import Data.Functor ((<&>))
import Language.Elna.Location (Identifier(..))
import Language.Elna.Frontend.Types (Type(..))
import Data.Foldable (traverse_)
import Control.Monad (foldM, unless)

data Error
    = UndefinedTypeError Identifier
    | UnexpectedTypeInfoError Info
    | IdentifierAlreadyDefinedError Identifier
    | UndefinedSymbolError Identifier
    | UnexpectedArrayByValue Identifier
    deriving Eq

instance Show Error
  where
    show (UndefinedTypeError identifier) =
        concat ["Type \"", show identifier, "\" is not defined"]
    show (UnexpectedTypeInfoError info) = show info
        <> " expected to be a type"
    show (IdentifierAlreadyDefinedError identifier) =
        concat ["The identifier \"", show identifier, "\" is already defined"]
    show (UndefinedSymbolError identifier) =
        concat ["Symbol \"", show identifier, "\" is not defined"]
    show (UnexpectedArrayByValue identifier) = concat
        [ "Array \""
        , show identifier
        , "\" cannot be passed by value, only by reference"
        ]

newtype NameAnalysis a = NameAnalysis
    { runNameAnalysis :: Except Error a
    }

instance Functor NameAnalysis
  where
    fmap f (NameAnalysis x) = NameAnalysis $ f <$> x

instance Applicative NameAnalysis
  where
    pure = NameAnalysis . pure
    (NameAnalysis f) <*> (NameAnalysis x) = NameAnalysis $ f <*> x

instance Monad NameAnalysis
  where
    (NameAnalysis x) >>= f = NameAnalysis $ x >>= (runNameAnalysis . f)

nameAnalysis :: AST.Program -> Either Error SymbolTable
nameAnalysis = runExcept
    . runNameAnalysis
    . program SymbolTable.builtInSymbolTable

program :: SymbolTable -> AST.Program -> NameAnalysis SymbolTable
program symbolTable (AST.Program declarations) = do
    globalTable <- foldM procedureDeclaration symbolTable declarations
    foldM declaration globalTable declarations

procedureDeclaration :: SymbolTable -> AST.Declaration -> NameAnalysis SymbolTable
procedureDeclaration globalTable = \case
    (AST.ProcedureDeclaration identifier parameters _ _)
        -> mapM (parameter globalTable) parameters
        >>= enterOrFail identifier
        . ProcedureInfo SymbolTable.empty
        . Vector.fromList
    (AST.TypeDefinition identifier typeExpression)
        -> dataType globalTable typeExpression
        >>= enterOrFail identifier . SymbolTable.TypeInfo
  where
    enterOrFail identifier declarationInfo =
        maybe (NameAnalysis $ throwE $ IdentifierAlreadyDefinedError identifier) pure
            $ SymbolTable.enter identifier declarationInfo globalTable

declaration :: SymbolTable -> AST.Declaration -> NameAnalysis SymbolTable
declaration globalTable (AST.ProcedureDeclaration identifier parameters variables body) = do
    variableInfo <- mapM (variableDeclaration globalTable) variables
    parameterInfo <- mapM (parameterToVariableInfo globalTable) parameters
    procedureTable <- fmap (SymbolTable.scope globalTable)
        $ either (NameAnalysis . throwE . IdentifierAlreadyDefinedError . NonEmpty.head) pure
        $ SymbolTable.fromList
        $ parameterInfo <> variableInfo
    traverse_ (statement procedureTable) body
    pure $ SymbolTable.update (updater procedureTable) identifier globalTable
  where
    updater procedureTable (ProcedureInfo _ parameters') = Just
        $ ProcedureInfo procedureTable parameters'
    updater _ _ = Nothing
declaration globalTable (AST.TypeDefinition _ _) = pure globalTable

parameterToVariableInfo :: SymbolTable -> AST.Parameter -> NameAnalysis (Identifier, Info)
parameterToVariableInfo symbolTable (AST.Parameter identifier typeExpression isReferenceParameter')
    = (identifier,) . VariableInfo isReferenceParameter'
    <$> dataType symbolTable typeExpression

variableDeclaration :: SymbolTable -> AST.VariableDeclaration -> NameAnalysis (Identifier, Info)
variableDeclaration globalTable (AST.VariableDeclaration identifier typeExpression)
    = (identifier,) . VariableInfo False
    <$> dataType globalTable typeExpression

parameter :: SymbolTable -> AST.Parameter -> NameAnalysis ParameterInfo
parameter environmentSymbolTable (AST.Parameter identifier typeExpression isReferenceParameter') = do
    parameterType <- dataType environmentSymbolTable typeExpression
    case parameterType of
        ArrayType _ _
            | not isReferenceParameter' -> NameAnalysis
                $ throwE $ UnexpectedArrayByValue identifier
        _ ->
            let parameterInfo = ParameterInfo
                    { name = identifier
                    , type' = parameterType
                    , isReferenceParameter = isReferenceParameter'
                    }
             in pure parameterInfo

dataType :: SymbolTable -> AST.TypeExpression -> NameAnalysis Type
dataType environmentSymbolTable (AST.NamedType baseType) = do
    case SymbolTable.lookup baseType environmentSymbolTable of
        Just baseInfo
            | TypeInfo baseType' <- baseInfo -> pure baseType'
            | otherwise -> NameAnalysis $ throwE $ UnexpectedTypeInfoError baseInfo
        _ -> NameAnalysis $ throwE $ UndefinedTypeError baseType
dataType environmentSymbolTable (AST.ArrayType arraySize baseType) =
    dataType environmentSymbolTable baseType <&> ArrayType (fromIntegral arraySize)

checkSymbol :: SymbolTable -> Identifier -> NameAnalysis ()
checkSymbol globalTable identifier
    = unless (SymbolTable.member identifier globalTable)
    $ NameAnalysis $ throwE
    $ UndefinedSymbolError identifier

expression :: SymbolTable -> AST.Expression -> NameAnalysis ()
expression _ (AST.LiteralExpression _) = pure ()
expression globalTable (AST.SumExpression lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
expression globalTable (AST.SubtractionExpression lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
expression globalTable (AST.NegationExpression negation) =
    expression globalTable negation
expression globalTable (AST.ProductExpression lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
expression globalTable (AST.DivisionExpression lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
expression globalTable (AST.VariableExpression variableExpression) =
    variableAccess globalTable variableExpression

statement :: SymbolTable -> AST.Statement -> NameAnalysis ()
statement _ AST.EmptyStatement = pure ()
statement globalTable (AST.CallStatement name arguments)
    = checkSymbol globalTable name
    >> traverse_ (expression globalTable) arguments
statement globalTable (AST.CompoundStatement statements) =
    traverse_ (statement globalTable) statements
statement globalTable (AST.IfStatement ifCondition ifStatement elseStatement)
    = condition globalTable ifCondition
    >> statement globalTable ifStatement
    >> maybe (pure ()) (statement globalTable) elseStatement
statement globalTable (AST.AssignmentStatement lvalue rvalue)
    = variableAccess globalTable lvalue
    >> expression globalTable rvalue
statement globalTable (AST.WhileStatement whileCondition loop)
    = condition globalTable whileCondition
    >> statement globalTable loop

condition :: SymbolTable -> AST.Condition -> NameAnalysis ()
condition globalTable (AST.EqualCondition lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
condition globalTable (AST.NonEqualCondition lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
condition globalTable (AST.LessCondition lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
condition globalTable (AST.GreaterCondition lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
condition globalTable (AST.LessOrEqualCondition lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
condition globalTable (AST.GreaterOrEqualCondition lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs

variableAccess :: SymbolTable -> AST.VariableAccess -> NameAnalysis ()
variableAccess globalTable (AST.VariableAccess identifier) =
    checkSymbol globalTable identifier
variableAccess globalTable (AST.ArrayAccess arrayExpression indexExpression)
    = variableAccess globalTable arrayExpression
    >> expression globalTable indexExpression