{-# 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,
      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

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 <- genTemp (sizeof TInt)
    addIns $ ILea ref n
    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 <- 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
convertStatement SDebugger nextnext = do
    addIns IDebugger
    setTerm $ IJmp 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 (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 <- 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"
        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 elemsz = sizeof $ 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 <- genTemp (sizeof TInt)
    arrsz <- genTemp (sizeof 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 elemsz))
    addIns $ IAri AAdd elemptr arrref offref
    ref <- genTemp elemsz
    addIns $ ILoad ref elemptr
    setTerm $ IJmp nextnext
    return ref
convertExpression (EGet st n t) nextnext = do
    let elemsz = sizeof $ fromJust t
    bl2 <- newBlockNoSwitch
    stref <- convertExpression st bl2
    switchBlock bl2
    let subtype = structMemberType (fromJust $ typeof st) n
    eref <- case subtype of
        TStruct _ -> genStructTemp elemsz
        _ -> genTemp elemsz
    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 <- 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)
    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
    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 <- genTemp (refSize 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 <- genTemp (sizeof TInt)
    elemptr <- genTemp (sizeof TInt)
    arrszptr <- genTemp (sizeof TInt)
    arrsz <- genTemp (sizeof 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 <- 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)
        elemptr <- genTemp (sizeof TInt)
        arrszptr <- genTemp (sizeof TInt)
        arrsz <- genTemp (sizeof 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 elemsz))
        addIns $ IAri AAdd elemptr ref offref
        dstref <- genTemp elemsz
        addIns $ ILoad dstref elemptr
        return dstref
    goLoad topae@(AEGet topup _ _) = do
        let (core, _, offset) = collectAESets topae
        coreref <- goLoad core
        ref <- genTemp (sizeof $ fromJust $ typeof topup)
        addIns $ IGet ref coreref offset
        return ref
getAESubscriptStoreRef _ = undefined