module Language.Elna.Frontend.SymbolTable
    ( SymbolTable
    , Info(..)
    , ParameterInfo(..)
    , builtInSymbolTable
    , empty
    , enter
    , fromList
    , lookup
    , member
    , scope
    , size
    , toMap
    , update
    ) where

import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import Data.List (sort)
import Data.List.NonEmpty (NonEmpty)
import qualified Data.List.NonEmpty as NonEmpty
import Data.Maybe (isJust)
import Data.Vector (Vector)
import qualified Data.Vector as Vector
import Language.Elna.Location (Identifier(..))
import Language.Elna.Frontend.Types (Type(..), intType)
import Prelude hiding (lookup)

data SymbolTable = SymbolTable (Maybe SymbolTable) (HashMap Identifier Info)
    deriving (Eq, Show)

empty :: SymbolTable
empty = SymbolTable Nothing HashMap.empty

update :: (Info -> Maybe Info) -> Identifier -> SymbolTable -> SymbolTable
update updater key (SymbolTable parent mappings) = SymbolTable parent
    $ HashMap.update updater key mappings

scope :: SymbolTable -> SymbolTable -> SymbolTable
scope parent (SymbolTable _ mappings) = SymbolTable (Just parent) mappings

builtInSymbolTable :: SymbolTable
builtInSymbolTable = SymbolTable Nothing $ HashMap.fromList
    [ ("printi", ProcedureInfo empty (Vector.singleton printiX))
    , ("printc", ProcedureInfo empty (Vector.singleton printcI))
    , ("exit", ProcedureInfo empty Vector.empty)
    , ("int", TypeInfo intType)
    ]
  where
    printiX = ParameterInfo
        { name = "x"
        , type' = intType
        , isReferenceParameter = False
        }
    printcI = ParameterInfo
        { name = "i"
        , type' = intType
        , isReferenceParameter = False
        }

toMap :: SymbolTable  -> HashMap Identifier Info
toMap (SymbolTable _ map') = map'

enter :: Identifier -> Info -> SymbolTable -> Maybe SymbolTable
enter identifier info table@(SymbolTable parent hashTable)
    | member identifier table = Nothing
    | otherwise = Just
        $ SymbolTable parent (HashMap.insert identifier info hashTable)

lookup :: Identifier -> SymbolTable -> Maybe Info
lookup identifier (SymbolTable parent table)
    | Just found <- HashMap.lookup identifier table = Just found
    | Just parent' <- parent = lookup identifier parent'
    | otherwise = Nothing

member :: Identifier -> SymbolTable -> Bool
member identifier table =
    isJust $ lookup identifier table

size :: SymbolTable -> Int
size (SymbolTable _ map') = HashMap.size map'

fromList :: [(Identifier, Info)] -> Either (NonEmpty Identifier) SymbolTable
fromList elements
    | Just identifierDuplicates' <- identifierDuplicates =
        Left identifierDuplicates'
    | otherwise = Right $ SymbolTable Nothing $ HashMap.fromList elements
  where
    identifierDuplicates = NonEmpty.nonEmpty
        $ fmap NonEmpty.head
        $ filter ((> 1) . NonEmpty.length)
        $ NonEmpty.group . sort
        $ fst <$> elements

data ParameterInfo = ParameterInfo
    { name :: Identifier
    , type' :: Type
    , isReferenceParameter :: Bool
    } deriving (Eq, Show)

data Info
    = TypeInfo Type
    | VariableInfo Bool Type
    | ProcedureInfo SymbolTable (Vector ParameterInfo)
    deriving (Eq, Show)