module Codegen(module Codegen, A.Module) where

import qualified Data.Map.Strict as Map
-- import qualified LLVM.General.AST.Type as A
-- import qualified LLVM.General.AST.Global as A
-- import qualified LLVM.General.AST.Constant as A.C
-- import qualified LLVM.General.AST.Operand as A
-- import qualified LLVM.General.AST.Name as A
-- import qualified LLVM.General.AST.Instruction as A
import qualified LLVM.General.AST as A

import AST


type Error a = Either String a


codegen :: Program  -- Program to compile
        -> String   -- Module name
        -> String   -- File name of source
        -> Error A.Module
codegen prog name fname = do
    defs <- generateDefs (preprocess prog)
    return $ A.defaultModule {
            A.moduleName = name,
            A.moduleSourceFileName = fname,
            A.moduleDefinitions = defs
        }

preprocess :: Program -> Program
preprocess prog@(Program decls) = mapProgram' filtered mapper
    where
        filtered = Program $ filter notTypedef decls
        mapper = defaultPM' {typeHandler' = typeReplacer (findTypeRenames prog)}

        notTypedef :: Declaration -> Bool
        notTypedef (DecTypedef _ _) = False
        notTypedef _ = True

        typeReplacer :: Map.Map Name Type -> Type -> Type
        typeReplacer m t@(TypeName n) = maybe t id $ Map.lookup n m
        typeReplacer _ t = t

        findTypeRenames :: Program -> Map.Map Name Type
        findTypeRenames (Program d) = foldl go Map.empty d
            where
                go :: Map.Map Name Type -> Declaration -> Map.Map Name Type
                go m (DecTypedef t n) = Map.insert n t m
                go m _ = m


generateDefs :: Program -> Error [A.Definition]
generateDefs prog = do
    checkUndefinedTypes prog
    checkUndefinedVars prog
    fail "TODO"
    return []

checkUndefinedTypes :: Program -> Error ()
checkUndefinedTypes prog = fmap (const ()) $ mapProgram prog $ defaultPM {typeHandler = check}
    where
        check :: Type -> Error Type
        check (TypeName n) = Left $ "Undefined type name '" ++ n ++ "'"
        check t = Right t

-- checkUndefinedVars :: Program -> Error ()
-- checkUndefinedVars prog = do


-- mapTypes' :: Program -> (Type -> Type) -> Program
-- mapTypes' prog f = (\(Right res) -> res) $ mapTypes prog (return . f)

-- mapTypes :: Program -> (Type -> Error Type) -> Error Program
-- mapTypes (Program decls) f = Program <$> sequence (map goD decls)
--     where
--         handler :: Type -> Error Type
--         handler (TypePtr t) = f t >>= f . TypePtr
--         handler t = f t

--         goD :: Declaration -> Error Declaration
--         goD (DecFunction t n a b) = do
--             rt <- handler t
--             ra <- sequence $ map (\(at,an) -> (\art -> (art,an)) <$> handler at) a
--             rb <- goB b
--             return $ DecFunction rt n ra rb
--         goD (DecVariable t n v) = (\rt -> DecVariable rt n v) <$> handler t
--         goD (DecTypedef t n) = (\rt -> DecTypedef rt n) <$> handler t

--         goB :: Block -> Error Block
--         goB (Block stmts) = Block <$> sequence (map goS stmts)

--         goS :: Statement -> Error Statement
--         goS (StBlock bl) = StBlock <$> goB bl
--         goS (StVarDeclaration t n e) = (\rt -> StVarDeclaration rt n e) <$> handler t
--         goS (StIf c t e) = do
--             rt <- goS t
--             re <- goS e
--             return $ StIf c rt re
--         goS (StWhile c b) = StWhile c <$> goS b
--         goS s = return s


type MapperHandler a = a -> Error a

data ProgramMapper = ProgramMapper
    {declarationHandler :: MapperHandler Declaration
    ,blockHandler :: MapperHandler Block
    ,typeHandler :: MapperHandler Type
    ,literalHandler :: MapperHandler Literal
    ,binOpHandler :: MapperHandler BinaryOperator
    ,unOpHandler :: MapperHandler UnaryOperator
    ,expressionHandler :: MapperHandler Expression
    ,statementHandler :: MapperHandler Statement
    ,nameHandler :: MapperHandler Name}

type MapperHandler' a = a -> a

data ProgramMapper' = ProgramMapper'
    {declarationHandler' :: MapperHandler' Declaration
    ,blockHandler' :: MapperHandler' Block
    ,typeHandler' :: MapperHandler' Type
    ,literalHandler' :: MapperHandler' Literal
    ,binOpHandler' :: MapperHandler' BinaryOperator
    ,unOpHandler' :: MapperHandler' UnaryOperator
    ,expressionHandler' :: MapperHandler' Expression
    ,statementHandler' :: MapperHandler' Statement
    ,nameHandler' :: MapperHandler' Name}

defaultPM :: ProgramMapper
defaultPM = ProgramMapper return return return return return return return return return

defaultPM' :: ProgramMapper'
defaultPM' = ProgramMapper' id id id id id id id id id

mapProgram' :: Program -> ProgramMapper' -> Program
mapProgram' prog mapper = (\(Right r) -> r) $ mapProgram prog $ ProgramMapper
    {declarationHandler = return . declarationHandler' mapper
    ,blockHandler = return . blockHandler' mapper
    ,typeHandler = return . typeHandler' mapper
    ,literalHandler = return . literalHandler' mapper
    ,binOpHandler = return . binOpHandler' mapper
    ,unOpHandler = return . unOpHandler' mapper
    ,expressionHandler = return . expressionHandler' mapper
    ,statementHandler = return . statementHandler' mapper
    ,nameHandler = return . nameHandler' mapper}

mapProgram :: Program -> ProgramMapper -> Error Program
mapProgram prog mapper = goP prog
    where
        h_d = declarationHandler mapper
        h_b = blockHandler mapper
        h_t = typeHandler mapper
        h_l = literalHandler mapper
        h_bo = binOpHandler mapper
        h_uo = unOpHandler mapper
        h_e = expressionHandler mapper
        h_s = statementHandler mapper
        h_n = nameHandler mapper

        goP :: Program -> Error Program
        goP (Program decls) = Program <$> sequence (map (\d -> goD d >>= h_d) decls)

        goD :: Declaration -> Error Declaration
        goD (DecFunction t n a b) = do
            rt <- goT t
            rn <- goN n
            ra <- sequence $ map (\(at,an) -> (,) <$> goT at <*> goN an) a
            rb <- goB b
            h_d $ DecFunction rt rn ra rb
        goD (DecVariable t n mv) = do
            rt <- goT t
            rn <- goN n
            rmv <- sequence $ fmap goE mv
            h_d $ DecVariable rt rn rmv
        goD (DecTypedef t n) = do
            rt <- goT t
            rn <- goN n
            h_d $ DecTypedef rt rn

        goT :: Type -> Error Type
        goT (TypePtr t) = goT t >>= (h_t . TypePtr)
        goT (TypeName n) = goN n >>= (h_t . TypeName)
        goT t = h_t t

        goN :: Name -> Error Name
        goN = h_n

        goB :: Block -> Error Block
        goB (Block sts) = (Block <$> sequence (map goS sts)) >>= h_b

        goE :: Expression -> Error Expression
        goE (ExLit l) = goL l >>= (h_e . ExLit)
        goE (ExBinOp bo e1 e2) = do
            rbo <- goBO bo
            re1 <- goE e1
            re2 <- goE e2
            h_e $ ExBinOp rbo re1 re2
        goE (ExUnOp uo e) = do
            ruo <- goUO uo
            re <- goE e
            h_e $ ExUnOp ruo re

        goS :: Statement -> Error Statement
        goS StEmpty = h_s StEmpty
        goS (StBlock b) = goB b >>= (h_s . StBlock)
        goS (StExpr e) = goE e >>= (h_s . StExpr)
        goS (StVarDeclaration t n me) = do
            rt <- goT t
            rn <- goN n
            rme <- sequence $ fmap goE me
            h_s $ StVarDeclaration rt rn rme
        goS (StAssignment n e) = do
            rn <- goN n
            re <- goE e
            h_s $ StAssignment rn re
        goS (StIf e s1 s2) = do
            re <- goE e
            rs1 <- goS s1
            rs2 <- goS s2
            h_s $ StIf re rs1 rs2
        goS (StWhile e s) = do
            re <- goE e
            rs <- goS s
            h_s $ StWhile re rs
        goS (StReturn e) = goE e >>= (h_s . StReturn)

        goL :: Literal -> Error Literal
        goL (LitVar n) = goN n >>= (h_l . LitVar)
        goL (LitCall n a) = do
            rn <- goN n
            ra <- sequence $ map goE a
            h_l $ LitCall rn ra

        goBO :: BinaryOperator -> Error BinaryOperator
        goBO = h_bo

        goUO :: UnaryOperator -> Error UnaryOperator
        goUO = h_uo