{-# LANGUAGE TupleSections, GeneralizedNewtypeDeriving #-}
module Codegen(codegen) where

import Control.Monad.State.Strict
import Control.Monad.Except
import Data.Maybe
import qualified Data.Map.Strict as Map
import qualified LLVM.General.AST.Type as A
import qualified LLVM.General.AST.Global as A.G
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 Debug.Trace

import AST


type Error a = Either String a

type LLName = String


data GenState
    = GenState {currentBlock :: Maybe LLName
               ,allBlocks :: Map.Map LLName A.BasicBlock
               ,nextId :: Integer
               ,definitions :: [A.Definition]
               ,variables :: Map.Map LLName LLName}
    deriving (Show)

initialGenState :: GenState
initialGenState
    = GenState {currentBlock = Nothing
               ,allBlocks = Map.empty
               ,nextId = 1
               ,definitions = []
               ,variables = Map.empty}

newtype CGMonad a = CGMonad {unMon :: ExceptT String (State GenState) a}
    deriving (Functor, Applicative, Monad, MonadState GenState, MonadError String)

runCGMonad :: CGMonad a -> Error (a, GenState)
runCGMonad m = let (e, s) = runState (runExceptT (unMon m)) initialGenState
               in either Left (\x -> Right (x, s)) e

getUniqueId :: CGMonad Integer
getUniqueId = state $ \s -> (nextId s, s {nextId = nextId s + 1})

getNewName :: String -> CGMonad String
getNewName base = fmap ((base++) . show) getUniqueId

newBlock :: CGMonad LLName
newBlock = do
    name <- getNewName "bb"
    state $ \s -> (name, s {
            currentBlock = Just name,
            allBlocks = Map.insert name (A.BasicBlock (A.Name name) [] undefined) $ allBlocks s
        })

newBlockJump :: LLName -> CGMonad LLName
newBlockJump next = do
    bb <- newBlock
    setTerminator $ A.Br (A.Name next) []
    return bb

changeBlock :: LLName -> CGMonad ()
changeBlock name = state $ \s -> ((), s {currentBlock = Just name})

addInstr :: A.Instruction -> CGMonad LLName
addInstr instr = do
    name <- getNewName "t"
    addNamedInstr $ A.Name name A.:= instr

addNamedInstr :: A.Named A.Instruction -> CGMonad LLName
addNamedInstr instr@(A.Name name A.:= _) = do
    let append (A.BasicBlock n il t) = A.BasicBlock n (il ++ [instr]) t
    state $ \s -> (name, s {allBlocks = Map.adjust append (fromJust (currentBlock s)) (allBlocks s)})
addNamedInstr _ = undefined

addNamedInstrList :: [A.Named A.Instruction] -> CGMonad LLName
addNamedInstrList l = mapM addNamedInstr l >>= return . last

setTerminator :: A.Terminator -> CGMonad ()
setTerminator term = do
    let replace (A.BasicBlock n il _) = A.BasicBlock n il (A.Do term)
    state $ \s -> ((), s {allBlocks = Map.adjust replace (fromJust (currentBlock s)) (allBlocks s)})

setVarLabel :: LLName -> LLName -> CGMonad ()
setVarLabel name label = do
    state $ \s -> ((), s {variables = Map.insert name label $ variables s})

lookupVar :: LLName -> CGMonad LLName
lookupVar name = liftM (fromJust . Map.lookup name . variables) get


namedName :: A.Named a -> LLName
namedName (A.Name name A.:= _) = name
namedName _ = undefined


codegen :: Program  -- Program to compile
        -> String   -- Module name
        -> String   -- File name of source
        -> Error A.Module
codegen prog name fname = do
    (defs, st) <- runCGMonad $ do
                    defs <- generateDefs prog
                    return defs

    traceShow st $ return ()

    return $ A.defaultModule {
            A.moduleName = name,
            A.moduleSourceFileName = fname,
            A.moduleDefinitions = defs
        }


generateDefs :: Program -> CGMonad [A.Definition]
generateDefs prog = do
    vardecls <- genGlobalVars prog
    fundecls <- genFunctions prog
    return $ vardecls ++ fundecls

genGlobalVars :: Program -> CGMonad [A.Definition]
genGlobalVars (Program decs) = mapM gen $ filter isDecVariable decs
  where
    gen :: Declaration -> CGMonad A.Definition
    gen (DecVariable t n Nothing) = return $ A.GlobalDefinition $
            A.globalVariableDefaults {
                A.G.name = A.Name n,
                A.G.type' = toLLVMType t,
                A.G.initializer = Just $ initializerFor t
            }
    gen (DecVariable _ _ (Just _)) = throwError $ "Initialised global variables not supported yet"
    gen _ = undefined

genFunctions :: Program -> CGMonad [A.Definition]
genFunctions (Program decs) = mapM gen $ filter isDecFunction decs
  where
    gen :: Declaration -> CGMonad A.Definition
    gen (DecFunction rettype name args body) = do
        firstbb <- genBlock' body
        blockmap <- liftM allBlocks get
        let bbs' = map snd $ filter (\x -> fst x /= firstbb) $ Map.toList blockmap
            bbs = fromJust (Map.lookup firstbb blockmap) : bbs'
        return $ A.GlobalDefinition $ A.functionDefaults {
                A.G.returnType = toLLVMType rettype,
                A.G.name = A.Name name,
                A.G.parameters = ([A.Parameter (toLLVMType t) (A.Name n) [] | (t,n) <- args], False),
                A.G.basicBlocks = bbs
            }
    gen _ = undefined



genBlock' :: Block -> CGMonad LLName
genBlock' bl = do
    termbb <- newBlock
    setTerminator $ A.Unreachable []
    genBlock bl termbb

genBlock :: Block
         -> LLName  -- name of BasicBlock following this Block
         -> CGMonad LLName  -- name of first BasicBlock
genBlock (Block []) following = genBlock (Block [StEmpty]) following
genBlock (Block [stmt]) following = do
    firstbb <- genSingle stmt following
    return firstbb
genBlock (Block (stmt:rest)) following = do
    restbb <- genBlock (Block rest) following
    firstbb <- genSingle stmt restbb
    return firstbb

genSingle :: Statement
          -> LLName  -- name of BasicBlock following this statement
          -> CGMonad LLName  -- name of first BasicBlock
genSingle StEmpty following = newBlockJump following
genSingle (StBlock block) following = genBlock block following
genSingle (StExpr expr) following = do
    bb <- newBlockJump following
    void $ genExpression expr >>= addNamedInstrList
    return bb
genSingle (StVarDeclaration t n Nothing) following = do
    bb <- newBlockJump following
    label <- addInstr $ A.Alloca (toLLVMType t) Nothing 0 []
    setVarLabel n label
    return bb
genSingle (StVarDeclaration _ _ (Just _)) _ = undefined

genExpression :: Expression -> CGMonad [A.Named A.Instruction]
genExpression (ExLit (LitInt i) (Just t@(TypeInt sz))) = do
    aname <- getNewName "t"
    rname <- getNewName "t"
    return [A.Name aname A.:= A.Alloca (toLLVMType t) Nothing 0 []
           ,A.Name rname A.:= A.Store False (A.LocalReference (toLLVMType t) (A.Name aname))
                                      (A.ConstantOperand (A.C.Int (fromIntegral sz) i)) Nothing 0 []]



toLLVMType :: Type -> A.Type
toLLVMType (TypeInt s) = A.IntegerType $ fromIntegral s
toLLVMType (TypeUInt s) = A.IntegerType $ fromIntegral s
toLLVMType TypeFloat = A.float
toLLVMType TypeDouble = A.double
toLLVMType (TypePtr t) = A.ptr $ toLLVMType t
toLLVMType (TypeName _) = undefined

initializerFor :: Type -> A.C.Constant
initializerFor (TypeInt s) = A.C.Int (fromIntegral s) 0
initializerFor (TypeUInt s) = A.C.Int (fromIntegral s) 0
initializerFor _ = undefined


isDecVariable :: Declaration -> Bool
isDecVariable (DecVariable {}) = True
isDecVariable _ = False

isDecFunction :: Declaration -> Bool
isDecFunction (DecFunction {}) = True
isDecFunction _ = False