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

import Control.Monad.Trans.Except (Except, runExcept, throwE)
import Control.Monad.Trans.Reader (ReaderT(..), ask, 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)

data Error
    = UndefinedTypeError Identifier
    | UnexpectedTypeInfoError Info
    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 x = NameAnalysis $ pure x
    (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)
    = flip (SymbolTable.enter identifier) globalTable . TypeInfo
    <$> withSymbolTable globalTable (dataType typeExpression)
declaration globalTable (AST.ProcedureDefinition identifier parameters variables _body) = do
    parametersInfo <- mapM parameter parameters
    variableInfo <- mapM variableDeclaration variables
    let localTable = SymbolTable.fromList $ parametersInfo <> variableInfo
    pure $ SymbolTable.enter identifier (ProcedureInfo localTable mempty) globalTable

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