{-# LANGUAGE DeriveFunctor, GeneralizedNewtypeDeriving, BangPatterns #-}
module Interpreter(interpret) where

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

-- import Debug.Trace

import AST
import qualified GCStore as GCS


interpret :: Program -> (Maybe String, String, VM)  -- (Maybe error, output, vm)
interpret prog =
    let res = evaluateProgram prog >> liftM show get >>= tellLine :: ExMonad ()
        m1 = unExMonad res :: ExceptT String (WriterT String (State VM)) ()
        m2 = runExceptT m1 :: WriterT String (State VM) (Either String ())
        m3 = runWriterT m2 :: State VM (Either String (), String)
        m4 = runState m3 makeVM :: ((Either String (), String), VM)
        vm = snd m4 :: VM
        merr = either Just (const Nothing) (fst (fst m4))
        output = snd (fst m4)
    in (merr, output, vm)


type Scope = Map.Map Name GCS.Id

data Value
    = VNum Double
    | VStr String
    | VBlock BlockType ArgList Block
    | VNil
  deriving (Show)

type ObjectStore = GCS.Store Value

data VM = VM {scopeStack :: [Scope],
              objStore :: ObjectStore,
              tempValue :: Int}
  deriving (Show)

newtype ExMonad a = ExContext {unExMonad :: ExceptT String (WriterT String (State VM)) a}
  deriving (Functor, Applicative, Monad, MonadState VM, MonadWriter String, MonadError String)

typeNameOf :: Value -> String
typeNameOf (VNum _) = "Number"
typeNameOf (VStr _) = "String"
typeNameOf (VBlock _ _ _) = "Block"
typeNameOf VNil = "Nil"

niceThrowError :: String -> ExMonad a
niceThrowError s = tellLine s >> throwError s

getTopScope :: String -> ExMonad Scope
getTopScope mname = get >>= \(VM stk _ _) -> case stk of
    [] -> niceThrowError $ mname ++ " on empty scope stack"
    (sc:_) -> return sc

modifyScopeStack :: ([Scope] -> [Scope]) -> ExMonad ()
modifyScopeStack f = modify $ \vm@(VM stk _ _) -> vm {scopeStack = f stk}

putScopeStack :: [Scope] -> ExMonad ()
putScopeStack = modifyScopeStack . const

modifyObjStore :: (ObjectStore -> ObjectStore) -> ExMonad ()
modifyObjStore f = modify $ \vm@(VM _ os _) -> vm {objStore = f os}

putObjStore :: ObjectStore -> ExMonad ()
putObjStore = modifyObjStore . const

modifyTempValue :: (Int -> Int) -> ExMonad ()
modifyTempValue func = modify $ \vm -> vm {tempValue = func (tempValue vm)}

getTempValue :: ExMonad Int
getTempValue = liftM tempValue get

getRefCount :: GCS.Id -> ExMonad Int
getRefCount gcsid = get >>= \vm -> return (GCS.refcount (objStore vm) gcsid)

derefTell :: GCS.Id -> ExMonad ()
derefTell gcsid = do
    get >>= \vm -> return (GCS.deref (objStore vm) gcsid)
        >>= putObjStore
    rc <- getRefCount gcsid
    tellLine $ "GCS deref " ++ show gcsid ++ " -> rc=" ++ show rc

cleanupScope :: Scope -> ExMonad ()
cleanupScope sc = do
    tellLine $ "## cleanupScope: " ++ show sc
    mapM_ cleanupGcsid (Map.elems sc)

cleanupGcsid :: GCS.Id -> ExMonad ()
cleanupGcsid gcsid = do
    vm <- get
    case GCS.retrieve (objStore vm) gcsid of
         Nothing -> niceThrowError $ "cleanupScope: gcsid " ++ show gcsid ++ " doesn't exist"
         Just (VBlock _ _ b) -> cleanupBlock b
         _ -> return ()
    derefTell gcsid

cleanupBlock :: Block -> ExMonad ()
cleanupBlock block = do
    let idset = collectBlock block
    mapM_ derefTell idset

collectBlock :: Block -> Set.Set GCS.Id
collectBlock (Block sts) = foldl Set.union $ map collectStatement sts

collectStatement :: Statement -> Set.Set GCS.Id
collectStatement (Declaration _ e) = collectExpression e
collectStatement (Assignment _ e) = collectExpression e
collectStatement (Condition e b1 b2) =
    foldl Set.union [collectExpression e, collectBlock b1, collectBlock b2]
collectStatement (Dive _ es b) =
    Set.union (foldl Set.union (map collectExpression es)) (collectBlock b)
collectStatement (Expr e) = collectExpression e

collectExpression :: Expression -> Set.Set GCS.Id
collectExpression (EBin _ e1 e2) = Set.union (collectExpression e1) (collectExpression e2)
collectExpression (EUn _ e) = collectExpression e
collectExpression (ELit li) = collectLiteral li

collectLiteral :: Literal -> Set.Set GCS.Id
collectLiteral (LBlock _ _ b) = collectBlock b
collectLiteral (LGCSId gcsid') = cleanupGcsid gcsid'
collectLiteral _ = return ()

cleanupValue :: Value -> Set.Set GCS.Id
cleanupValue (VBlock _ _ b) = cleanupBlock b
cleanupValue _ = return ()

pushScope :: ExMonad ()
pushScope = tellLine "pushScope" >> modifyScopeStack (Map.empty :)

popScope :: ExMonad ()
popScope = do
    vm <- get
    case scopeStack vm of
        [] -> niceThrowError $ "popScope on empty scope stack"
        (sc:rest) -> do
            cleanupScope sc
            putScopeStack rest
            tellLine "popScope"

findVar :: Name -> ExMonad (Maybe (Int, GCS.Id, Value))
findVar n = do
    (VM stk os _) <- get
    let mbs = map (\sc -> Map.lookup n sc) stk
        dr = dropWhile isNothing mbs
        gcsid = fromJust (head dr)  -- not evaluated if dr == []
    if null dr
        then return Nothing
        else maybe (niceThrowError $ "findVar: gcsid " ++ show gcsid ++ " doesn't exist") return
                   (GCS.retrieve os gcsid)
                >>= \value -> return (Just (length mbs - length dr, gcsid, value))

createVarInScope :: Scope -> ObjectStore -> Name -> Value -> ExMonad (Scope, ObjectStore)
createVarInScope sc os n val =
    let (gcsid, os') = GCS.store os val
        sc' = Map.insert n gcsid sc
    in tellLine ("GCS store in " ++ show gcsid ++ " value " ++ show val) >>
           if isNothing (Map.lookup n sc)
               then return (sc', os')
               else error $ "Var " ++ n ++ " already exists in createVarInScope"

updateVarInScope :: Scope -> ObjectStore -> Name -> Value -> ExMonad ObjectStore
updateVarInScope sc os n val = do
    let gcsid = maybe (error $ "Var " ++ n ++ " not found in updateVarInScope") id (Map.lookup n sc)
        prev = maybe (error "Retrieve=Nothing in updateVarInScope") id (GCS.retrieve os gcsid)
    cleanupValue prev
    tellLine ("GCS update " ++ show gcsid ++ " to value " ++ show val)
    return (GCS.update os gcsid val)

createVarLocal :: Name -> Value -> ExMonad ()
createVarLocal n val = do
    vm <- get
    let (VM (sc:rest) os _) = vm
    (sc', os') <- createVarInScope sc os n val
    put $ vm {scopeStack = sc' : rest, objStore = os'}

updateVarAt :: Int -> Name -> Value -> ExMonad ()
updateVarAt idx n val = do
    vm <- get
    let (VM stk os _) = vm
    updateVarInScope (stk !! idx) os n val >>= putObjStore

tellLine :: String -> ExMonad ()
tellLine s = get >>= \vm ->
    let lvl = length (scopeStack vm)
        ind = if lvl >= 1 then take (4 * (lvl - 1)) (cycle "|   ") else error "ZERO LEVEL"
    in tell $ ind ++ replace "\n" ('\n' : ind) s ++ "\n"

tellLineStandout :: String -> ExMonad ()
tellLineStandout s = tell $ "\x1B[1m" ++ s ++ "\x1B[0m\n"

replace :: String -> String -> String -> String
replace _ _ "" = ""
replace a b subj =
    if take (length a) subj == a
        then b ++ replace a b (drop (length a) subj)
        else head subj : replace a b (tail subj)


makeVM :: VM
makeVM = VM [Map.empty] GCS.empty 0 

evaluateProgram :: Program -> ExMonad ()
evaluateProgram (Program (Block sts)) = mapM_ evStatement sts

evStatement :: Statement -> ExMonad ()
-- evStatement a | traceShow a False = unreachable
evStatement a = tellLine ("-# evStatement: " ++ astPretty a) >> evStatement' a

evStatement' :: Statement -> ExMonad ()
evStatement' (Declaration n e) =
    getTopScope "evStatement" >>=
        maybe (evExpression e >>= createVarLocal n)
              (const $ niceThrowError $ "Variable " ++ n ++ " already exists in scope")
          . Map.lookup n
evStatement' (Assignment n e) =
    findVar n >>= 
        maybe (niceThrowError $ "Variable " ++ n ++ " assigned to but not found")
              (\(idx, _, _) -> evExpression e >>= updateVarAt idx n)
evStatement' (Condition cond b1 b2) =
    evExpression cond >>= truthValue >>= \bool -> evBlock $ if bool then b1 else b2
evStatement' (Dive "print" al b) = do
    (show <$> mapM evExpression al) >>= tellLineStandout
    evBlock b
evStatement' (Dive n al b) =
    findVar n >>=
        maybe (niceThrowError $ "Variable " ++ n ++ " dived into but not found")
              (\(_, _, value) -> case value of
                  (VBlock BT2 val vb)
                      | length al == length val -> do
                          pushScope
                          evBlockNoScope $ Block [Declaration dn de | (dn, de) <- zip val al]
                          evBlockNoScope vb
                          evBlockNoScope b
                          popScope
                      | otherwise ->
                          niceThrowError $ "Invalid number of arguments to dived-in block " ++ n
                  _ -> niceThrowError $ "Cannot dive into " ++ n ++ " of invalid type " ++ typeNameOf value)
evStatement' (Expr e) = evExpression e >>= cleanupValue

evExpression :: Expression -> ExMonad Value
-- evExpression a | traceShow a False = unreachable
evExpression a = tellLine ("-# evExpression: " ++ astPretty a) >> evExpression' a

evExpression' :: Expression -> ExMonad Value
evExpression' (EBin bo e1 e2) = evBO bo e1 e2
evExpression' (EUn uo e) = evExpression e >>= evUO uo
evExpression' (ELit li) = evLiteral li

evBlock :: Block -> ExMonad ()
-- evBlock a | traceShow a False = unreachable
evBlock bl = do
    pushScope
    evBlockNoScope bl
    popScope

evBlockNoScope :: Block -> ExMonad ()
evBlockNoScope a = tellLine ("-# evBlockNoScope: " ++ astPretty a) >> evBlockNoScope' a

evBlockNoScope' :: Block -> ExMonad ()
evBlockNoScope' (Block sts) = mapM_ evStatement sts

evLiteral :: Literal -> ExMonad Value
evLiteral (LNum m) = return $ VNum m
evLiteral (LStr s) = return $ VStr s
evLiteral (LVar n) =
    findVar n >>= maybe (niceThrowError $ "Variable " ++ n ++ " referenced but not found") (return . thrd)
evLiteral (LBlock BT0 al bl)
    | length al == 0 = evBlock bl >> return VNil
    | otherwise = niceThrowError $ "Immediately invoked block literal cannot have parameters"
evLiteral (LBlock bt al bl) = liftM (VBlock bt al) $ processBlock bl
evLiteral (LGCSId gcsid) = get >>= \(VM _ os _) -> maybe (niceThrowError $ "evLiteral: gcsid " ++ show gcsid ++ " doesn't exist") return (GCS.retrieve os gcsid)
evLiteral LNil = return VNil

evUO :: UnaryOp -> Value -> ExMonad Value
evUO UONeg (VNum m) = return $ VNum (-m)
evUO UONeg value = niceThrowError $ "Operator '(-)' does not take a value of type " ++ typeNameOf value
evUO UONot value = (bool2VNum . not) <$> truthValue value

classifyBO :: BinaryOp -> (Bool, Bool, Bool)
classifyBO bo = case bo of
    BOPlus -> (True, f, f)
    BOMinus -> (True, f, f)
    BOMul -> (True, f, f)
    BODiv -> (True, f, f)
    BOMod -> (True, f, f)
    BOPow -> (True, f, f)
    BOLess -> (f, True, f)
    BOGreater -> (f, True, f)
    BOEqual -> (f, True, f)
    BOLEq -> (f, True, f)
    BOGEq -> (f, True, f)
    BOBoolAnd -> (f, f, True)
    BOBoolOr -> (f, f, True)
  where f = False

isArithBO, isCompBO {-, isBoolBO-} :: BinaryOp -> Bool
isArithBO bo = let (r, _, _) = classifyBO bo in r
isCompBO bo = let (_, r, _) = classifyBO bo in r
-- isBoolBO bo = let (_, _, r) = classifyBO bo in r

evArithBO :: BinaryOp -> Double -> Double -> Double
evArithBO BOPlus = (+)
evArithBO BOMinus = (-)
evArithBO BOMul = (*)
evArithBO BODiv = (/)
evArithBO BOMod = \a b -> a - fromIntegral (floor (a / b) :: Integer) * b
evArithBO BOPow = (**)
evArithBO _ = unreachable

evOrderingBO :: Ord a => BinaryOp -> a -> a -> Value
evOrderingBO bo a b =
    let c = compare a b
    in bool2VNum $ case bo of
        BOLess -> c == LT
        BOGreater -> c == GT
        BOEqual -> c == EQ
        BOLEq -> c == LT || c == EQ
        BOGEq -> c == GT || c == EQ
        _ -> unreachable

evBO :: BinaryOp -> Expression -> Expression -> ExMonad Value
evBO bo e1 e2
    | isArithBO bo = evExpression e1 >>= \v1 -> evExpression e2 >>= \v2 -> case (v1, v2) of
        (VNum m1, VNum m2) -> return $ VNum $ evArithBO bo m1 m2
        (_, _) -> niceThrowError $ "Operator '" ++ astPretty bo ++ "' does not take types " ++
                                   typeNameOf v1 ++ " and " ++ typeNameOf v2
    | isCompBO bo = evExpression e1 >>= \v1 -> evExpression e2 >>= \v2 -> case (v1, v2) of
        (VNum m1, VNum m2) -> return $ evOrderingBO bo m1 m2
        (VStr s1, VStr s2) -> return $ evOrderingBO bo s1 s2
        (_, _) -> niceThrowError $ "Operator '" ++ astPretty bo ++ "' does not take types " ++
                                   typeNameOf v1 ++ " and " ++ typeNameOf v2
    | bo == BOBoolAnd =
        evExpression e1 >>= truthValue >>= \v1 ->
            if not v1 then return (VNum 0)
                      else evExpression e2 >>= truthValue >>= return . bool2VNum
    | bo == BOBoolOr =
        evExpression e1 >>= truthValue >>= \v1 ->
            if v1 then return (VNum 1)
                  else evExpression e2 >>= truthValue >>= return . bool2VNum
    | otherwise = unreachable


truthValue :: Value -> ExMonad Bool
truthValue (VNum 0) = return False
truthValue (VNum _) = return True
truthValue (VStr "") = return False
truthValue (VStr _) = return True
truthValue (VBlock _ _ _) = niceThrowError "Block not valid as truth value"
truthValue VNil = return False

bool2VNum :: Bool -> Value
bool2VNum True = VNum 1
bool2VNum False = VNum 0


data ProcessState = ProcessState {psStack :: [Set.Set Name], psSet :: Set.Set GCS.Id}

newtype ProcessMonad a = ProcessMonad {unProcessMonad :: State ProcessState a}
  deriving (Functor, Applicative, Monad, MonadState ProcessState)

pmModifyStack :: ([Set.Set Name] -> [Set.Set Name]) -> ProcessMonad ()
pmModifyStack f = modify $ \ps@(ProcessState s _) -> ps {psStack = f s}

pmModifyStackTop :: (Set.Set Name -> Set.Set Name) -> ProcessMonad ()
pmModifyStackTop f = modify $ \ps@(ProcessState (s:ss) _) -> ps {psStack = f s : ss}

pmModifySet :: (Set.Set GCS.Id -> Set.Set GCS.Id) -> ProcessMonad ()
pmModifySet f = modify $ \ps@(ProcessState _ s) -> ps {psSet = f s}

processBlockCollect :: [Scope] -> Block -> ProcessMonad Block
processBlockCollect gStack = goBlock
  where
    goBlock :: Block -> ProcessMonad Block
    goBlock (Block sts) = do
        pmModifyStack (Set.empty :)
        b <- Block <$> mapM goStatement sts
        pmModifyStack tail
        return b

    goStatement :: Statement -> ProcessMonad Statement
    goStatement (Declaration n e) = do
        d <- Declaration n <$> goExpression e
        pmModifyStackTop (Set.insert n)
        return d
    goStatement (Assignment n e) = Assignment n <$> goExpression e
    goStatement (Condition e b1 b2) = Condition <$> goExpression e <*> goBlock b1 <*> goBlock b2
    goStatement (Dive n al b) = Dive n <$> mapM goExpression al <*> goBlock b
    goStatement (Expr e) = Expr <$> goExpression e

    goExpression :: Expression -> ProcessMonad Expression
    goExpression (EBin bo e1 e2) = EBin bo <$> goExpression e1 <*> goExpression e2
    goExpression (EUn uo e) = EUn uo <$> goExpression e
    goExpression (ELit l) = ELit <$> goLiteral l

    goLiteral :: Literal -> ProcessMonad Literal
    goLiteral (LVar name) = do
        (ProcessState st _) <- get
        let midx1 = findIndex (Set.member name) st
            midx2 = findIndex (Map.member name) gStack
        case (midx1, midx2) of
            (Nothing, Nothing) -> return (LVar name)
            (Just _, _) -> return (LVar name)
            (Nothing, Just idx) ->
                let gcsid = fromJust $ Map.lookup name (gStack !! idx)
                in pmModifySet (Set.insert gcsid) >> return (LGCSId gcsid)
    goLiteral (LBlock bt al bl) = LBlock bt al <$> goBlock bl
    goLiteral (LGCSId gcsid) = do
        pmModifySet (Set.insert gcsid)
        return (LGCSId gcsid)
    goLiteral l = return l

processBlock :: Block -> ExMonad Block
processBlock block = do
    tellLine ("## processBlock: " ++ astPretty block)
    (VM stk os _) <- get
    let (block', ProcessState _ psset) =
            runState (unProcessMonad $ processBlockCollect stk block)
                     (ProcessState [] Set.empty)
    let refitems = Set.toList psset
        os' = foldl GCS.ref os psset
    mapM_ (\gcsid -> tellLine ("GCS ref " ++ show gcsid)) refitems
    putObjStore os'
    return block'

-- processBlock :: Block -> ExMonad Block
-- processBlock b = tellLine ("## processBlock: " ++ show b) >> modifyTempValue (const 0) >> goBlock b
--   where
--     goStatement :: Statement -> ExMonad Statement
--     goStatement (Declaration n e) = do
--         d <- Declaration n <$> goExpression e
--         createVarLocal n VNil
--         return d
--     goStatement (Assignment n e) = Assignment n <$> goExpression e
--     goStatement (Condition e b1 b2) = Condition <$> goExpression e <*> goBlock b1 <*> goBlock b2
--     goStatement (Dive n al b') = Dive n <$> mapM goExpression al <*> goBlock b'
--     goStatement (Expr e) = Expr <$> goExpression e

--     goBlock :: Block -> ExMonad Block
--     goBlock (Block sts) = do
--         pushScope
--         modifyTempValue succ
--         b' <- Block <$> mapM goStatement sts
--         modifyTempValue pred
--         popScope
--         return b'

--     goExpression :: Expression -> ExMonad Expression
--     goExpression (EBin bo e1 e2) = EBin bo <$> goExpression e1 <*> goExpression e2
--     goExpression (EUn uo e) = EUn uo <$> goExpression e
--     goExpression (ELit l) = ELit <$> goLiteral l

--     goLiteral :: Literal -> ExMonad Literal
--     goLiteral (LVar name) = do
--         tv <- getTempValue
--         findVar name >>=
--             maybe (niceThrowError $ "Unknown variable " ++ name ++ " while processing block")
--                   (\(idx, gcsid, _) ->
--                       if idx >= tv
--                           then modifyObjStore (\os -> GCS.ref os gcsid)
--                                  >> tellLine ("GCS ref " ++ show gcsid)
--                                  >> return (LGCSId gcsid)
--                           else return (LVar name))
--     goLiteral (LBlock bt al bl) = LBlock bt al <$> goBlock bl
--     goLiteral l = return l


thrd :: (a, b, c) -> c
thrd (_, _, c) = c

unreachable :: a
unreachable = error "Unreachable"