module Language.Elna.Glue
    ( glue
    ) where

import Control.Monad.Trans.State (State, gets, modify', runState)
import Data.Bifunctor (Bifunctor(..))
import Data.Foldable (Foldable(..), traverse_)
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import Data.Maybe (catMaybes)
import Data.Vector (Vector)
import qualified Data.Text.Lazy.Builder.Int as Text.Builder
import qualified Data.Text.Lazy.Builder as Text.Builder
import qualified Data.Text.Lazy as Text.Lazy
import qualified Data.Vector as Vector
import Data.Word (Word32)
import qualified Language.Elna.Frontend.AST as AST
import Language.Elna.Frontend.Types (Type(..))
import Language.Elna.Backend.Intermediate
    ( Label(..)
    , Operand(..)
    , Quadruple(..)
    , Variable(..)
    )
import Language.Elna.Frontend.SymbolTable (Info(..), SymbolTable)
import qualified Language.Elna.Frontend.SymbolTable as SymbolTable
import GHC.Records (HasField(..))
import Language.Elna.Frontend.AST (Identifier(..))
import Debug.Trace (traceShow)

data Paste = Paste
    { temporaryCounter :: Word32
    , labelCounter :: Word32
    , localMap :: HashMap Identifier Variable
    }

newtype Glue a = Glue
    { runGlue :: State Paste a }

instance Functor Glue
  where
    fmap f (Glue x) = Glue $ f <$> x

instance Applicative Glue
  where
    pure = Glue . pure
    (Glue f) <*> (Glue x) = Glue $ f <*> x

instance Monad Glue
  where
    (Glue x) >>= f = Glue $ x >>= (runGlue . f)

glue :: SymbolTable -> AST.Program -> HashMap Identifier (Vector (Quadruple Variable))
glue globalTable
    = fst
    . flip runState emptyPaste
    . runGlue
    . program globalTable
  where
    emptyPaste = Paste
        { temporaryCounter = 0
        , labelCounter = 0
        , localMap = mempty
        }

program :: SymbolTable -> AST.Program -> Glue (HashMap Identifier (Vector (Quadruple Variable)))
program globalTable (AST.Program declarations)
    = HashMap.fromList . catMaybes
    <$> traverse (declaration globalTable) declarations

declaration
    :: SymbolTable
    -> AST.Declaration
    -> Glue (Maybe (AST.Identifier, Vector (Quadruple Variable)))
declaration globalTable (AST.ProcedureDeclaration procedureName parameters variableDeclarations statements) =
    let Just (ProcedureInfo localTable _) = SymbolTable.lookup procedureName globalTable
     in Glue (modify' resetTemporaryCounter)
        >> traverseWithIndex registerVariable variableDeclarations
        >> traverseWithIndex registerParameter (reverse parameters)
        >> nameQuadruplesTuple <$> traverse (statement localTable) statements
  where
    traverseWithIndex f = traverse_ (uncurry f) . zip [0..]
    registerParameter index (AST.Parameter identifier _ _) =
        Glue $ modify' $ modifier identifier $ ParameterVariable index
    registerVariable index (AST.VariableDeclaration identifier _) =
        Glue $ modify' $ modifier identifier $ LocalVariable index
    modifier identifier currentCounter generator = generator
        { localMap = HashMap.insert identifier currentCounter
            $ getField @"localMap" generator
        }
    nameQuadruplesTuple quadrupleList = Just
        ( procedureName
        , Vector.cons StartQuadruple
            $ flip Vector.snoc StopQuadruple
            $ fold quadrupleList
        )
    resetTemporaryCounter paste = paste
        { temporaryCounter = 0
        , localMap = mempty
        }
declaration _ (AST.TypeDefinition _ _) = pure Nothing

statement :: SymbolTable -> AST.Statement -> Glue (Vector (Quadruple Variable))
statement _ AST.EmptyStatement = pure mempty
statement localTable (AST.CallStatement (AST.Identifier callName) arguments) = do
    visitedArguments <- traverse (expression localTable) arguments
    let (parameterStatements, argumentStatements)
            = bimap (Vector.fromList . fmap ParameterQuadruple) Vector.concat
            $ unzip visitedArguments
     in pure
        $ Vector.snoc (argumentStatements <> parameterStatements)
        $ CallQuadruple callName
        $ fromIntegral
        $ length arguments
statement localTable (AST.CompoundStatement statements) =
    fold <$> traverse (statement localTable) statements
statement localTable (AST.IfStatement ifCondition ifStatement elseStatement) = do
    (conditionStatements, jumpConstructor) <- condition localTable ifCondition
    ifLabel <- createLabel
    endLabel <- createLabel
    ifStatements <- statement localTable ifStatement
    possibleElseStatements <- traverse (statement localTable) elseStatement
    pure $ conditionStatements <> case possibleElseStatements of
        Just elseStatements -> Vector.cons (jumpConstructor ifLabel) elseStatements
            <> Vector.fromList [GoToQuadruple endLabel, LabelQuadruple ifLabel]
            <> Vector.snoc ifStatements (LabelQuadruple endLabel)
        Nothing -> Vector.fromList [jumpConstructor ifLabel, GoToQuadruple endLabel, LabelQuadruple ifLabel]
            <> Vector.snoc ifStatements (LabelQuadruple endLabel)
statement localTable (AST.AssignmentStatement variableAccess' assignee) = do
    (rhsOperand, rhsStatements) <- expression localTable assignee
    let variableType' = variableType variableAccess' localTable
    accessResult <- variableAccess localTable variableAccess' Nothing variableType' mempty
    lhsStatements <- case accessResult of
            (identifier, Just accumulatedIndex, accumulatedStatements)
                -> Vector.snoc accumulatedStatements
                . ArrayAssignQuadruple rhsOperand accumulatedIndex
                <$> lookupLocal identifier
            (identifier, Nothing, accumulatedStatements)
                -> Vector.snoc accumulatedStatements
                . AssignQuadruple rhsOperand
                <$> lookupLocal identifier
    pure $ rhsStatements <> lhsStatements
statement localTable (AST.WhileStatement whileCondition whileStatement) = do
    (conditionStatements, jumpConstructor) <- condition localTable whileCondition
    startLabel <- createLabel
    endLabel <- createLabel
    conditionLabel <- createLabel
    whileStatements <- statement localTable whileStatement
    pure $ Vector.fromList [LabelQuadruple conditionLabel]
        <> conditionStatements
        <> Vector.fromList [jumpConstructor startLabel, GoToQuadruple endLabel, LabelQuadruple startLabel]
        <> whileStatements
        <> Vector.fromList [GoToQuadruple conditionLabel, LabelQuadruple endLabel]

createTemporary :: Glue Variable
createTemporary = do
    currentCounter <- Glue $ gets $ getField @"temporaryCounter"
    Glue $ modify' modifier
    pure $ TempVariable currentCounter
  where
    modifier generator = generator
        { temporaryCounter = getField @"temporaryCounter" generator + 1
        }

lookupLocal :: Identifier -> Glue Variable
lookupLocal identifier =
    fmap (HashMap.! identifier) $ Glue $ gets $ getField @"localMap"

createLabel :: Glue Label
createLabel = do
    currentCounter <- Glue $ gets $ getField @"labelCounter"
    Glue $ modify' modifier
    pure $ Label
        $ Text.Lazy.toStrict
        $ Text.Builder.toLazyText
        $ ".L" <> Text.Builder.decimal currentCounter
  where
    modifier generator = generator
        { labelCounter = getField @"labelCounter" generator + 1
        }

condition
    :: SymbolTable
    -> AST.Condition
    -> Glue (Vector (Quadruple Variable), Label -> Quadruple Variable)
condition localTable (AST.EqualCondition lhs rhs) = do
    (lhsOperand, lhsStatements) <- expression localTable lhs
    (rhsOperand, rhsStatements) <- expression localTable rhs
    pure
        ( lhsStatements <> rhsStatements
        , EqualQuadruple lhsOperand rhsOperand
        )
condition localTable (AST.NonEqualCondition lhs rhs) = do
    (lhsOperand, lhsStatements) <- expression localTable lhs
    (rhsOperand, rhsStatements) <- expression localTable rhs
    pure
        ( lhsStatements <> rhsStatements
        , NonEqualQuadruple lhsOperand rhsOperand
        )
condition localTable (AST.LessCondition lhs rhs) = do
    (lhsOperand, lhsStatements) <- expression localTable lhs
    (rhsOperand, rhsStatements) <- expression localTable rhs
    pure (lhsStatements <> rhsStatements, LessQuadruple lhsOperand rhsOperand)
condition localTable (AST.GreaterCondition lhs rhs) = do
    (lhsOperand, lhsStatements) <- expression localTable lhs
    (rhsOperand, rhsStatements) <- expression localTable rhs
    pure
        ( lhsStatements <> rhsStatements
        , GreaterQuadruple lhsOperand rhsOperand
        )
condition localTable (AST.LessOrEqualCondition lhs rhs) = do
    (lhsOperand, lhsStatements) <- expression localTable lhs
    (rhsOperand, rhsStatements) <- expression localTable rhs
    pure
        ( lhsStatements <> rhsStatements
        , LessOrEqualQuadruple lhsOperand rhsOperand
        )
condition localTable (AST.GreaterOrEqualCondition lhs rhs) = do
    (lhsOperand, lhsStatements) <- expression localTable lhs
    (rhsOperand, rhsStatements) <- expression localTable rhs
    pure
        ( lhsStatements <> rhsStatements
        , GreaterOrEqualQuadruple lhsOperand rhsOperand
        )

variableAccess
    :: SymbolTable
    -> AST.VariableAccess
    -> Maybe (Operand Variable)
    -> Type
    -> Vector (Quadruple Variable)
    -> Glue (AST.Identifier, Maybe (Operand Variable), Vector (Quadruple Variable))
variableAccess _ (AST.VariableAccess identifier) accumulatedIndex _ accumulatedStatements =
    pure (identifier, accumulatedIndex, accumulatedStatements)
variableAccess localTable accessKind accumulatedIndex arrayType statements
    | (AST.ArrayAccess access1 index1) <- accessKind
    , (ArrayType arraySize baseType) <- arrayType = do
        (indexPlace, statements') <- expression localTable index1
        case accumulatedIndex  of
            Just baseIndex -> do
                resultVariable <- createTemporary
                let resultOperand = VariableOperand resultVariable
                    indexCalculation = Vector.fromList
                        [ ProductQuadruple (IntOperand $ fromIntegral arraySize) baseIndex resultVariable
                        , AddQuadruple indexPlace resultOperand resultVariable
                        ]
                 in variableAccess localTable access1 (Just resultOperand) baseType
                    $ statements <> indexCalculation <> statements'
            Nothing ->
                variableAccess localTable access1 (Just indexPlace) baseType statements'
variableAccess _ _ _ _ _ = error "Array access operator doesn't match the type."

variableType :: AST.VariableAccess -> SymbolTable -> Type
variableType (AST.VariableAccess identifier) symbolTable
    | Just (TypeInfo type') <- SymbolTable.lookup identifier symbolTable = type'
    | Just (VariableInfo _ type') <- SymbolTable.lookup identifier symbolTable = type'
    | otherwise = traceShow identifier $ error "Undefined type."
variableType (AST.ArrayAccess arrayAccess' _) symbolTable =
    variableType arrayAccess' symbolTable

expression :: SymbolTable -> AST.Expression -> Glue (Operand Variable, Vector (Quadruple Variable))
expression localTable = \case
    (AST.LiteralExpression literal') -> pure (literal literal', mempty)
    (AST.SumExpression lhs rhs) -> binaryExpression AddQuadruple lhs rhs
    (AST.SubtractionExpression lhs rhs) ->
        binaryExpression SubtractionQuadruple lhs rhs
    (AST.NegationExpression negation) -> do
        (operand, statements) <- expression localTable negation
        tempVariable <- createTemporary
        let negationQuadruple = NegationQuadruple operand tempVariable
        pure
            ( VariableOperand tempVariable
            , Vector.snoc statements negationQuadruple
            )
    (AST.ProductExpression lhs rhs) ->
        binaryExpression ProductQuadruple lhs rhs
    (AST.DivisionExpression lhs rhs) ->
        binaryExpression DivisionQuadruple lhs rhs
    (AST.VariableExpression variableExpression) -> do
        let variableType' = variableType variableExpression localTable
        variableAccess' <- variableAccess localTable variableExpression Nothing variableType' mempty
        case variableAccess' of
            (identifier, Nothing, statements)
                -> (, statements) . VariableOperand 
                <$> lookupLocal identifier
            (identifier, Just operand, statements) -> do
                arrayAddress <- createTemporary
                localVariable <- lookupLocal identifier
                let arrayStatement = ArrayQuadruple localVariable operand arrayAddress
                pure
                    ( VariableOperand arrayAddress
                    , Vector.snoc statements arrayStatement
                    )
  where
    binaryExpression f lhs rhs = do
        (lhsOperand, lhsStatements) <- expression localTable lhs
        (rhsOperand, rhsStatements) <- expression localTable rhs
        tempVariable <- createTemporary
        let newQuadruple = f lhsOperand rhsOperand tempVariable
        pure
            ( VariableOperand tempVariable
            , Vector.snoc (lhsStatements <> rhsStatements) newQuadruple
            )

literal :: AST.Literal -> Operand Variable
literal (AST.DecimalLiteral integer) = IntOperand integer
literal (AST.HexadecimalLiteral integer) = IntOperand integer
literal (AST.CharacterLiteral character) = IntOperand $ fromIntegral character