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

import Control.Monad.Trans.Except (Except, runExcept, throwE)
import Control.Monad.Trans.Reader
    ( ReaderT(..)
    , ask
    , asks
    , runReaderT
    , withReaderT
    )
import Data.Functor ((<&>))
import qualified Language.Elna.AST as AST
import Language.Elna.Location (Identifier(..))
import Language.Elna.SymbolTable (Info(..), SymbolTable, builtInSymbolTable)
import qualified Language.Elna.SymbolTable as SymbolTable
import Language.Elna.Types (Type(..))
import Control.Monad.Trans.Class (MonadTrans(..))
import Control.Monad (foldM, unless)
import qualified Data.List.NonEmpty as NonEmpty
import Data.Foldable (traverse_)

data Error
    = UndefinedTypeError Identifier
    | UnexpectedTypeInfoError Info
    | IdentifierAlreadyDefinedError Identifier
    | UndefinedSymbolError Identifier
    deriving (Eq, Show)

newtype NameAnalysis a = NameAnalysis
    { runNameAnalysis :: ReaderT SymbolTable (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
    . flip runReaderT builtInSymbolTable
    . runNameAnalysis
    . program

program :: AST.Program -> NameAnalysis SymbolTable
program (AST.Program declarations)
    = NameAnalysis ask
    >>= flip (foldM declaration) declarations

declaration :: SymbolTable -> AST.Declaration -> NameAnalysis SymbolTable
declaration globalTable (AST.TypeDefinition identifier typeExpression)
    = withSymbolTable globalTable (dataType typeExpression)
    >>= flip (enter identifier) globalTable . TypeInfo
declaration globalTable (AST.ProcedureDefinition identifier parameters variables body) = do
        parametersInfo <- mapM parameter parameters
        variableInfo <- mapM variableDeclaration variables
        newTable <- either (identifierAlreadyDefinedError . NonEmpty.head) pure
            $ SymbolTable.fromList
            $ parametersInfo <> variableInfo
        traverse_ (statement globalTable) body
        enter identifier (ProcedureInfo newTable mempty) globalTable

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

checkSymbol :: Identifier -> SymbolTable -> NameAnalysis ()
checkSymbol identifier globalTable =
    let undefinedSymbolError = NameAnalysis
            $ lift
            $ throwE
            $ UndefinedSymbolError identifier
        isDefined = SymbolTable.member identifier globalTable
     in NameAnalysis (asks (SymbolTable.member identifier))
        >>= (flip unless undefinedSymbolError . (isDefined ||))

expression :: SymbolTable -> AST.Expression -> NameAnalysis ()
expression globalTable (AST.VariableExpression identifier) =
    checkSymbol identifier globalTable
expression _ (AST.LiteralExpression _) = pure ()
expression globalTable (AST.NegationExpression negation) =
    expression globalTable negation
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.ProductExpression lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
expression globalTable (AST.DivisionExpression lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
expression globalTable (AST.EqualExpression lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
expression globalTable (AST.NonEqualExpression lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
expression globalTable (AST.LessExpression lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
expression globalTable (AST.GreaterExpression lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
expression globalTable (AST.LessOrEqualExpression lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
expression globalTable (AST.GreaterOrEqualExpression lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
expression globalTable (AST.ArrayExpression arrayExpression indexExpression)
    = expression globalTable arrayExpression
    >> expression globalTable indexExpression

enter :: Identifier -> Info -> SymbolTable -> NameAnalysis SymbolTable
enter identifier info table
    = maybe (identifierAlreadyDefinedError identifier) pure
    $ SymbolTable.enter identifier info table

identifierAlreadyDefinedError :: Identifier -> NameAnalysis a
identifierAlreadyDefinedError = NameAnalysis
    . lift
    . throwE
    . IdentifierAlreadyDefinedError

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

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

withSymbolTable :: forall a. SymbolTable -> NameAnalysis a -> NameAnalysis a
withSymbolTable symbolTable' = NameAnalysis
    . withReaderT (const symbolTable')
    . runNameAnalysis

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