{-# 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
import Utils


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

data BuildState = BuildState
    { nextId :: Id,
      scopeStack :: [Scope],
      loopStack :: [Id],
      currentBlock :: Id,
      errorBlock :: Id,
      blockMap :: Map.Map Id BB,
      internedStrings :: [(Name, String)] }

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

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

genStructTemp :: Size -> BuildM Ref
genStructTemp sz = liftM (StructTemp sz) genId

genTempForType :: Type -> BuildM Ref
genTempForType t@(TStruct _) = genStructTemp (sizeof t)
genTempForType t = genTemp (sizeof t)

genTempLike :: Ref -> BuildM Ref
genTempLike (Temp sz _) = genTemp sz
genTempLike (StructTemp sz _) = genStructTemp sz
genTempLike _ = undefined

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

setErrorBlock :: Id -> BuildM ()
setErrorBlock i = modify $ \s -> s {errorBlock = 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

internString :: String -> BuildM Ref
internString str = do
    i <- genId
    let n = "__str_cnst_" ++ show i
    ref' <- genTempForType TInt
    ref <- genTempForType TInt
    addIns $ ILea ref' n
    addIns $ IAri AAdd ref ref' (Constant (refSize ref') (fromIntegral $ sizeof TInt))
    state $ \s -> (ref, s {internedStrings = internedStrings s ++ [(n, str)]})


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
        makeErrorBlock >>= setErrorBlock
        switchBlock firstid
        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

    makeErrorBlock :: BuildM Id
    makeErrorBlock = do
        bl <- newBlock
        addIns $ ICall "_builtin_outofbounds" []
        setTerm IUnreachable
        return bl

    result :: BuildM IRProgram
    result = do
        withScope $ do
            mapM_ (\(DVar t n _) -> scopeInsert n (Global (sizeof t) n) t) vars
            funcs' <- mapM goDFunc funcs
            ns <- gets internedStrings
            let strvars = flip map ns $ \(n, str) ->
                            let t = TArr TChar (Just $ fromIntegral $ length str)
                            in DVar t n (ELit (LStr str) (Just t))
            return $ IRProgram (vars ++ strvars) funcs'
buildIR _ = undefined

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 <- genTempForType 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
convertStatement SDebugger nextnext = do
    addIns IDebugger
    setTerm $ IJmp nextnext

convertExpression :: Expression -> Id -> BuildM Ref
convertExpression (ELit (LInt n) _) nextnext = do
    ref <- genTempForType TInt
    addIns $ IMov ref (Constant (sizeof TInt) (fromInteger n))
    setTerm $ IJmp nextnext
    return ref
convertExpression (ELit (LChar c) _) nextnext = do
    ref <- genTempForType 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 <- genTempForType 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 <- genTempForType typ
            addIns $ ICallr r n refs
            return r
    setTerm $ IJmp nextnext
    return destref
convertExpression (ELit (LStr s) _) nextnext = do
    ref <- internString s
    setTerm $ IJmp nextnext
    return ref
convertExpression (ELit (LStruct ms) stype) nextnext = do
    ref <- genStructTemp (sizeof $ fromJust stype)
    forM_ ms $ \(n,e) -> do
        bl <- newBlockNoSwitch
        r <- convertExpression e bl
        switchBlock bl
        addIns $ ISet ref (offsetInStruct (fromJust stype) n) r
    setTerm $ IJmp nextnext
    return ref
convertExpression (EBin BOAnd e1 e2 _) nextnext = do
    destref <- genTempForType 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 <- genTempForType (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"
        BOBitAnd -> addIns $ IAri AAnd ref ref1 ref2
        BOBitOr -> addIns $ IAri AOr ref ref1 ref2
        BOBitXor -> addIns $ IAri AXor ref ref1 ref2
        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 elemtype = fromJust t
    bl2 <- newBlockNoSwitch
    arrref <- convertExpression arr bl2
    switchBlock bl2
    bl3 <- newBlockNoSwitch
    subref <- convertExpression sub bl3
    switchBlock bl3
    offref <- genTemp (refSize subref)
    elemptr <- genTemp (refSize arrref)
    arrszptr <- genTempForType TInt
    arrsz <- genTempForType TInt

    errbl <- gets errorBlock

    addIns $ IAri ASub arrszptr arrref (Constant (refSize arrref) (fromIntegral $ sizeof TInt))
    addIns $ ILoad arrsz arrszptr
    bl4 <- newBlockNoSwitch
    setTerm $ IJcc CUGeq subref arrsz errbl bl4

    switchBlock bl4
    addIns $ IAri AMul offref subref (Constant (refSize subref) (fromIntegral $ sizeof elemtype))
    addIns $ IAri AAdd elemptr arrref offref
    ref <- genTempForType elemtype
    addIns $ ILoad ref elemptr
    setTerm $ IJmp nextnext
    return ref
convertExpression (EGet st n t) nextnext = do
    let elemtype = fromJust t
    assertM $ structMemberType (fromJust $ typeof st) n == elemtype
    bl2 <- newBlockNoSwitch
    stref <- convertExpression st bl2
    switchBlock bl2
    eref <- genTempForType elemtype
    addIns $ IGet eref stref (offsetInStruct (fromJust $ typeof st) n)
    setTerm $ IJmp nextnext
    return eref
convertExpression (ECast dt e) nextnext = do
    let typ = case typeof e of
                  Nothing -> error $ "Cast subject " ++ show e ++ " has Nothing type"
                  Just t -> t
    if typ == dt
        then convertExpression e nextnext
        else do
            when (not $ isIntegralType typ && isIntegralType dt) $
                error $ "convertExpression: unimplemented cast from " ++ pretty typ ++
                        " to " ++ pretty dt
            ref <- genTempForType 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' <- genTempForType (TArr t Nothing)
    ref <- genTempForType (TArr t Nothing)
    argref' <- genTempForType TInt
    argref <- genTempForType 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
    addIns $ IAri AAdd ref ref' (Constant (refSize ref') (fromIntegral $ sizeof TInt))
    setTerm $ IJmp nextnext
    return ref

convertAsExpression :: AsExpression -> Ref -> Id -> BuildM ()
convertAsExpression aevar@(AEVar _ _) valueref nextnext = do
    vref <- getAEVarRef aevar
    addIns $ IMov vref valueref
    setTerm $ IJmp nextnext
convertAsExpression aesubscript@(AESubscript _ _ _) valueref nextnext = do
    elemptr <- getAESubscriptStoreRef aesubscript
    addIns $ IStore elemptr valueref
    setTerm $ IJmp nextnext
convertAsExpression topae@(AEGet _ _ _) valueref nextnext = do
    let (core, _, offset) = collectAESets topae
    case core of
        aevar@(AEVar _ _) -> do
            vref <- getAEVarRef aevar
            addIns $ ISet vref offset valueref
        aesubscript@(AESubscript _ _ _) -> do
            elemptr <- getAESubscriptStoreRef aesubscript
            fieldptr <- genTempLike elemptr
            addIns $ IAri AAdd fieldptr elemptr (Constant (refSize elemptr) (fromIntegral offset))
            addIns $ IStore fieldptr valueref
        AEGet _ _ _ -> undefined
    setTerm $ IJmp nextnext

collectAESets :: AsExpression -> (AsExpression, [AsExpression], Offset)
collectAESets ae@(AEGet ae2 n _) =
    let (core, sets, offset) = collectAESets ae2
    in (core, ae : sets, offset + offsetInStruct (fromJust $ typeof ae2) n)
collectAESets ae = (ae, [], 0)

getAEVarRef :: AsExpression -> BuildM Ref
getAEVarRef (AEVar n _) = do
    mres <- findVar n
    case mres of
        Just (_, (r, _)) -> return r
        Nothing -> throwError $ "Undefined variable '" ++ n ++ "'" ++
                                " used in assignment expression"
getAEVarRef _ = undefined

getAESubscriptStoreRef :: AsExpression -> BuildM Ref
getAESubscriptStoreRef (AESubscript ae2 expr mrt) = do
    let elemsz = sizeof $ fromJust mrt
    ae2ref <- goLoad ae2
    bl2 <- newBlockNoSwitch
    subref <- convertExpression expr bl2
    switchBlock bl2
    offref <- genTempForType TInt
    elemptr <- genTempForType TInt
    arrszptr <- genTempForType TInt
    arrsz <- genTempForType TInt

    errbl <- gets errorBlock

    addIns $ IAri ASub arrszptr ae2ref (Constant (refSize ae2ref) (fromIntegral $ sizeof TInt))
    addIns $ ILoad arrsz arrszptr
    bl3 <- newBlockNoSwitch
    setTerm $ IJcc CUGeq subref arrsz errbl bl3

    switchBlock bl3
    addIns $ IAri AMul offref subref (Constant (refSize subref) (fromIntegral elemsz))
    addIns $ IAri AAdd elemptr ae2ref offref
    return elemptr
  where
    -- evaluate as if it were an Expression
    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 <- genTempForType t
        addIns $ IMov ref vref
        return ref
    goLoad (AESubscript ae expr' _) = do
        let elemtype = fromJust $ typeof ae
        ref <- goLoad ae
        bl2 <- newBlockNoSwitch
        eref <- convertExpression expr' bl2
        switchBlock bl2
        offref <- genTempForType TInt
        elemptr <- genTempForType TInt
        arrszptr <- genTempForType TInt
        arrsz <- genTempForType TInt

        errbl <- gets errorBlock

        addIns $ IAri ASub arrszptr ref (Constant (refSize ref) (fromIntegral $ sizeof TInt))
        addIns $ ILoad arrsz arrszptr
        bl3 <- newBlockNoSwitch
        setTerm $ IJcc CUGeq eref arrsz errbl bl3

        switchBlock bl3
        addIns $ IAri AMul offref eref (Constant (sizeof TInt) (fromIntegral $ sizeof elemtype))
        addIns $ IAri AAdd elemptr ref offref
        dstref <- genTempForType elemtype
        addIns $ ILoad dstref elemptr
        return dstref
    goLoad topae@(AEGet topup _ _) = do
        let (core, _, offset) = collectAESets topae
        coreref <- goLoad core
        ref <- genTempForType (fromJust $ typeof topup)
        addIns $ IGet ref coreref offset
        return ref
getAESubscriptStoreRef _ = undefined