{-
-- EPITECH PROJECT, 2025
-- Glados
-- File description:
-- Semantic Checker Engine with AST Transformation (Type Inference)
-}

module AST.Semantics.Check (checkAst, checkExpr, checkStmt) where

import qualified Data.Map.Strict as Map
import qualified Data.Text as DT
import Data.List (isPrefixOf)

import AST.Ast (Ast(..))
import AST.Semantics.Type
import AST.Semantics.CheckCall (checkCall)
import Common.Type.Integer (IntValue(..))
import Control.Monad

-- | Main entry point: validates AND transforms a list of AST nodes.
--
-- This function iterates over the top-level AST nodes, threading the
-- environment state. It returns a new AST where 'auto' types have been
-- resolved to concrete types.
--
-- @param asts: The list of raw AST nodes from the parser.
-- @return: Either an error message or the transformed AST list.
checkAst :: [Ast] -> Either String [Ast]
checkAst :: [Ast] -> Either String [Ast]
checkAst [Ast]
asts = do
    (finalAsts, _) <- (([Ast], CheckEnv) -> Ast -> Either String ([Ast], CheckEnv))
-> ([Ast], CheckEnv) -> [Ast] -> Either String ([Ast], CheckEnv)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM ([Ast], CheckEnv) -> Ast -> Either String ([Ast], CheckEnv)
processNode ([], CheckEnv
emptyEnv) [Ast]
asts
    return (reverse finalAsts)
  where
    -- acc = Accumulator (list of processed nodes)
    -- env = Current Environment state
    processNode :: ([Ast], CheckEnv) -> Ast -> Either String ([Ast], CheckEnv)
processNode ([Ast]
acc, CheckEnv
env) Ast
node = do
        (newNode, newEnv) <- CheckEnv -> Ast -> Either String (Ast, CheckEnv)
checkStmt CheckEnv
env Ast
node
        return (newNode : acc, newEnv)

-- | Helper to insert a variable into the environment.
-- Used to update the scope with new variable definitions.
insertVar :: CheckEnv -> DT.Text -> Type -> Either String CheckEnv
insertVar :: CheckEnv -> Text -> Type -> Either String CheckEnv
insertVar CheckEnv
env Text
name Type
t =
    let newVars :: Map Text Type
newVars = Text -> Type -> Map Text Type -> Map Text Type
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Text
name Type
t (CheckEnv -> Map Text Type
envVars CheckEnv
env)
    in CheckEnv -> Either String CheckEnv
forall a b. b -> Either a b
Right (CheckEnv -> Either String CheckEnv)
-> CheckEnv -> Either String CheckEnv
forall a b. (a -> b) -> a -> b
$ CheckEnv
env { envVars = newVars }

-- | Check Expression: Verifies an expression and returns its Semantic Type.
--
-- Does not modify the AST, only extracts the type.
-- Handles literals (Int, Bool), variables (Symbol), and complex structures.
checkExpr :: CheckEnv -> Ast -> Either String Type
checkExpr :: CheckEnv -> Ast -> Either String Type
checkExpr CheckEnv
_ (AInteger IntValue
_) = Type -> Either String Type
forall a b. b -> Either a b
Right Type
TyInt
checkExpr CheckEnv
_ (ABool Bool
_)    = Type -> Either String Type
forall a b. b -> Either a b
Right Type
TyBool
checkExpr CheckEnv
_ Ast
AVoid        = Type -> Either String Type
forall a b. b -> Either a b
Right Type
TyVoid
checkExpr CheckEnv
_ (ABlock []) = Type -> Either String Type
forall a b. b -> Either a b
Right Type
TyVoid
checkExpr CheckEnv
env (ASymbol Text
name) = CheckEnv -> Text -> Either String Type
checkSymbol CheckEnv
env Text
name
checkExpr CheckEnv
env (AIf Ast
c Ast
t Ast
e)    = CheckEnv -> Ast -> Ast -> Ast -> Either String Type
checkIfExpr CheckEnv
env Ast
c Ast
t Ast
e
checkExpr CheckEnv
env (ASetStruct Text
name [(Text, Ast)]
fields) =
    CheckEnv -> Text -> [(Text, Ast)] -> Either String Type
checkStructInstantiation CheckEnv
env Text
name [(Text, Ast)]
fields
checkExpr CheckEnv
env (ACall Ast
func [Ast]
args) = (CheckEnv -> Ast -> Either String Type)
-> CheckEnv -> Ast -> [Ast] -> Either String Type
checkCall CheckEnv -> Ast -> Either String Type
checkExpr CheckEnv
env Ast
func [Ast]
args
checkExpr CheckEnv
env (AAccessStruct Ast
obj Text
field) = CheckEnv -> Ast -> Text -> Either String Type
checkStructAccess CheckEnv
env Ast
obj Text
field
checkExpr CheckEnv
env (AList [Ast]
list) = CheckEnv -> [Ast] -> Either String Type
checkListExpr CheckEnv
env [Ast]
list
checkExpr CheckEnv
env (ABlock [Ast]
xs) = CheckEnv -> [Ast] -> Either String Type
checkBlockExpr CheckEnv
env [Ast]
xs
checkExpr CheckEnv
env (ASetVar Text
name Text
typeStr Ast
expr) = do
    actualType <- CheckEnv -> Ast -> Either String Type
checkExpr CheckEnv
env Ast
expr
    let expectedType = Text -> Type
parseType Text
typeStr
    if areTypesCompatible expectedType actualType
        then Right actualType
        else Left $ "Type mismatch in expression assignment for variable '"
            ++ DT.unpack name ++ "'"
checkExpr CheckEnv
env (ADefineLambda [Text]
args Ast
body) = CheckEnv -> [Text] -> Ast -> Either String Type
checkLambda CheckEnv
env [Text]
args Ast
body
checkExpr CheckEnv
env (APos Int
line Int
_ Ast
ast) =
    case CheckEnv -> Ast -> Either String Type
checkExpr CheckEnv
env Ast
ast of
        Left String
err -> 
            if String
"Error line" String -> String -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` String
err 
                then String -> Either String Type
forall a b. a -> Either a b
Left String
err
                else String -> Either String Type
forall a b. a -> Either a b
Left (String -> Either String Type) -> String -> Either String Type
forall a b. (a -> b) -> a -> b
$ String
"Error line " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
line String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
": " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
err
        Right Type
t  -> Type -> Either String Type
forall a b. b -> Either a b
Right Type
t
checkExpr CheckEnv
_ Ast
_ = String -> Either String Type
forall a b. a -> Either a b
Left String
"Error: Expression type not supported"

-- | Helper to check a block used as an expression.
-- Returns the type of the *last* instruction in the block.
checkBlockExpr :: CheckEnv -> [Ast] -> Either String Type
checkBlockExpr :: CheckEnv -> [Ast] -> Either String Type
checkBlockExpr CheckEnv
_ [] = Type -> Either String Type
forall a b. b -> Either a b
Right Type
TyVoid
checkBlockExpr CheckEnv
env [Ast
x] = CheckEnv -> Ast -> Either String Type
checkExpr CheckEnv
env Ast
x
checkBlockExpr CheckEnv
env (Ast
x:[Ast]
xs) = do
    _ <- CheckEnv -> Ast -> Either String Type
checkExpr CheckEnv
env Ast
x
    checkBlockExpr env xs

-- | Looks up a variable symbol in the environment.
checkSymbol :: CheckEnv -> DT.Text -> Either String Type
checkSymbol :: CheckEnv -> Text -> Either String Type
checkSymbol CheckEnv
env Text
name = case Text -> Map Text Type -> Maybe Type
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
name (CheckEnv -> Map Text Type
envVars CheckEnv
env) of
    Just Type
t  -> Type -> Either String Type
forall a b. b -> Either a b
Right Type
t
    Maybe Type
Nothing -> String -> Either String Type
forall a b. a -> Either a b
Left (String -> Either String Type) -> String -> Either String Type
forall a b. (a -> b) -> a -> b
$ String
"Undefined variable '" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Text -> String
DT.unpack Text
name String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"'"

-- | Checks a Lambda expression.
-- Infers the return type based on the body.
-- Assumes arguments are Integers (default behavior for untyped lambdas in this grammar).
checkLambda :: CheckEnv -> [DT.Text] -> Ast -> Either String Type
checkLambda :: CheckEnv -> [Text] -> Ast -> Either String Type
checkLambda CheckEnv
env [Text]
args Ast
body = do
    let argTypes :: [Type]
argTypes = (Text -> Type) -> [Text] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (\Text
_ -> Type
TyInt) [Text]
args

    envWithArgs <- (CheckEnv -> (Text, Type) -> Either String CheckEnv)
-> CheckEnv -> [(Text, Type)] -> Either String CheckEnv
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\CheckEnv
e (Text
n, Type
t) -> CheckEnv -> Text -> Type -> Either String CheckEnv
insertVar CheckEnv
e Text
n Type
t) CheckEnv
env ([Text] -> [Type] -> [(Text, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Text]
args [Type]
argTypes)
    retType <- checkExpr envWithArgs body
    Right (TyFunc argTypes retType)

-- | Check Statement: Verifies AND transforms a statement.
--
-- Dispatcher function that routes AST nodes to specific handlers.
-- Returns a tuple containing the transformed AST node (with inferred types)
-- and the updated environment.
--
-- @param env: Current Semantic Environment.
-- @param node: AST node to process.
checkStmt :: CheckEnv -> Ast -> Either String (Ast, CheckEnv)
checkStmt :: CheckEnv -> Ast -> Either String (Ast, CheckEnv)
checkStmt CheckEnv
env (ASetVar Text
name Text
typeStr Ast
expr) =
    CheckEnv -> Text -> Text -> Ast -> Either String (Ast, CheckEnv)
processSetVar CheckEnv
env Text
name Text
typeStr Ast
expr
checkStmt CheckEnv
env (ADefineFunc Text
name [(Text, Text)]
args Text
ret Ast
body) =
    CheckEnv
-> Text
-> [(Text, Text)]
-> Text
-> Ast
-> Either String (Ast, CheckEnv)
processFuncDef CheckEnv
env Text
name [(Text, Text)]
args Text
ret Ast
body
checkStmt CheckEnv
env (AIf Ast
c Ast
t Ast
e) = CheckEnv -> Ast -> Ast -> Ast -> Either String (Ast, CheckEnv)
processIfStmt CheckEnv
env Ast
c Ast
t Ast
e
checkStmt CheckEnv
env (AWhile Ast
c Ast
body) = CheckEnv -> Ast -> Ast -> Either String (Ast, CheckEnv)
processWhileStmt CheckEnv
env Ast
c Ast
body
checkStmt CheckEnv
env (AFor Ast
i Ast
c Ast
u Ast
b) = CheckEnv
-> Ast -> Ast -> Ast -> Ast -> Either String (Ast, CheckEnv)
processForLoop CheckEnv
env Ast
i Ast
c Ast
u Ast
b
checkStmt CheckEnv
env (AList [Ast]
list) = do
    (newList, newEnv) <- (([Ast], CheckEnv) -> Ast -> Either String ([Ast], CheckEnv))
-> ([Ast], CheckEnv) -> [Ast] -> Either String ([Ast], CheckEnv)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM ([Ast], CheckEnv) -> Ast -> Either String ([Ast], CheckEnv)
processList ([], CheckEnv
env) [Ast]
list
    return (AList (reverse newList), newEnv)
  where
    processList :: ([Ast], CheckEnv) -> Ast -> Either String ([Ast], CheckEnv)
processList ([Ast]
acc, CheckEnv
e) Ast
node = do
        (newNode, nEnv) <- CheckEnv -> Ast -> Either String (Ast, CheckEnv)
checkStmt CheckEnv
e Ast
node
        return (newNode : acc, nEnv)
checkStmt CheckEnv
env (ABlock [Ast]
list) = do
    (newList, _) <- (([Ast], CheckEnv) -> Ast -> Either String ([Ast], CheckEnv))
-> ([Ast], CheckEnv) -> [Ast] -> Either String ([Ast], CheckEnv)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM ([Ast], CheckEnv) -> Ast -> Either String ([Ast], CheckEnv)
processList ([], CheckEnv
env) [Ast]
list
    return (ABlock (reverse newList), env)
  where
    processList :: ([Ast], CheckEnv) -> Ast -> Either String ([Ast], CheckEnv)
processList ([Ast]
acc, CheckEnv
e) Ast
node = do
        (newNode, nEnv) <- CheckEnv -> Ast -> Either String (Ast, CheckEnv)
checkStmt CheckEnv
e Ast
node
        return (newNode : acc, nEnv)
checkStmt CheckEnv
env (APos Int
l Int
c Ast
ast) = 
    case CheckEnv -> Ast -> Either String (Ast, CheckEnv)
checkStmt CheckEnv
env Ast
ast of
        Left String
err -> 
            if String
"Error line" String -> String -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` String
err
                then String -> Either String (Ast, CheckEnv)
forall a b. a -> Either a b
Left String
err
                else String -> Either String (Ast, CheckEnv)
forall a b. a -> Either a b
Left (String -> Either String (Ast, CheckEnv))
-> String -> Either String (Ast, CheckEnv)
forall a b. (a -> b) -> a -> b
$ String
"Error line " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
l String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
": " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
err
        Right (Ast
newAst, CheckEnv
newEnv) -> (Ast, CheckEnv) -> Either String (Ast, CheckEnv)
forall a b. b -> Either a b
Right (Int -> Int -> Ast -> Ast
APos Int
l Int
c Ast
newAst, CheckEnv
newEnv)
checkStmt CheckEnv
env (ADefineStruct Text
name [(Text, Text)]
fields) = do
    newEnv <- CheckEnv -> Text -> [(Text, Text)] -> Either String CheckEnv
defineStruct CheckEnv
env Text
name [(Text, Text)]
fields
    return (ADefineStruct name fields, newEnv)
checkStmt CheckEnv
env (AReturn Ast
expr) = do
    actualType <- CheckEnv -> Ast -> Either String Type
checkExpr CheckEnv
env Ast
expr
    case Map.lookup (DT.pack "$RET") (envVars env) of
        Just Type
expectedType -> 
            Bool -> Either String () -> Either String ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Type -> Type -> Bool
areTypesCompatible Type
expectedType Type
actualType) (Either String () -> Either String ())
-> Either String () -> Either String ()
forall a b. (a -> b) -> a -> b
$
                String -> Either String ()
forall a b. a -> Either a b
Left (String -> Either String ()) -> String -> Either String ()
forall a b. (a -> b) -> a -> b
$ String
"Return type mismatch: expected " String -> String -> String
forall a. [a] -> [a] -> [a]
++
                       Type -> String
typeToString Type
expectedType String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" but got " String -> String -> String
forall a. [a] -> [a] -> [a]
++
                       Type -> String
typeToString Type
actualType
        Maybe Type
Nothing -> () -> Either String ()
forall a b. b -> Either a b
Right ()
    return (AReturn expr, env)
checkStmt CheckEnv
env (AImport Text
_) = (Ast, CheckEnv) -> Either String (Ast, CheckEnv)
forall a b. b -> Either a b
Right (Ast
AVoid, CheckEnv
env)
checkStmt CheckEnv
env (AExprStmt Ast
stmt) = do
    (newStmt, newEnv) <- CheckEnv -> Ast -> Either String (Ast, CheckEnv)
checkStmt CheckEnv
env Ast
stmt
    return (AExprStmt newStmt, newEnv)
checkStmt CheckEnv
env Ast
ABreak = (Ast, CheckEnv) -> Either String (Ast, CheckEnv)
forall a b. b -> Either a b
Right (Ast
ABreak, CheckEnv
env)
checkStmt CheckEnv
env Ast
ast = do
    _ <- CheckEnv -> Ast -> Either String Type
checkExpr CheckEnv
env Ast
ast
    return (ast, env)

-- | Helper to validate boolean/int condition for loops/if.
--
-- @param c: The Condition AST to check.
-- @param context: String name of the structure (e.g. "if", "while").
validateCondition :: CheckEnv -> Ast -> String -> Either String ()
validateCondition :: CheckEnv -> Ast -> String -> Either String ()
validateCondition CheckEnv
env Ast
c String
context = do
    tCond <- CheckEnv -> Ast -> Either String Type
checkExpr CheckEnv
env Ast
c
    unless (areTypesCompatible TyBool tCond ||
        areTypesCompatible TyInt tCond) $
        Left $ "Error: '" ++ context ++
               "' condition must be boolean or integer"

-- | Process variable assignment (handles auto).
-- Calculates the actual type of the expression and replaces 'auto'.
processSetVar :: CheckEnv -> DT.Text -> DT.Text -> Ast ->
    Either String (Ast, CheckEnv)
processSetVar :: CheckEnv -> Text -> Text -> Ast -> Either String (Ast, CheckEnv)
processSetVar CheckEnv
env Text
name Text
typeStr Ast
expr = do
    actualType <- CheckEnv -> Ast -> Either String Type
checkExpr CheckEnv
env Ast
expr
    let (declaredType, finalTypeStr) = 
            if typeStr == DT.pack "auto"
            then (actualType, DT.pack (typeToString actualType))
            else (parseType typeStr, typeStr)
    newEnv <- applyAssignment env name declaredType actualType
    return (ASetVar name finalTypeStr expr, newEnv)

-- | Process function definition.
--
-- 1. Prepares environment (adds function name + arguments).
-- 2. Checks function body.
-- 3. Verifies return statements recursively.
processFuncDef :: CheckEnv -> DT.Text -> [(DT.Text, DT.Text)] -> DT.Text ->
    Ast -> Either String (Ast, CheckEnv)
processFuncDef :: CheckEnv
-> Text
-> [(Text, Text)]
-> Text
-> Ast
-> Either String (Ast, CheckEnv)
processFuncDef CheckEnv
env Text
name [(Text, Text)]
args Text
ret Ast
body = do
    let retTy :: Type
retTy = Text -> Type
parseType Text
ret
    (envWithFunc, envForBody) <- CheckEnv
-> Text
-> [(Text, Text)]
-> Type
-> Either String (CheckEnv, CheckEnv)
prepareFuncEnv CheckEnv
env Text
name [(Text, Text)]
args Type
retTy
    let envWithRet = case CheckEnv -> Text -> Type -> Either String CheckEnv
insertVar CheckEnv
envForBody (String -> Text
DT.pack String
"$RET") Type
retTy of
            Right CheckEnv
e -> CheckEnv
e
            Left String
_ -> CheckEnv
envForBody 
    (newBody, _) <- checkStmt envWithRet body
    _ <- checkMissingReturn name retTy newBody
    return (ADefineFunc name args ret newBody, envWithFunc)

-- | Helper to setup the environment for a function.
-- Adds the function itself (for recursion) and its arguments.
prepareFuncEnv :: CheckEnv -> DT.Text -> [(DT.Text, DT.Text)] -> Type ->
    Either String (CheckEnv, CheckEnv)
prepareFuncEnv :: CheckEnv
-> Text
-> [(Text, Text)]
-> Type
-> Either String (CheckEnv, CheckEnv)
prepareFuncEnv CheckEnv
env Text
name [(Text, Text)]
args Type
retTy = do
    let argTypes :: [Type]
argTypes = ((Text, Text) -> Type) -> [(Text, Text)] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Text -> Type
parseType (Text -> Type) -> ((Text, Text) -> Text) -> (Text, Text) -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text, Text) -> Text
forall a b. (a, b) -> b
snd) [(Text, Text)]
args
    let funcType :: Type
funcType = [Type] -> Type -> Type
TyFunc [Type]
argTypes Type
retTy
    envWithFunc <- CheckEnv -> Text -> Type -> Either String CheckEnv
insertVar CheckEnv
env Text
name Type
funcType
    envForBody <- foldM (\CheckEnv
e (Text
n, Text
t) -> CheckEnv -> Text -> Type -> Either String CheckEnv
insertVar CheckEnv
e Text
n (Text -> Type
parseType Text
t))
                        envWithFunc args
    return (envWithFunc, envForBody)

-- | Process If statement.
--
-- Abbreviations used:
-- c = Condition AST
-- t = Then branch AST
-- e = Else branch AST
processIfStmt :: CheckEnv -> Ast -> Ast -> Ast -> Either String (Ast, CheckEnv)
processIfStmt :: CheckEnv -> Ast -> Ast -> Ast -> Either String (Ast, CheckEnv)
processIfStmt CheckEnv
env Ast
c Ast
t Ast
e = do
    CheckEnv -> Ast -> String -> Either String ()
validateCondition CheckEnv
env Ast
c String
"if"
    (newT, _) <- CheckEnv -> Ast -> Either String (Ast, CheckEnv)
checkStmt CheckEnv
env Ast
t
    (newE, _) <- checkStmt env e
    return (AIf c newT newE, env)

-- | Process While statement.
--
-- Abbreviations used:
-- c = Condition AST
processWhileStmt :: CheckEnv -> Ast -> Ast -> Either String (Ast, CheckEnv)
processWhileStmt :: CheckEnv -> Ast -> Ast -> Either String (Ast, CheckEnv)
processWhileStmt CheckEnv
env Ast
c Ast
body = do
    CheckEnv -> Ast -> String -> Either String ()
validateCondition CheckEnv
env Ast
c String
"while"
    (newBody, _) <- CheckEnv -> Ast -> Either String (Ast, CheckEnv)
checkStmt CheckEnv
env Ast
body
    return (AWhile c newBody, env)

-- | Process For loop.
--
-- Abbreviations used:
-- i = Initialization AST (e.g. i = 0)
-- c = Condition AST (e.g. i < 10)
-- u = Update AST (e.g. i = i + 1)
-- b = Body AST
processForLoop :: CheckEnv -> Ast -> Ast -> Ast -> Ast ->
    Either String (Ast, CheckEnv)
processForLoop :: CheckEnv
-> Ast -> Ast -> Ast -> Ast -> Either String (Ast, CheckEnv)
processForLoop CheckEnv
env Ast
i Ast
c Ast
u Ast
b = do
    (newInit, envAfterInit) <- CheckEnv -> Ast -> Either String (Ast, CheckEnv)
checkStmt CheckEnv
env Ast
i
    validateCondition envAfterInit c "for"
    (newUpdate, _) <- checkStmt envAfterInit u
    (newBody, _) <- checkStmt envAfterInit b
    return (AFor newInit c newUpdate newBody, env)

-- | Registers a new Structure definition in the environment.
defineStruct :: CheckEnv -> DT.Text -> [(DT.Text, DT.Text)] ->
    Either String CheckEnv
defineStruct :: CheckEnv -> Text -> [(Text, Text)] -> Either String CheckEnv
defineStruct CheckEnv
env Text
name [(Text, Text)]
fields =
    case Text -> Map Text StructDef -> Maybe StructDef
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
name (CheckEnv -> Map Text StructDef
envStructs CheckEnv
env) of
        Just StructDef
_ -> String -> Either String CheckEnv
forall a b. a -> Either a b
Left (String -> Either String CheckEnv)
-> String -> Either String CheckEnv
forall a b. (a -> b) -> a -> b
$ String
"Struct '" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Text -> String
DT.unpack Text
name String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"' is already defined"
        Maybe StructDef
Nothing ->
            let parsedFields :: [(Text, Type)]
parsedFields = ((Text, Text) -> (Text, Type)) -> [(Text, Text)] -> [(Text, Type)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Text
n, Text
t) -> (Text
n, Text -> Type
parseType Text
t)) [(Text, Text)]
fields
                fieldsMap :: Map Text Type
fieldsMap = [(Text, Type)] -> Map Text Type
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(Text, Type)]
parsedFields
                newDef :: StructDef
newDef = Text -> Map Text Type -> StructDef
StructDef Text
name Map Text Type
fieldsMap
                newStructs :: Map Text StructDef
newStructs = Text -> StructDef -> Map Text StructDef -> Map Text StructDef
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Text
name StructDef
newDef (CheckEnv -> Map Text StructDef
envStructs CheckEnv
env)
            in CheckEnv -> Either String CheckEnv
forall a b. b -> Either a b
Right (CheckEnv -> Either String CheckEnv)
-> CheckEnv -> Either String CheckEnv
forall a b. (a -> b) -> a -> b
$ CheckEnv
env { envStructs = newStructs }

-- | Applies assignment logic.
-- Verifies that the actual value type matches the declared variable type.
applyAssignment :: CheckEnv -> DT.Text -> Type -> Type ->
    Either String CheckEnv
applyAssignment :: CheckEnv -> Text -> Type -> Type -> Either String CheckEnv
applyAssignment CheckEnv
env Text
name Type
TyAuto Type
actual = CheckEnv -> Text -> Type -> Either String CheckEnv
insertVar CheckEnv
env Text
name Type
actual
applyAssignment CheckEnv
env Text
name Type
declared Type
actual
    | Type -> Type -> Bool
areTypesCompatible Type
declared Type
actual = CheckEnv -> Text -> Type -> Either String CheckEnv
insertVar CheckEnv
env Text
name Type
declared
    | Bool
otherwise = String -> Either String CheckEnv
forall a b. a -> Either a b
Left (String -> Either String CheckEnv)
-> String -> Either String CheckEnv
forall a b. (a -> b) -> a -> b
$ String
"Variable '" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Text -> String
DT.unpack Text
name String -> String -> String
forall a. [a] -> [a] -> [a]
++ 
                         String
"' declared as " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
typeToString Type
declared String -> String -> String
forall a. [a] -> [a] -> [a]
++ 
                         String
" but assigned " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
typeToString Type
actual

-- | Handles 'if' used as an expression (e.g. x = if c then 1 else 2).
checkIfExpr :: CheckEnv -> Ast -> Ast -> Ast -> Either String Type
checkIfExpr :: CheckEnv -> Ast -> Ast -> Ast -> Either String Type
checkIfExpr CheckEnv
env Ast
cond Ast
thenB Ast
elseB = do
    tCond <- CheckEnv -> Ast -> Either String Type
checkExpr CheckEnv
env Ast
cond
    case tCond of
        Type
TyBool -> CheckEnv -> Ast -> Ast -> Either String Type
checkBranches CheckEnv
env Ast
thenB Ast
elseB
        Type
TyInt  -> CheckEnv -> Ast -> Ast -> Either String Type
checkBranches CheckEnv
env Ast
thenB Ast
elseB
        Type
_ -> String -> Either String Type
forall a b. a -> Either a b
Left String
"Error: 'if' condition must be boolean"

-- | Verifies that both branches of an 'if' expression return compatible types.
checkBranches :: CheckEnv -> Ast -> Ast -> Either String Type
checkBranches :: CheckEnv -> Ast -> Ast -> Either String Type
checkBranches CheckEnv
env Ast
thenB Ast
elseB = do
    tThen <- CheckEnv -> Ast -> Either String Type
checkExpr CheckEnv
env Ast
thenB
    tElse <- checkExpr env elseB
    if areTypesCompatible tThen tElse
        then Right tThen
        else Left $ "Type mismatch in 'if' branches (" ++
                    typeToString tThen ++ " vs " ++
                    typeToString tElse ++ ")"

-- | Checks existence of a field within a specific struct definition.
validateStructField :: CheckEnv -> DT.Text -> DT.Text -> Either String Type
validateStructField :: CheckEnv -> Text -> Text -> Either String Type
validateStructField CheckEnv
env Text
sName Text
field = do
    def <- CheckEnv -> Text -> Either String StructDef
getStructDef CheckEnv
env Text
sName
    case Map.lookup field (structFields def) of
        Just Type
t -> Type -> Either String Type
forall a b. b -> Either a b
Right Type
t
        Maybe Type
Nothing -> String -> Either String Type
forall a b. a -> Either a b
Left (String -> Either String Type) -> String -> Either String Type
forall a b. (a -> b) -> a -> b
$ String
"Field '" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Text -> String
DT.unpack Text
field String -> String -> String
forall a. [a] -> [a] -> [a]
++ 
                          String
"' not found in struct '" String -> String -> String
forall a. [a] -> [a] -> [a]
++ 
                          Text -> String
DT.unpack Text
sName String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"'"

-- | Validates access to a struct field (e.g. object.field).
checkStructAccess :: CheckEnv -> Ast -> DT.Text -> Either String Type
checkStructAccess :: CheckEnv -> Ast -> Text -> Either String Type
checkStructAccess CheckEnv
env Ast
obj Text
field = do
    typeObj <- CheckEnv -> Ast -> Either String Type
checkExpr CheckEnv
env Ast
obj
    case typeObj of
        TyStruct Text
name -> CheckEnv -> Text -> Text -> Either String Type
validateStructField CheckEnv
env Text
name Text
field
        Type
_ -> String -> Either String Type
forall a b. a -> Either a b
Left (String -> Either String Type) -> String -> Either String Type
forall a b. (a -> b) -> a -> b
$ String
"Cannot access field '" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Text -> String
DT.unpack Text
field String -> String -> String
forall a. [a] -> [a] -> [a]
++ 
                    String
"' on non-struct type " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
typeToString Type
typeObj

-- | Validates a "new Struct {...}" instantiation.
-- Checks for extra fields, missing fields, and field types.
checkStructInstantiation :: CheckEnv -> DT.Text -> [(DT.Text, Ast)] ->
    Either String Type
checkStructInstantiation :: CheckEnv -> Text -> [(Text, Ast)] -> Either String Type
checkStructInstantiation CheckEnv
env Text
name [(Text, Ast)]
args = do
    def <- CheckEnv -> Text -> Either String StructDef
getStructDef CheckEnv
env Text
name
    let expected = StructDef -> Map Text Type
structFields StructDef
def
    let providedNames = ((Text, Ast) -> Text) -> [(Text, Ast)] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text, Ast) -> Text
forall a b. (a, b) -> a
fst [(Text, Ast)]
args
    let expectedNames = Map Text Type -> [Text]
forall k a. Map k a -> [k]
Map.keys Map Text Type
expected

    checkExtraFields name expectedNames providedNames
    checkMissingFields name expectedNames providedNames
    mapM_ (validateField env expected) args
    Right (TyStruct name)

-- | Retrieves a Struct definition from the environment.
getStructDef :: CheckEnv -> DT.Text -> Either String StructDef
getStructDef :: CheckEnv -> Text -> Either String StructDef
getStructDef CheckEnv
env Text
name = case Text -> Map Text StructDef -> Maybe StructDef
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
name (CheckEnv -> Map Text StructDef
envStructs CheckEnv
env) of
    Just StructDef
d -> StructDef -> Either String StructDef
forall a b. b -> Either a b
Right StructDef
d
    Maybe StructDef
Nothing -> String -> Either String StructDef
forall a b. a -> Either a b
Left (String -> Either String StructDef)
-> String -> Either String StructDef
forall a b. (a -> b) -> a -> b
$ String
"Error: Undefined struct '" String -> String -> String
forall a. [a] -> [a] -> [a]
++ 
                      Text -> String
DT.unpack Text
name String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"'"

-- | Errors if fields are provided that don't exist in the struct definition.
checkExtraFields :: DT.Text -> [DT.Text] -> [DT.Text] -> Either String ()
checkExtraFields :: Text -> [Text] -> [Text] -> Either String ()
checkExtraFields Text
name [Text]
expected [Text]
provided =
    let unknown :: [Text]
unknown = (Text -> Bool) -> [Text] -> [Text]
forall a. (a -> Bool) -> [a] -> [a]
filter (Text -> [Text] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Text]
expected) [Text]
provided
    in [Text] -> (Text -> Either String (ZonkAny 1)) -> Either String ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Int -> [Text] -> [Text]
forall a. Int -> [a] -> [a]
take Int
1 [Text]
unknown) ((Text -> Either String (ZonkAny 1)) -> Either String ())
-> (Text -> Either String (ZonkAny 1)) -> Either String ()
forall a b. (a -> b) -> a -> b
$ \Text
u ->
        String -> Either String (ZonkAny 1)
forall a b. a -> Either a b
Left (String -> Either String (ZonkAny 1))
-> String -> Either String (ZonkAny 1)
forall a b. (a -> b) -> a -> b
$ String
"Error: Unknown field '" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Text -> String
DT.unpack Text
u String -> String -> String
forall a. [a] -> [a] -> [a]
++ 
                  String
"' in struct '" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Text -> String
DT.unpack Text
name String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"'"

-- | Errors if mandatory fields are missing from the instantiation.
checkMissingFields :: DT.Text -> [DT.Text] -> [DT.Text] -> Either String ()
checkMissingFields :: Text -> [Text] -> [Text] -> Either String ()
checkMissingFields Text
name [Text]
expected [Text]
provided =
    let missing :: [Text]
missing = (Text -> Bool) -> [Text] -> [Text]
forall a. (a -> Bool) -> [a] -> [a]
filter (Text -> [Text] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Text]
provided) [Text]
expected
    in [Text] -> (Text -> Either String (ZonkAny 0)) -> Either String ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Int -> [Text] -> [Text]
forall a. Int -> [a] -> [a]
take Int
1 [Text]
missing) ((Text -> Either String (ZonkAny 0)) -> Either String ())
-> (Text -> Either String (ZonkAny 0)) -> Either String ()
forall a b. (a -> b) -> a -> b
$ \Text
m ->
        String -> Either String (ZonkAny 0)
forall a b. a -> Either a b
Left (String -> Either String (ZonkAny 0))
-> String -> Either String (ZonkAny 0)
forall a b. (a -> b) -> a -> b
$ String
"Error: Missing field '" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Text -> String
DT.unpack Text
m String -> String -> String
forall a. [a] -> [a] -> [a]
++ 
                  String
"' in construction of '" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Text -> String
DT.unpack Text
name String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"'"

-- | Validates the type of a single field assignment in a constructor.
validateField :: CheckEnv -> Map.Map DT.Text Type -> (DT.Text, Ast) ->
    Either String ()
validateField :: CheckEnv -> Map Text Type -> (Text, Ast) -> Either String ()
validateField CheckEnv
env Map Text Type
expectedFields (Text
fieldName, Ast
expr) = do
    actualType <- CheckEnv -> Ast -> Either String Type
checkExpr CheckEnv
env Ast
expr
    let expectedType = Map Text Type
expectedFields Map Text Type -> Text -> Type
forall k a. Ord k => Map k a -> k -> a
Map.! Text
fieldName

    unless (areTypesCompatible expectedType actualType) $
        Left $ "Error: Field '" ++ DT.unpack fieldName ++ 
               "' expected " ++ typeToString expectedType ++ 
               " but got " ++ typeToString actualType

-- | Checks if an AST node represents a character literal.
--
-- Used to distinguish between standard integers (I8, I64...) and characters (IChar)
-- within the 'AInteger' constructor.
--
-- @param ast: The AST node to inspect.
-- @return: True if the node is an AInteger wrapping an IChar, False otherwise.
isCharLiteral :: Ast -> Bool
isCharLiteral :: Ast -> Bool
isCharLiteral (AInteger (IChar Int8
_)) = Bool
True
isCharLiteral (APos Int
_ Int
_ Ast
ast) = Ast -> Bool
isCharLiteral Ast
ast
isCharLiteral Ast
_ = Bool
False

-- | Checks that all elements in a list literal have the same type AND same literal structure.
--
-- This function iterates over the list and delegates validation to 'checkElement'.
-- It determines the expected type and literal structure (Char vs Int) based on the first element.
--
-- @param env: The current semantic environment.
-- @param list: The list of AST nodes to check.
-- @return: The inferred type of the list (e.g., TyList TyInt) or an error.
checkListExpr :: CheckEnv -> [Ast] -> Either String Type
checkListExpr :: CheckEnv -> [Ast] -> Either String Type
checkListExpr CheckEnv
_ [] = Type -> Either String Type
forall a b. b -> Either a b
Right (Type -> Type
TyList Type
TyVoid)
checkListExpr CheckEnv
env (Ast
x:[Ast]
xs) = do
    expectedType <- CheckEnv -> Ast -> Either String Type
checkExpr CheckEnv
env Ast
x
    let firstIsChar = Ast -> Bool
isCharLiteral Ast
x
    forM_ xs $ \Ast
elemAst -> 
        CheckEnv -> Type -> Bool -> Ast -> Either String ()
checkElement CheckEnv
env Type
expectedType Bool
firstIsChar Ast
elemAst
    return (TyList expectedType)

-- | Validates a single list element against the expected constraints.
--
-- It performs two checks:
-- 1. Standard Type Compatibility (e.g., Int vs Int).
-- 2. Strict Literal Structure (prevents mixing 'char' literals with standard integers).
--
-- @param env: Current checking environment.
-- @param expectedType: The type derived from the first element of the list.
-- @param firstIsChar: Boolean flag indicating if the first element was a Char literal.
-- @param elemAst: The current element being validated.
checkElement :: CheckEnv -> Type -> Bool -> Ast -> Either String ()
checkElement :: CheckEnv -> Type -> Bool -> Ast -> Either String ()
checkElement CheckEnv
env Type
expectedType Bool
firstIsChar Ast
elemAst = do
    elemType <- CheckEnv -> Ast -> Either String Type
checkExpr CheckEnv
env Ast
elemAst
    unless (areTypesCompatible expectedType elemType) $
        Left $ "List type mismatch: expected " ++
               typeToString expectedType ++
               " but got " ++ typeToString elemType
    when (firstIsChar /= isCharLiteral elemAst) $
        Left $ "List literal mismatch: cannot mix 'char' and " ++
            "'int' in the same list"

-- | Checks if the AST contains a Return statement.
hasReturn :: Ast -> Bool
hasReturn :: Ast -> Bool
hasReturn (AReturn Ast
_) = Bool
True
hasReturn (ABlock [Ast]
stmts) = (Ast -> Bool) -> [Ast] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Ast -> Bool
hasReturn [Ast]
stmts
hasReturn (AList [Ast]
stmts) = (Ast -> Bool) -> [Ast] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Ast -> Bool
hasReturn [Ast]
stmts
hasReturn (AIf Ast
_ Ast
t Ast
e) = Ast -> Bool
hasReturn Ast
t Bool -> Bool -> Bool
|| Ast -> Bool
hasReturn Ast
e
hasReturn (AWhile Ast
_ Ast
body) = Ast -> Bool
hasReturn Ast
body
hasReturn (AFor Ast
_ Ast
_ Ast
_ Ast
body) = Ast -> Bool
hasReturn Ast
body
hasReturn (APos Int
_ Int
_ Ast
ast) = Ast -> Bool
hasReturn Ast
ast
hasReturn Ast
_ = Bool
False

-- | Validates that a non-void function actually returns.
checkMissingReturn :: DT.Text -> Type -> Ast -> Either String ()
checkMissingReturn :: Text -> Type -> Ast -> Either String ()
checkMissingReturn Text
_ Type
TyVoid Ast
_ = () -> Either String ()
forall a b. b -> Either a b
Right () -- Les fonctions void n'ont pas besoin de return
checkMissingReturn Text
name Type
retTy Ast
body = 
    if Ast -> Bool
hasReturn Ast
body
    then () -> Either String ()
forall a b. b -> Either a b
Right ()
    else String -> Either String ()
forall a b. a -> Either a b
Left (String -> Either String ()) -> String -> Either String ()
forall a b. (a -> b) -> a -> b
$ String
"Error: Function '" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Text -> String
DT.unpack Text
name String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"' declared as " String -> String -> String
forall a. [a] -> [a] -> [a]
++ 
                Type -> String
typeToString Type
retTy String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" but has no return statement"