{-# LANGUAGE DeriveFunctor, GeneralizedNewtypeDeriving #-}

module BuildIR(buildIR) where

import Control.Monad.Except
import Control.Monad.State.Strict
import Data.Char
import Data.List
import qualified Data.Map.Strict as Map
import Data.Maybe

import AST
import Defs
import Intermediate
import Pretty
import TypeRules


type Scope = Map.Map Name (Ref, Type)

data BuildState = BuildState
    { nextId :: Id,
      scopeStack :: [Scope],
      loopStack :: [Id],
      currentBlock :: Id,
      blockMap :: Map.Map Id BB }

initBuildState :: BuildState
initBuildState = BuildState
    { nextId = 0,
      scopeStack = [],
      loopStack = [],
      currentBlock = undefined,
      blockMap = Map.empty }

newtype BuildM a = BuildM {unBuildM :: StateT BuildState (Except String) a}
  deriving (Functor, Applicative, Monad, MonadState BuildState, MonadError String)

genId :: BuildM Id
genId = state $ \s -> (nextId s, s {nextId = nextId s + 1})

genTemp :: Size -> BuildM Ref
genTemp sz = liftM (Temp sz) genId

newBlock :: BuildM Id
newBlock = do
    i <- genId
    let block = BB i [] ITermNone
    modify $ \s -> s {currentBlock = i, blockMap = Map.insert i block (blockMap s)}
    return i

newBlockNoSwitch :: BuildM Id
newBlockNoSwitch = do
    i <- genId
    let block = BB i [] ITermNone
    modify $ \s -> s {blockMap = Map.insert i block (blockMap s)}
    return i

addIns :: IRIns -> BuildM ()
addIns ins =
    modify $ \s -> s {
        blockMap = Map.adjust bbAddIns (currentBlock s) (blockMap s)}
  where
    bbAddIns :: BB -> BB
    bbAddIns (BB bid inss term) = BB bid (inss ++ [ins]) term

setTerm :: IRTerm -> BuildM ()
setTerm term =
    modify $ \s -> s {
        blockMap = Map.adjust bbSetTerm (currentBlock s) (blockMap s)}
  where
    bbSetTerm :: BB -> BB
    bbSetTerm (BB bid inss oldterm) = case oldterm of
        ITermNone -> BB bid inss term
        _ -> error "setTerm: oldterm /= ITermNone"

clearBlockMap :: BuildM ()
clearBlockMap = modify $ \s -> s {currentBlock = undefined, blockMap = Map.empty}

getAllBlocks :: BuildM [BB]
getAllBlocks = liftM Map.elems (gets blockMap)

switchBlock :: Id -> BuildM ()
switchBlock bid = modify $ \s -> s {currentBlock = bid}

withLoop :: Id -> BuildM a -> BuildM a
withLoop i act = do
    modify $ \s -> s {loopStack = i : loopStack s}
    res <- act
    modify $ \s -> s {loopStack = tail (loopStack s)}
    return res

withScope :: BuildM a -> BuildM a
withScope act = do
    modify $ \s -> s {scopeStack = Map.empty : scopeStack s}
    res <- act
    modify $ \s -> s {scopeStack = tail (scopeStack s)}
    return res

modifyScope :: (Scope -> Scope) -> BuildM ()
modifyScope f =
    modify $ \s -> s {scopeStack = f (head (scopeStack s)) : tail (scopeStack s)}

scopeInsert :: Name -> Ref -> Type -> BuildM ()
scopeInsert n ref t = modifyScope $ Map.insert n (ref, t)

findVar :: Name -> BuildM (Maybe (Int, (Ref, Type)))
findVar n = do
    stk <- gets scopeStack
    let results = map (Map.lookup n) stk
    return $ fmap (\idx -> (idx, fromJust (results !! idx))) $
        findIndex isJust results


buildIR :: Program -> Error IRProgram
buildIR (Program vars funcs) =
    runExcept $ evalStateT (unBuildM result) initBuildState
  where
    goDFunc :: DFunc -> BuildM IRFunc
    goDFunc (DFunc rt n al bl) = do
        clearBlockMap 
        firstid <- newBlock
        lastid <- newBlockNoSwitch
        withScope $ do
            forM_ al $ \(at, an) -> scopeInsert an (Argument (sizeof at) an) at
            convertBlock bl lastid
        switchBlock lastid
        setTerm IRet
        bblist <- getAllBlocks
        return $ IRFunc rt n al bblist firstid

    result :: BuildM IRProgram
    result = do
        withScope $ do
            mapM_ (\(DVar t n _) -> scopeInsert n (Global (sizeof t) n) t) vars
            IRProgram vars <$> mapM goDFunc funcs

convertBlock :: Block -> Id -> BuildM ()
convertBlock (Block sts) nextnext = do
    withScope $ forM_ sts $ \st -> do
        endid <- newBlockNoSwitch
        convertStatement st endid
        switchBlock endid
    setTerm $ IJmp nextnext

convertStatement :: Statement -> Id -> BuildM ()
convertStatement (SDecl t n e) nextnext = do
    endid <- newBlockNoSwitch
    ref <- convertExpression e endid
    varref <- genTemp (sizeof t)
    scopeInsert n varref t
    switchBlock endid
    addIns $ IMov varref ref
    setTerm $ IJmp nextnext
convertStatement (SAs ae e) nextnext = do
    bl2 <- newBlockNoSwitch
    eref <- convertExpression e bl2
    switchBlock bl2
    convertAsExpression ae eref nextnext
convertStatement (SIf c b1 b2) nextnext = do
    cend <- newBlockNoSwitch
    blThen <- newBlockNoSwitch
    blElse <- newBlockNoSwitch

    cref <- convertExpression c cend
    switchBlock cend
    setTerm $ IJcc CNeq cref (Constant (refSize cref) 0) blThen blElse
    switchBlock blThen
    convertBlock b1 nextnext
    switchBlock blElse
    convertBlock b2 nextnext
convertStatement (SWhile c b) nextnext = do
    cond <- newBlockNoSwitch
    setTerm $ IJmp cond
    cend <- newBlockNoSwitch
    body <- newBlockNoSwitch
    bodyend <- newBlockNoSwitch

    switchBlock cond
    cref <- convertExpression c cend
    switchBlock cend
    setTerm $ IJcc CNeq cref (Constant (refSize cref) 0) body nextnext
    switchBlock body
    withLoop nextnext $ convertBlock b bodyend
    switchBlock bodyend
    setTerm $ IJmp cond
convertStatement (SBreak n) _ = do
    ls <- gets loopStack
    setTerm $ IJmp (ls !! n)
convertStatement (SReturn Nothing) _ = do
    setTerm IRet
convertStatement (SReturn (Just e)) _ = do
    bl <- newBlockNoSwitch
    ref <- convertExpression e bl
    switchBlock bl
    setTerm $ IRetr ref
convertStatement (SExpr e) nextnext = do
    void $ convertExpression e nextnext

convertExpression :: Expression -> Id -> BuildM Ref
convertExpression (ELit (LInt n) _) nextnext = do
    ref <- genTemp (sizeof TInt)
    addIns $ IMov ref (Constant (sizeof TInt) (fromInteger n))
    setTerm $ IJmp nextnext
    return ref
convertExpression (ELit (LChar c) _) nextnext = do
    ref <- genTemp (sizeof TChar)
    addIns $ IMov ref (Constant (sizeof TChar) (fromIntegral $ ord c))
    setTerm $ IJmp nextnext
    return ref
convertExpression (ELit (LVar n) _) nextnext = do
    mres <- findVar n
    case mres of
        Just (_, (r, t)) -> do
            ref <- genTemp (sizeof t)
            addIns $ IMov ref r
            setTerm $ IJmp nextnext
            return ref
        Nothing -> throwError $ "Undefined variable '" ++ n ++ "' referenced"
convertExpression (ELit (LCall n al) mrt) nextnext = do
    refs <- withScope $ forM al $ \arg -> do
        endid <- newBlockNoSwitch
        r <- convertExpression arg endid
        switchBlock endid
        return r
    destref <- case mrt of
        Nothing -> do
            addIns $ ICall n refs
            return $ Temp 0 (-1)
        Just typ -> do
            r <- genTemp (sizeof typ)
            addIns $ ICallr r n refs
            return r
    setTerm $ IJmp nextnext
    return destref
convertExpression (EBin BOAnd e1 e2 _) nextnext = do
    destref <- genTemp (sizeof TInt)
    bl2 <- newBlockNoSwitch
    blTryR <- newBlockNoSwitch
    bl3 <- newBlockNoSwitch
    blNope <- newBlockNoSwitch
    blYes <- newBlockNoSwitch

    ref1 <- convertExpression e1 bl2

    switchBlock bl2
    setTerm $ IJcc CNeq ref1 (Constant (refSize ref1) 0) blTryR blNope

    switchBlock blTryR
    ref2 <- convertExpression e2 bl3

    switchBlock bl3
    setTerm $ IJcc CNeq ref2 (Constant (refSize ref2) 0) blYes blNope

    switchBlock blYes
    addIns $ IMov destref (Constant (refSize destref) 1)
    setTerm $ IJmp nextnext

    switchBlock blNope
    addIns $ IMov destref (Constant (refSize destref) 0)
    setTerm $ IJmp nextnext

    return destref
convertExpression (EBin bo e1 e2 _) nextnext = do
    bl2 <- newBlockNoSwitch
    ref1 <- convertExpression e1 bl2
    switchBlock bl2
    bl3 <- newBlockNoSwitch
    ref2 <- convertExpression e2 bl3
    switchBlock bl3
    ref <- genTemp (sizeof $ fromJust $ retTypeBO bo (fromJust $ typeof e1) (fromJust $ typeof e2))
    case bo of
        BOAdd -> addIns $ IAri AAdd ref ref1 ref2
        BOSub -> addIns $ IAri ASub ref ref1 ref2
        BOMul -> addIns $ IAri AMul ref ref1 ref2
        BODiv -> addIns $ IAri ADiv ref ref1 ref2
        BOMod -> addIns $ IAri AMod ref ref1 ref2
        BOEq -> addIns $ IAri AEq ref ref1 ref2
        BONeq -> addIns $ IAri ANeq ref ref1 ref2
        BOGt -> addIns $ IAri AGt ref ref1 ref2
        BOLt -> addIns $ IAri ALt ref ref1 ref2
        BOGeq -> addIns $ IAri AGeq ref ref1 ref2
        BOLeq -> addIns $ IAri ALeq ref ref1 ref2
        BOPow -> error $ "Pow operator not implemented"
        BOAnd -> undefined
        BOOr -> undefined
    setTerm $ IJmp nextnext
    return ref
convertExpression (EUn UONot e mt) nextnext =
    convertExpression (EBin BOEq e (ELit (LInt 0) (typeof e)) mt) nextnext
convertExpression (EUn UONeg e mt) nextnext =
    convertExpression (EBin BOSub (ELit (LInt 0) (typeof e)) e mt) nextnext
convertExpression (ESubscript arr sub t) nextnext = do
    let elemsz = sizeof $ fromJust t
    bl2 <- newBlockNoSwitch
    arrref <- convertExpression arr bl2
    switchBlock bl2
    bl3 <- newBlockNoSwitch
    subref <- convertExpression sub bl3
    switchBlock bl3
    offref <- genTemp (refSize subref)
    off8ref <- genTemp (refSize subref)
    elemptr <- genTemp (refSize arrref)
    addIns $ IAri AMul offref subref (Constant (refSize subref) (fromIntegral elemsz))
    addIns $ IAri AAdd off8ref offref (Constant (refSize subref) (fromIntegral $ sizeof TInt))
    addIns $ IAri AAdd elemptr arrref off8ref
    ref <- genTemp elemsz
    addIns $ ILoad ref elemptr
    setTerm $ IJmp nextnext
    return ref
convertExpression (ECast dt e) nextnext = do
    let typ = case typeof e of
                  Nothing -> error $ "Cast subject " ++ show e ++ " has Nothing type"
                  Just t -> t
    when (not $ isIntegralType typ && isIntegralType dt) $
        error $ "convertExpression: unimplemented cast from " ++ pretty typ ++ " to " ++ pretty dt
    ref <- genTemp (sizeof dt)
    bl <- newBlockNoSwitch
    eref <- convertExpression e bl
    switchBlock bl
    addIns $ IResize ref eref
    setTerm $ IJmp nextnext
    return ref
convertExpression (ENew t sze) nextnext = do
    when (not $ isBasicType t) $
        throwError $ "Array element type in 'new' expression is not a basic type (" ++ pretty t ++ ")"
    bl2 <- newBlockNoSwitch
    szref <- convertExpression sze bl2
    switchBlock bl2
    ref <- genTemp (sizeof $ TArr t Nothing)
    argref' <- genTemp (sizeof TInt)
    argref <- genTemp (sizeof TInt)
    addIns $ IAri AMul argref' szref (Constant (sizeof TInt) (fromIntegral $ sizeof t))
    addIns $ IAri AAdd argref argref' (Constant (sizeof TInt) (fromIntegral $ sizeof TInt))
    addIns $ ICallr ref "_builtin_malloc" [argref]
    addIns $ IStore ref szref
    setTerm $ IJmp nextnext
    return ref

convertAsExpression :: AsExpression -> Ref -> Id -> BuildM ()
convertAsExpression (AEVar n _) valueref nextnext = do
    mres <- findVar n
    vref <- case mres of
                Just (_, (r, _)) -> return r
                Nothing -> throwError $ "Undefined variable '" ++ n ++ "'" ++
                                        " used in assignment expression"
    addIns $ IMov vref valueref
    setTerm $ IJmp nextnext
convertAsExpression (AESubscript ae2 expr mrt) valueref nextnext = do
    let elemsz = sizeof $ fromJust mrt
    ae2ref <- goLoad ae2
    bl2 <- newBlockNoSwitch
    subref <- convertExpression expr bl2
    switchBlock bl2
    offref' <- genTemp (sizeof TInt)
    offref <- genTemp (sizeof TInt)
    elemptr <- genTemp (sizeof TInt)
    -- TODO: do bounds checking
    addIns $ IAri AMul offref' subref (Constant (sizeof TInt) (fromIntegral elemsz))
    addIns $ IAri AAdd offref offref' (Constant (sizeof TInt) (fromIntegral $ sizeof TInt))
    addIns $ IAri AAdd elemptr ae2ref offref
    addIns $ IStore elemptr valueref
    setTerm $ IJmp nextnext
  where
    goLoad :: AsExpression -> BuildM Ref
    goLoad (AEVar n _) = do
        mres <- findVar n
        (vref, t) <- case mres of
                        Just (_, (r, t)) -> return (r, t)
                        Nothing -> throwError $ "Undefined variable '" ++ n ++ "'" ++
                                                " used in assignment expression"
        ref <- genTemp (sizeof t)
        addIns $ IMov ref vref
        return ref
    goLoad (AESubscript ae expr' _) = do
        let elemsz = sizeof $ fromJust $ typeof ae
        ref <- goLoad ae
        bl2 <- newBlockNoSwitch
        eref <- convertExpression expr' bl2
        switchBlock bl2
        offref' <- genTemp (sizeof TInt)
        offref <- genTemp (sizeof TInt)
        elemptr <- genTemp (sizeof TInt)
        -- TODO: do bounds checking
        addIns $ IAri AMul offref' eref (Constant (sizeof TInt) (fromIntegral elemsz))
        addIns $ IAri AAdd offref offref' (Constant (sizeof TInt) (fromIntegral $ sizeof TInt))
        addIns $ IAri AAdd elemptr ref offref
        dstref <- genTemp elemsz
        addIns $ ILoad dstref elemptr
        return dstref