module Ast.Parser.State where

import qualified Ast.Types as AT
import qualified Data.Set as Set

-- | Represents the state of type definitions in the parser.
type TypeState = [(String, AT.Type)]

-- | Represents the state of variable definitions in the parser.
type VarState = [(String, AT.Type)]

-- | Tracks preprocessor-specific information, such as visited imports and import depth.
data PreProcessorState = PreProcessorState
  { PreProcessorState -> Set String
visitedImports :: Set.Set String,
    PreProcessorState -> Int
importDepth :: Int
  }
  deriving (Int -> PreProcessorState -> ShowS
[PreProcessorState] -> ShowS
PreProcessorState -> String
(Int -> PreProcessorState -> ShowS)
-> (PreProcessorState -> String)
-> ([PreProcessorState] -> ShowS)
-> Show PreProcessorState
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PreProcessorState -> ShowS
showsPrec :: Int -> PreProcessorState -> ShowS
$cshow :: PreProcessorState -> String
show :: PreProcessorState -> String
$cshowList :: [PreProcessorState] -> ShowS
showList :: [PreProcessorState] -> ShowS
Show, PreProcessorState -> PreProcessorState -> Bool
(PreProcessorState -> PreProcessorState -> Bool)
-> (PreProcessorState -> PreProcessorState -> Bool)
-> Eq PreProcessorState
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PreProcessorState -> PreProcessorState -> Bool
== :: PreProcessorState -> PreProcessorState -> Bool
$c/= :: PreProcessorState -> PreProcessorState -> Bool
/= :: PreProcessorState -> PreProcessorState -> Bool
Eq)

-- | Represents the stack of deferred expressions, grouped by scope.
type DeferState = [[AT.Expr]]

-- | The complete state of the parser, including types, variables, preprocessor state, and deferred expressions.
data ParserState = ParserState
  { ParserState -> TypeState
typeState :: TypeState,
    ParserState -> TypeState
varState :: VarState,
    ParserState -> PreProcessorState
preProcessorState :: PreProcessorState,
    ParserState -> DeferState
deferState :: DeferState
  }
  deriving (Int -> ParserState -> ShowS
[ParserState] -> ShowS
ParserState -> String
(Int -> ParserState -> ShowS)
-> (ParserState -> String)
-> ([ParserState] -> ShowS)
-> Show ParserState
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ParserState -> ShowS
showsPrec :: Int -> ParserState -> ShowS
$cshow :: ParserState -> String
show :: ParserState -> String
$cshowList :: [ParserState] -> ShowS
showList :: [ParserState] -> ShowS
Show, ParserState -> ParserState -> Bool
(ParserState -> ParserState -> Bool)
-> (ParserState -> ParserState -> Bool) -> Eq ParserState
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ParserState -> ParserState -> Bool
== :: ParserState -> ParserState -> Bool
$c/= :: ParserState -> ParserState -> Bool
/= :: ParserState -> ParserState -> Bool
Eq)

-- | The initial state of the parser.
parserState :: ParserState
parserState :: ParserState
parserState = TypeState
-> TypeState -> PreProcessorState -> DeferState -> ParserState
ParserState [] [] (Set String -> Int -> PreProcessorState
PreProcessorState Set String
forall a. Set a
Set.empty Int
0) []

-- | Inserts a custom type into the environment.
-- If the type already exists, it overwrites it.
insertType :: String -> AT.Type -> ParserState -> ParserState
insertType :: String -> Type -> ParserState -> ParserState
insertType String
name Type
t ParserState
s = ParserState
s {typeState :: TypeState
typeState = (String
name, Type
t) (String, Type) -> TypeState -> TypeState
forall a. a -> [a] -> [a]
: ParserState -> TypeState
typeState ParserState
s}

-- | Looks up a custom type in the environment by its name.
lookupType :: String -> ParserState -> Maybe AT.Type
lookupType :: String -> ParserState -> Maybe Type
lookupType String
name (ParserState {typeState :: ParserState -> TypeState
typeState = TypeState
types}) = String -> TypeState -> Maybe Type
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup String
name TypeState
types

-- | Inserts a variable into the environment.
-- If the variable already exists, it overwrites it.
insertVar :: String -> AT.Type -> ParserState -> ParserState
insertVar :: String -> Type -> ParserState -> ParserState
insertVar String
name Type
t ParserState
s = ParserState
s {varState :: TypeState
varState = (String
name, Type
t) (String, Type) -> TypeState -> TypeState
forall a. a -> [a] -> [a]
: ParserState -> TypeState
varState ParserState
s}

-- | Looks up a variable in the environment by its name.
-- Returns `Nothing` if the variable is not found.
lookupVar :: String -> ParserState -> Maybe AT.Type
lookupVar :: String -> ParserState -> Maybe Type
lookupVar String
name (ParserState {varState :: ParserState -> TypeState
varState = TypeState
vars}) = String -> TypeState -> Maybe Type
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup String
name TypeState
vars

-- | Marks an import as visited by adding it to the set of visited imports.
insertImport :: String -> ParserState -> ParserState
insertImport :: String -> ParserState -> ParserState
insertImport String
i ParserState
s =
  ParserState
s {preProcessorState :: PreProcessorState
preProcessorState = (ParserState -> PreProcessorState
preProcessorState ParserState
s) {visitedImports :: Set String
visitedImports = String -> Set String -> Set String
forall a. Ord a => a -> Set a -> Set a
Set.insert String
i (Set String -> Set String) -> Set String -> Set String
forall a b. (a -> b) -> a -> b
$ PreProcessorState -> Set String
visitedImports (PreProcessorState -> Set String)
-> (ParserState -> PreProcessorState) -> ParserState -> Set String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ParserState -> PreProcessorState
preProcessorState (ParserState -> Set String) -> ParserState -> Set String
forall a b. (a -> b) -> a -> b
$ ParserState
s}}

-- | Checks if a file has already been imported.
-- Returns `True` if the file is in the set of visited imports.
lookupImport :: String -> ParserState -> Bool
lookupImport :: String -> ParserState -> Bool
lookupImport String
i (ParserState {preProcessorState :: ParserState -> PreProcessorState
preProcessorState = PreProcessorState {visitedImports :: PreProcessorState -> Set String
visitedImports = Set String
vi}})
  | String -> Set String -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member String
i Set String
vi = Bool
True
  | Bool
otherwise = Bool
False

-- | Sets the depth of nested imports in the preprocessor state.
setImportDepth :: Int -> ParserState -> ParserState
setImportDepth :: Int -> ParserState -> ParserState
setImportDepth Int
d ParserState
s = ParserState
s {preProcessorState :: PreProcessorState
preProcessorState = (ParserState -> PreProcessorState
preProcessorState ParserState
s) {importDepth :: Int
importDepth = Int
d}}

-- | Gets the current depth of nested imports from the preprocessor state.
getImportDepth :: ParserState -> Int
getImportDepth :: ParserState -> Int
getImportDepth (ParserState {preProcessorState :: ParserState -> PreProcessorState
preProcessorState = PreProcessorState {importDepth :: PreProcessorState -> Int
importDepth = Int
d}}) = Int
d

-- | Pushes a deferred expression onto the current defer stack.
pushDefered :: AT.Expr -> ParserState -> ParserState
pushDefered :: Expr -> ParserState -> ParserState
pushDefered Expr
e s :: ParserState
s@(ParserState {deferState :: ParserState -> DeferState
deferState = DeferState
ds}) = case DeferState
ds of
  [] -> ParserState
s {deferState :: DeferState
deferState = [[Expr
e]]}
  ([Expr]
current : DeferState
rest) -> ParserState
s {deferState :: DeferState
deferState = (Expr
e Expr -> [Expr] -> [Expr]
forall a. a -> [a] -> [a]
: [Expr]
current) [Expr] -> DeferState -> DeferState
forall a. a -> [a] -> [a]
: DeferState
rest}

-- | Pushes a new empty array onto the defer state to represent entering a new scope.
pushScope :: ParserState -> ParserState
pushScope :: ParserState -> ParserState
pushScope s :: ParserState
s@(ParserState {deferState :: ParserState -> DeferState
deferState = DeferState
ds}) = ParserState
s {deferState :: DeferState
deferState = [] [Expr] -> DeferState -> DeferState
forall a. a -> [a] -> [a]
: DeferState
ds}

-- | Pops the top scope expressions stack.
-- Returns the popped stack and the updated ParserState.
popScope :: ParserState -> ([AT.Expr], ParserState)
popScope :: ParserState -> ([Expr], ParserState)
popScope s :: ParserState
s@(ParserState {deferState :: ParserState -> DeferState
deferState = DeferState
ds}) =
  case DeferState
ds of
    [] -> ([], ParserState
s)
    ([Expr]
top : DeferState
rest) -> ([Expr]
top, ParserState
s {deferState :: DeferState
deferState = DeferState
rest})