{-
-- EPITECH PROJECT, 2025
-- Glados
-- File description:
-- Statement Parser (Functions, Returns, Types)
-}

{-|
Module      : Parser.Statement
Description : Parser for language statements and declarations.
Stability   : stable

This module defines parsers for the structural elements of the language:
- Variable definitions
- Function declarations
- Return statements
- Type annotations
- Code blocks
-}
module Parser.Statement (parseALL, pStatement) where

import Text.Megaparsec
import qualified Data.Text as DT
import Text.Megaparsec.Char (char)
import qualified Text.Megaparsec.Char.Lexer as L
import AST.Ast (Ast(..))
import Parser.Lexer
import Parser.Expression (pExpr)
import Data.Void (Void)
import Parser.Conditions (pIf, pWhile, pFor)
import Data.Char (ord)
import Common.Type.Integer (IntValue(..))

-- | Intermediate representation for assignment targets
-- Allows handling chains like x.y[0].z = 10
data Accessor
    = AccIndex Ast      --  Array index: [expr]
    | AccField DT.Text  --  Struct field: .name

-- | Parse a list type syntax (e.g., [int]).
--
-- Recursively parses the inner type.
pListType :: Parser DT.Text
pListType :: Parser Text
pListType = do
    _ <- Text -> Parser Text
symbol (String -> Text
DT.pack String
"[")
    innerType <- pType
    _ <- symbol (DT.pack "]")
    return (DT.pack "[" <> innerType <> DT.pack "]")

-- | Parse a type annotation.
--
-- Supports:
-- * Primitive types: int, bool, void
-- * List types: [int], [[char]]
-- * Custom types (via identifiers)
pType :: Parser DT.Text
pType :: Parser Text
pType = [Parser Text] -> Parser Text
forall (f :: * -> *) (m :: * -> *) a.
(Foldable f, Alternative m) =>
f (m a) -> m a
choice
    [ Parser Text
pListType
    , Parser Text -> Parser Text
forall a.
ParsecT Void Text Identity a -> ParsecT Void Text Identity a
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try (Text -> Parser Text
pKeyword (String -> Text
DT.pack String
"int"))
    , Parser Text -> Parser Text
forall a.
ParsecT Void Text Identity a -> ParsecT Void Text Identity a
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try (Text -> Parser Text
pKeyword (String -> Text
DT.pack String
"bool"))
    , Parser Text -> Parser Text
forall a.
ParsecT Void Text Identity a -> ParsecT Void Text Identity a
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try (Text -> Parser Text
pKeyword (String -> Text
DT.pack String
"void"))
    , Parser Text
pIdentifier
    ] Parser Text -> String -> Parser Text
forall e s (m :: * -> *) a.
MonadParsec e s m =>
m a -> String -> m a
<?> String
"type"

-- | Parse a function argument declaration (name and optional type).
--
-- Example: 
--   x: int  -> ("x", "int")
--   x       -> ("x", "auto")
pArgDeclaration :: Parser (DT.Text, DT.Text)
pArgDeclaration :: Parser (Text, Text)
pArgDeclaration = do
    name <- Parser Text
pIdentifier
    maybeType <- optional (colon >> pType)
    let finalType = case Maybe Text
maybeType of
            Just Text
t  -> Text
t
            Maybe Text
Nothing -> String -> Text
DT.pack String
"auto"
    return (name, finalType)

-- | Parse a return statement.
--
-- Example: ret 10;
pReturn :: Parser Ast
pReturn :: Parser Ast
pReturn = do
    _ <- Text -> Parser Text
pKeyword (String -> Text
DT.pack String
"ret")
    val <- pExpr
    _ <- semicolon
    return (AReturn val)

-- | Parse a block of code enclosed in braces.
--
pBlock :: Parser Ast
pBlock :: Parser Ast
pBlock = Parser Ast -> Parser Ast
forall a.
ParsecT Void Text Identity a -> ParsecT Void Text Identity a
braces (Parser Ast -> Parser Ast) -> Parser Ast -> Parser Ast
forall a b. (a -> b) -> a -> b
$ do
    stmts <- Parser Ast -> ParsecT Void Text Identity [Ast]
forall (m :: * -> *) a. MonadPlus m => m a -> m [a]
many Parser Ast
pStatement
    case stmts of
        [] -> Ast -> Parser Ast
forall a. a -> ParsecT Void Text Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return Ast
AVoid
        [Ast]
xs -> Ast -> Parser Ast
forall a. a -> ParsecT Void Text Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Ast] -> Ast
ABlock [Ast]
xs)

-- | Parse a function definition.
--
-- Syntax: func name(arg1: type, ...) -> retType { ... }
-- The return type is optional and defaults to "Void" if omitted.
pFunc :: Parser Ast
pFunc :: Parser Ast
pFunc = do
    _ <- Text -> Parser Text
pKeyword (String -> Text
DT.pack String
"func")
    name <- pIdentifier
    args <- parens (pArgDeclaration `sepBy` comma)
    maybeRetType <- optional (symbol (DT.pack "->") >> pType)
    let retType = case Maybe Text
maybeRetType of
            Just Text
t  -> Text
t
            Maybe Text
Nothing -> String -> Text
DT.pack String
"Void"
    body <- pBlock
    return (ADefineFunc name args retType body)

-- | Helper to construct a binary operator call for compound assignments.
--
-- Used for +=, -=, *=, /= to transform 'x += 1' into 'x = x + 1'.
makeOpCall :: DT.Text -> DT.Text -> Ast -> Ast
makeOpCall :: Text -> Text -> Ast -> Ast
makeOpCall Text
op Text
name Ast
expr = Ast -> [Ast] -> Ast
ACall (Text -> Ast
ASymbol Text
op) [Text -> Ast
ASymbol Text
name, Ast
expr]

-- | Parse assignment operators and return a transformation function.
--
-- Supports:
-- * Standard assignment (=) -> returns identity
-- * Compound assignment (+=, -=, *=, /=) -> returns transformation logic
pAssignOp :: DT.Text -> Parser (Ast -> Ast)
pAssignOp :: Text -> Parser (Ast -> Ast)
pAssignOp Text
name = [Parser (Ast -> Ast)] -> Parser (Ast -> Ast)
forall (f :: * -> *) (m :: * -> *) a.
(Foldable f, Alternative m) =>
f (m a) -> m a
choice
    [ (\Ast
e -> Ast
e) (Ast -> Ast) -> Parser Text -> Parser (Ast -> Ast)
forall a b.
a -> ParsecT Void Text Identity b -> ParsecT Void Text Identity a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Text -> Parser Text
symbol (String -> Text
DT.pack String
"=")
    , Text -> Text -> Ast -> Ast
makeOpCall (String -> Text
DT.pack String
"+") Text
name (Ast -> Ast) -> Parser Text -> Parser (Ast -> Ast)
forall a b.
a -> ParsecT Void Text Identity b -> ParsecT Void Text Identity a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Text -> Parser Text
symbol (String -> Text
DT.pack String
"+=")
    , Text -> Text -> Ast -> Ast
makeOpCall (String -> Text
DT.pack String
"-") Text
name (Ast -> Ast) -> Parser Text -> Parser (Ast -> Ast)
forall a b.
a -> ParsecT Void Text Identity b -> ParsecT Void Text Identity a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Text -> Parser Text
symbol (String -> Text
DT.pack String
"-=")
    , Text -> Text -> Ast -> Ast
makeOpCall (String -> Text
DT.pack String
"*") Text
name (Ast -> Ast) -> Parser Text -> Parser (Ast -> Ast)
forall a b.
a -> ParsecT Void Text Identity b -> ParsecT Void Text Identity a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Text -> Parser Text
symbol (String -> Text
DT.pack String
"*=")
    , Text -> Text -> Ast -> Ast
makeOpCall (String -> Text
DT.pack String
"div") Text
name (Ast -> Ast) -> Parser Text -> Parser (Ast -> Ast)
forall a b.
a -> ParsecT Void Text Identity b -> ParsecT Void Text Identity a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Text -> Parser Text
symbol (String -> Text
DT.pack String
"/=")
    , Text -> Text -> Ast -> Ast
makeOpCall (String -> Text
DT.pack String
"mod") Text
name (Ast -> Ast) -> Parser Text -> Parser (Ast -> Ast)
forall a b.
a -> ParsecT Void Text Identity b -> ParsecT Void Text Identity a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Text -> Parser Text
symbol (String -> Text
DT.pack String
"%=")
    ]

-- | Converts a field name (Text) into an AST list of integers.
-- This allows passing the field name as an argument to the 'attr_update' runtime function.
fieldToAst :: DT.Text -> Ast
fieldToAst :: Text -> Ast
fieldToAst Text
txt = [Ast] -> Ast
AList ([Ast] -> Ast) -> [Ast] -> Ast
forall a b. (a -> b) -> a -> b
$ (Char -> Ast) -> String -> [Ast]
forall a b. (a -> b) -> [a] -> [b]
map Char -> Ast
charToAst (Text -> String
DT.unpack Text
txt)
    where
        charToAst :: Char -> Ast
charToAst Char
c = IntValue -> Ast
AInteger (Int8 -> IntValue
IChar (Int -> Int8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Char -> Int
ord Char
c)))

-- | Recursively constructs the chain of 'update' and 'attr_update' calls.
-- This function handles nested modifications for both arrays and structures.
recursiveUpdate :: Ast -> [Accessor] -> Ast -> Ast
recursiveUpdate :: Ast -> [Accessor] -> Ast -> Ast
recursiveUpdate Ast
base [AccIndex Ast
idx] Ast
val =
    -- Case 1: End of chain on an Array (e.g. x[i] = val)
    Ast -> [Ast] -> Ast
ACall (Text -> Ast
ASymbol (String -> Text
DT.pack String
"nth_update")) [Ast
base, Ast
idx, Ast
val]

-- Case 2: End of chain on a Structure (e.g. x.field = val)
recursiveUpdate Ast
base [AccField Text
field] Ast
val =
    Ast -> [Ast] -> Ast
ACall (Text -> Ast
ASymbol (String -> Text
DT.pack String
"attr_update")) [Ast
base, Text -> Ast
fieldToAst Text
field, Ast
val]

-- Case 3: Recursion on Array (e.g. x[i]... = val)
-- We fetch the inner element using 'nth', update it recursively, and put it back using 'update'.
recursiveUpdate Ast
base (AccIndex Ast
idx : [Accessor]
rest) Ast
val =
    let inner :: Ast
inner = Ast -> [Ast] -> Ast
ACall (Text -> Ast
ASymbol (String -> Text
DT.pack String
"nth")) [Ast
base, Ast
idx]
        newVal :: Ast
newVal = Ast -> [Accessor] -> Ast -> Ast
recursiveUpdate Ast
inner [Accessor]
rest Ast
val
    in Ast -> [Ast] -> Ast
ACall (Text -> Ast
ASymbol (String -> Text
DT.pack String
"nth_update")) [Ast
base, Ast
idx, Ast
newVal]

-- Case 4: Recursion on Structure (e.g. x.field... = val)
-- We access the field, update it recursively, and put it back using 'attr_update'.
recursiveUpdate Ast
base (AccField Text
field : [Accessor]
rest) Ast
val =
    let inner :: Ast
inner = Ast -> Text -> Ast
AAccessStruct Ast
base Text
field
        newVal :: Ast
newVal = Ast -> [Accessor] -> Ast -> Ast
recursiveUpdate Ast
inner [Accessor]
rest Ast
val
    in Ast -> [Ast] -> Ast
ACall (Text -> Ast
ASymbol (String -> Text
DT.pack String
"attr_update")) 
       [Ast
base, Text -> Ast
fieldToAst Text
field, Ast
newVal]

recursiveUpdate Ast
_ [] Ast
_ = String -> Ast
forall a. HasCallStack => String -> a
error String
"Should not happen in buildUpdateChain"

-- | Entry point to construct a chain of updates for arrays and structures.
-- Transforms a complex assignment like `x.y[0] = 5` into nested function calls.
buildUpdateChain :: DT.Text -> [Accessor] -> Ast -> Ast
buildUpdateChain :: Text -> [Accessor] -> Ast -> Ast
buildUpdateChain Text
name [Accessor]
accessors Ast
finalVal =
        Ast -> [Accessor] -> Ast -> Ast
recursiveUpdate (Text -> Ast
ASymbol Text
name) [Accessor]
accessors Ast
finalVal

-- | Parse a standard variable definition or assignment.
--
-- Handles optional type annotation.
-- Example: x: int = 10; OR x = 10;
pSimpleDef :: DT.Text -> Parser Ast
pSimpleDef :: Text -> Parser Ast
pSimpleDef Text
name = do
    varType <- Parser Text -> ParsecT Void Text Identity (Maybe Text)
forall (f :: * -> *) a. Alternative f => f a -> f (Maybe a)
optional (Text -> Parser Text
symbol (String -> Text
DT.pack String
":") Parser Text -> Parser Text -> Parser Text
forall a b.
ParsecT Void Text Identity a
-> ParsecT Void Text Identity b -> ParsecT Void Text Identity b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Parser Text
pType)
    makeValue <- pAssignOp name
    val <- pExpr
    _ <- semicolon <?> "\";\" at the end of statement"
    let finalType = Text -> (Text -> Text) -> Maybe Text -> Text
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (String -> Text
DT.pack String
"auto") Text -> Text
forall a. a -> a
id Maybe Text
varType
    return (ASetVar name finalType (makeValue val))

-- | Parse a complex update (array or struct).
pComplexUpdate :: DT.Text -> [Accessor] -> Parser Ast
pComplexUpdate :: Text -> [Accessor] -> Parser Ast
pComplexUpdate Text
name [Accessor]
accessors = do
    _ <- Text -> Parser Text
symbol (String -> Text
DT.pack String
"=")
    val <- pExpr
    _ <- semicolon
    let updateExpr = Text -> [Accessor] -> Ast -> Ast
buildUpdateChain Text
name [Accessor]
accessors Ast
val
    return (ASetVar name (DT.pack "auto") updateExpr)

-- | Parser for array index accessor: [expr]
pAccIndex :: Parser Accessor
pAccIndex :: Parser Accessor
pAccIndex = do
    _ <- Text -> Parser Text
symbol (String -> Text
DT.pack String
"[")
    expr <- pExpr
    _ <- symbol (DT.pack "]")
    return (AccIndex expr)

-- | Parser for struct field accessor: .identifier
pAccField :: Parser Accessor
pAccField :: Parser Accessor
pAccField = do
    _ <- Text -> Parser Text
symbol (String -> Text
DT.pack String
".")
    field <- pIdentifier
    return (AccField field)

-- | Parse a variable definition or assignment.
-- Handles mixed chains: x = 1, x[0] = 1, x.y = 1, x.y[0].z = 1
pVarDef :: Parser Ast
pVarDef :: Parser Ast
pVarDef = do
    name <- Parser Text
pIdentifier
    -- Parse mixed chain of [indices] and .fields
    accessors <- many (pAccIndex <|> pAccField)
    if null accessors
        then pSimpleDef name
        else pComplexUpdate name accessors

-- | Parse a single field definition within a structure.
--
-- Syntax: fieldName: type;
pStructField :: Parser (DT.Text, DT.Text)
pStructField :: Parser (Text, Text)
pStructField = do
    name <- Parser Text
pIdentifier
    _ <- colon
    fType <- pType
    _ <- semicolon
    return (name, fType)

-- | Parse a structure definition.
--
-- Syntax: struct Name { fields... }
pStruct :: Parser Ast
pStruct :: Parser Ast
pStruct = do
    _ <- Text -> Parser Text
pKeyword (String -> Text
DT.pack String
"struct")
    name <- pIdentifier
    fields <- braces (many pStructField)
    return (ADefineStruct name fields)

-- | Parse the file path string for an import.
--
-- Expects a string enclosed in double quotes.
pImportPath :: Parser DT.Text
pImportPath :: Parser Text
pImportPath = Parser Text -> Parser Text
forall a.
ParsecT Void Text Identity a -> ParsecT Void Text Identity a
lexeme (Parser Text -> Parser Text) -> Parser Text -> Parser Text
forall a b. (a -> b) -> a -> b
$ do
    _ <- Token Text -> ParsecT Void Text Identity (Token Text)
forall e s (m :: * -> *).
(MonadParsec e s m, Token s ~ Char) =>
Token s -> m (Token s)
char Char
Token Text
'"'
    content <- manyTill L.charLiteral (char '"')
    return (DT.pack content)

-- | Parse an import directive.
--
-- Syntax: import "path/to/file";
pImport :: Parser Ast
pImport :: Parser Ast
pImport = do
    _ <- Text -> Parser Text
pKeyword (String -> Text
DT.pack String
"import")
    path <- pImportPath
    _ <- semicolon
    return (AImport path)

-- | Group of control flow parsers (If, While, For).
pControlFlow :: [Parser Ast]
pControlFlow :: [Parser Ast]
pControlFlow =
    [ Parser Ast -> Parser Ast
forall a.
ParsecT Void Text Identity a -> ParsecT Void Text Identity a
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try (Parser Ast -> Parser Ast -> Parser Ast
pIf Parser Ast
pVarDef Parser Ast
pBlock)
    , Parser Ast -> Parser Ast
forall a.
ParsecT Void Text Identity a -> ParsecT Void Text Identity a
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try (Parser Ast -> Parser Ast
pWhile Parser Ast
pBlock)
    , Parser Ast -> Parser Ast
forall a.
ParsecT Void Text Identity a -> ParsecT Void Text Identity a
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try (Parser Ast -> Parser Ast -> Parser Ast
pFor Parser Ast
pVarDef Parser Ast
pBlock)
    ]

-- | Group of top-level declaration parsers (Import, Struct, Func).
pDeclarations :: [Parser Ast]
pDeclarations :: [Parser Ast]
pDeclarations =
    [ Parser Ast
pImport
    , Parser Ast
pStruct
    , Parser Ast
pFunc
    ]

-- | Parses a variable definition or assignment and wraps it in an expression statement.
--
-- This ensures that variable-related instructions are stored as 'AExprStmt' nodes
-- at the top level, maintaining a consistent instruction structure.
pVarDefStmt :: Parser Ast
pVarDefStmt :: Parser Ast
pVarDefStmt = do
    stmt <- Parser Ast
pVarDef
    return (AExprStmt stmt)

-- | Parses an expression followed by a mandatory semicolon.
--
-- Example: 1 + 1;
-- It wraps the expression in an 'AExprStmt' node to distinguish 
-- standalone expressions from other types of statements.
pExprStatement :: Parser Ast
pExprStatement :: Parser Ast
pExprStatement = do
    expr <- Parser Ast
pExpr
    _ <- semicolon <?> "\";\" after expression"
    return (AExprStmt expr)

-- | Parses a 'break' statement.
--
-- Syntax: break;
--
-- It wraps the 'ABreak' node in an AExprStmt to ensure consistency
-- with other statements that are expressions (like function calls).
-- This allows the 'break' to be treated as a standard instruction
-- in the AST processing pipeline.
pBreak :: Parser Ast
pBreak :: Parser Ast
pBreak = do
    _ <- Text -> Parser Text
pKeyword (String -> Text
DT.pack String
"break")
    _ <- semicolon
    return (AExprStmt ABreak)

-- | Group of basic statement parsers (Return, Variable, Expression).
pBasic :: [Parser Ast]
pBasic :: [Parser Ast]
pBasic =
    [ Parser Ast
pReturn
    , Parser Ast
pBreak
    , Parser Ast -> Parser Ast
forall a.
ParsecT Void Text Identity a -> ParsecT Void Text Identity a
forall e s (m :: * -> *) a. MonadParsec e s m => m a -> m a
try Parser Ast
pVarDefStmt
    , Parser Ast
pExprStatement
    ]

-- | Main statement parser.
--
-- Aggregates all statement types (control flow, declarations, basic instructions)
-- into a single choice. This is the top-level parser for a line of code.
pStatement :: Parser Ast
pStatement :: Parser Ast
pStatement = Parser Ast -> Parser Ast
withPos (Parser Ast -> Parser Ast) -> Parser Ast -> Parser Ast
forall a b. (a -> b) -> a -> b
$ [Parser Ast] -> Parser Ast
forall (f :: * -> *) (m :: * -> *) a.
(Foldable f, Alternative m) =>
f (m a) -> m a
choice ([Parser Ast]
pControlFlow [Parser Ast] -> [Parser Ast] -> [Parser Ast]
forall a. [a] -> [a] -> [a]
++ [Parser Ast]
pDeclarations [Parser Ast] -> [Parser Ast] -> [Parser Ast]
forall a. [a] -> [a] -> [a]
++ [Parser Ast]
pBasic)

-- | Main entry point for the parser.
--
-- Parses a list of statements from the input text until EOF.
-- The filename argument is used for error reporting.
parseALL :: DT.Text -> Either (ParseErrorBundle DT.Text Void) [Ast]
parseALL :: Text -> Either (ParseErrorBundle Text Void) [Ast]
parseALL = ParsecT Void Text Identity [Ast]
-> String -> Text -> Either (ParseErrorBundle Text Void) [Ast]
forall e s a.
Parsec e s a -> String -> s -> Either (ParseErrorBundle s e) a
parse (Parser ()
sc Parser ()
-> ParsecT Void Text Identity [Ast]
-> ParsecT Void Text Identity [Ast]
forall a b.
ParsecT Void Text Identity a
-> ParsecT Void Text Identity b -> ParsecT Void Text Identity b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Parser Ast -> ParsecT Void Text Identity [Ast]
forall (m :: * -> *) a. MonadPlus m => m a -> m [a]
many Parser Ast
pStatement ParsecT Void Text Identity [Ast]
-> Parser () -> ParsecT Void Text Identity [Ast]
forall a b.
ParsecT Void Text Identity a
-> ParsecT Void Text Identity b -> ParsecT Void Text Identity a
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* Parser ()
forall e s (m :: * -> *). MonadParsec e s m => m ()
eof) String
"ParseALL"