{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE DeriveFunctor, DeriveFoldable, DeriveTraversable #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE DeriveGeneric, TypeOperators #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE OverloadedStrings #-}

module Hop.Hoplite.Eval where


import Hop.Hoplite.Heap
import Control.Monad.RWS.Class
import Control.Monad.RWS.Strict

import Data.Typeable
import Data.Data
import GHC.Generics
import qualified Data.Map.Strict as Map
import Numeric.Natural
import Data.Text (Text,pack)
import Bound
import Hop.Hoplite.Types
import Control.Lens (view,_1,makeLenses)
import Hop.Hoplite.STExcept


data Cmd = Cmd {
      from :: Text
    , to :: Text
    , posAmount :: Rational
    , fakeCryptoSig:: Text }
  deriving (Eq,Ord,Show,Read,Data,Typeable)

data ExpContext ty a = SCEmpty
                     | LetContext
                       (Maybe Text)
                       (Scope (Maybe Text) (Exp ty) a)
                       (ExpContext ty a)
                     | FunAppCtxt [Ref] [Exp ty a] (ExpContext ty a)
                     | PrimAppCtxt PrimOpId [Ref] [Exp ty a] (ExpContext ty a)
   deriving (Typeable
            ,Functor
            ,Foldable
            ,Traversable
            ,Generic
            ,Data
            ,Eq
            ,Ord
            ,Show)


data InterpreterError
  = PrimopTypeMismatch
  | NonClosureInApplicationPosition
  | ArityMismatchFailure
  | HeapLookupFailure
  | MalformedClosure
  | MismatchedStackContext
  | PrimFailure String
  | UnsupportedTermConstructFailure String
  deriving (Eq,Ord,Show,Typeable,Data)

newtype TransactionId = TransactionId Natural
  deriving (Eq, Show, Num, Enum, Ord)

newtype OpId = OpId Natural
  deriving (Eq, Show, Num, Enum, Ord)

type OrderedOp = (OpId, Cmd)

data PersistentState
  = PersistentState { _persistentNextTxId :: TransactionId
                    , _persistentBalances :: Map.Map Text Rational
                    , _persistentTxes :: [(TransactionId, [OrderedOp])]
                    }
  deriving (Eq, Show)

data InterpreterOverlay
  = InterpreterOverlay { _overlayNextOpId :: OpId
                       , _overlayBalances :: Map.Map Text Rational
                       }
  deriving (Eq, Show)

newtype BalancesDiff
  = BalancesDiff { _balancesDiff :: Map.Map Text (Rational, Rational) } -- (before, after)
  deriving (Eq, Show)

data InterpreterDiff
  = InterpreterDiff { _interpDiffBalances :: BalancesDiff
                    , _interpDiffOps :: [OrderedOp] }
  deriving (Eq, Show)

data InterpreterOutput
  = InterpreterOutput { _outputDiff :: InterpreterDiff
                      , _outputNextState :: PersistentState
                      }
  deriving (Eq, Show)

$(makeLenses ''PersistentState)
$(makeLenses ''InterpreterOverlay)
$(makeLenses ''BalancesDiff)
$(makeLenses ''InterpreterDiff)
$(makeLenses ''InterpreterOutput)

initialState :: PersistentState
initialState = PersistentState 0 Map.empty []

emptyOverlay :: InterpreterOverlay
emptyOverlay = InterpreterOverlay 0 Map.empty

--
-- TODO: test mkOutput
--
mkOutput :: InterpreterOverlay -> [OrderedOp] -> PersistentState -> InterpreterOutput
mkOutput (InterpreterOverlay _ overlayBals) ops (PersistentState txId bals txes) =
  InterpreterOutput diff nextState
    where
      balChanges = Map.intersectionWith (\ a b -> (a, b)) bals overlayBals
      diff = InterpreterDiff (BalancesDiff balChanges) ops
      nextState = PersistentState (succ txId)
                                  (overlayBals `mappend` bals)
                                  ((txId, ops) : txes)

runExpr :: Natural
        -> PersistentState
        -> (forall v . Exp () v) -- guarantees we have no free variables
        -> Either (() :+ InterpreterError :+ HeapError)
                  InterpreterOutput
runExpr step st expr = fmap projectOutput
                     $ handleSTE id
                     $ runEmptyHeap step
                     $ runRWST (evalExp SCEmpty expr) env emptyOverlay
  where
    env = _persistentBalances st
    projectOutput ((_heapVal, overlay, ops), _heap) = mkOutput overlay ops st

type InterpStack s ty b a
  = RWST (Map.Map Text Rational)
         [OrderedOp]
         InterpreterOverlay
         (HeapStepCounterM (Exp ty)
                           (STE (b :+ InterpreterError :+ HeapError) s))
         a

evalExp :: (Ord ty,Show ty) => ExpContext ty Ref -> Exp ty Ref
        -> InterpStack s ty b (HeapVal (Exp ty), Ref)
evalExp stk (V rf) = do rp@(_hpval,_ref) <- lift $ heapRefLookupTransitive rf
                        applyStack stk rp
evalExp _stk (Force _e)= throwInterpError $ UnsupportedTermConstructFailure "Force"
evalExp _stk (Delay _e) = throwInterpError $ UnsupportedTermConstructFailure "Delay"
evalExp stk (funE :@ args) = evalClosureApp (FunAppCtxt [] (funE:args) stk )
evalExp stk (ELit l) = do ref <- lift $ heapAllocate (VLitF l) ; applyStack stk (VLitF l, ref)
evalExp stk (PrimApp name args) = evalPrimApp (PrimAppCtxt name [] args stk )
evalExp stk (Lam ts bod) = do let val = DirectClosureF (MkClosure (map (ArityBoxed . view _1 )ts ) bod)
                              ref <- lift $ heapAllocate val
                              applyStack stk (val,ref)
evalExp stk (Let mv _mty rhsExp scp) = evalExp (LetContext mv scp stk) rhsExp

noBlackholeArgs :: (Ord ty,Show ty) => [Ref] -> InterpStack s ty b ()
noBlackholeArgs rls = lift $ void $ traverse heapRefLookupTransitive rls

throwInterpError :: InterpreterError
                 -> InterpStack s ty b a
throwInterpError e =lift $ lift $ throwSTE $ (InL . InR) e

evalClosureApp :: forall ty s b . (Show ty, Ord ty)
               => ExpContext ty Ref
               -> InterpStack s ty b (HeapVal (Exp ty), Ref)
evalClosureApp (FunAppCtxt ls [] stk) = do
  (funRef:argsRef) <- return $ reverse ls
  (val,_ref) <- lift $ heapRefLookupTransitive funRef
  noBlackholeArgs argsRef
  case val of
    (DirectClosureF (MkClosure wrpNames scp))
      | length argsRef == length wrpNames
      -> let nmMap = Map.fromList $ zip (map _extractArityInfo wrpNames) argsRef
             substApply :: Var Text (Exp ty2 Ref)
                        -> InterpStack s ty b (Exp ty2 Ref)
             substApply var = case var of
                     (B nm) -> maybe (throwInterpError MalformedClosure)
                                     (\ rf -> lift $ return $ V rf)
                                     (Map.lookup nm nmMap)
                     (F term) -> return term
         in do
           nextExp <- traverse substApply $ unscope scp
           evalExp stk $ join nextExp
      | otherwise -> throwInterpError ArityMismatchFailure
    _ -> throwInterpError NonClosureInApplicationPosition
evalClosureApp (FunAppCtxt ls (h : t) stk) = evalExp (FunAppCtxt ls t stk) h
evalClosureApp LetContext {} = throwInterpError MismatchedStackContext
evalClosureApp PrimAppCtxt {} = throwInterpError MismatchedStackContext
evalClosureApp SCEmpty = throwInterpError MismatchedStackContext


evalPrimApp ::(Show ty, Ord ty) => ExpContext ty Ref -> InterpStack s ty b (HeapVal (Exp ty), Ref)
evalPrimApp (PrimAppCtxt nm args [] stk) = do noBlackholeArgs args ; applyPrim stk nm $ reverse args
evalPrimApp (PrimAppCtxt nm args (h:t) stk) = evalExp (PrimAppCtxt nm args t stk) h
evalPrimApp LetContext {} = throwInterpError MismatchedStackContext
evalPrimApp FunAppCtxt {} = throwInterpError MismatchedStackContext
evalPrimApp SCEmpty = throwInterpError MismatchedStackContext


applyStack :: (Ord ty,Show ty) =>
              ExpContext ty Ref -> (HeapVal (Exp ty),Ref) ->
              InterpStack s ty b (HeapVal (Exp ty), Ref)
applyStack SCEmpty p = return p
applyStack (LetContext _mv scp stk) (_v,ref) = evalExp stk (instantiate1 (V ref) scp)
applyStack (FunAppCtxt ls [] stk) (_,ref) = evalClosureApp (FunAppCtxt (ref : ls) [] stk)
applyStack (FunAppCtxt ls (h:t) stk) (_,ref) = evalExp (FunAppCtxt (ref:ls) t stk) h
applyStack (PrimAppCtxt nm revArgs [] stk) (_,ref) = evalPrimApp (PrimAppCtxt nm (ref:revArgs) [] stk)
applyStack (PrimAppCtxt nm revargs (h:t) stk) (_,ref) = evalExp (PrimAppCtxt nm (ref : revargs) t stk) h


lookupPrimAccountBalance :: (Ord ty,Show ty) => Text -> InterpStack s ty b (Maybe Rational)
lookupPrimAccountBalance acctNam = do
      InterpreterOverlay _nextOpId localizedMap <- get
      case Map.lookup acctNam localizedMap of
          Just v -> return $ Just v
          Nothing -> do
              snapshotMap :: (Map.Map Text Rational ) <- ask
              case Map.lookup acctNam snapshotMap of
                Just a -> return $ Just a
                Nothing -> return Nothing

-- this may abort if target account doesn't exist
updatePrimAccountBalanceByAdding :: (Ord ty,Show ty) => Text -> Rational -> InterpStack s ty b ()
updatePrimAccountBalanceByAdding nm amt = do
      currentBalance <- lookupPrimAccountBalance nm
      case currentBalance of
          Nothing -> throwInterpError $ PrimFailure "account doesn't exist "
          Just current -> do
            InterpreterOverlay nextOpId localizedMap <- get
            put $ InterpreterOverlay nextOpId $ (Map.insert nm $! (current + amt)) localizedMap


applyPrim :: (Ord ty,Show ty) => ExpContext ty Ref -> PrimOpId -> [Ref]
          -> InterpStack s ty b (HeapVal (Exp ty), Ref)
applyPrim ctxt opid argsRef = do
  resPair <- applyPrimDemo opid argsRef
  applyStack ctxt resPair

applyRationalArithmetic :: (Ord ty, Show ty)
                        => (Rational -> Rational -> Rational)
                        -> [Ref]
                        -> InterpStack s ty b (HeapVal (Exp ty), Ref)
applyRationalArithmetic f [ref1, ref2] = do
  (ratLit1, _posRatRef) <- lift $ heapRefLookupTransitive ref1
  (ratLit2, _posRatRef) <- lift $ heapRefLookupTransitive ref2
  case (ratLit1, ratLit2) of
    (VLitF (LRational rat1), VLitF (LRational rat2)) -> do
      let val = VLitF $ LRational $ rat1 `f` rat2
      ref <- lift $ heapAllocate val
      return (val,ref)
    b -> error $ "deep invariant failure: bad args to arithmetic primop, the arguments were" ++ show b
applyRationalArithmetic _ refs =
    error $ "deep invariant failure: wrong number of args to arithmetic primop (2 instead of " ++ show (length refs) ++ ")"

applyPrimDemo :: (Ord ty,Show ty) => PrimOpId -> [Ref] -> InterpStack s ty b (HeapVal (Exp ty), Ref)
applyPrimDemo (PrimopId "transfer") [fromRef,toRef,posRatRef,fakeCryptoSigRef] = do
  (fromVal,_reffrom) <- lift $ heapRefLookupTransitive fromRef
  (toVal,_refto) <- lift $ heapRefLookupTransitive toRef
  (posRatVal,_posRatRef) <- lift $ heapRefLookupTransitive posRatRef
  (fakeSigVal,_fakeSigRef) <- lift $ heapRefLookupTransitive fakeCryptoSigRef
  case (fromVal,toVal,posRatVal,fakeSigVal) of
    (VLitF (LText fromNm), VLitF (LText toNm), VLitF (LRational amt), VLitF (LText demoSig))
      | amt < 0 -> error "transfer command was invoked with a negative amount"
      | otherwise -> do
          sourceBalanceM <- lookupPrimAccountBalance fromNm
          targetBalanceM <- lookupPrimAccountBalance toNm
          case (sourceBalanceM,targetBalanceM) of
              (Just _srcB ,Just _targetB) ->
                    do
                      updatePrimAccountBalanceByAdding fromNm (-amt)
                      updatePrimAccountBalanceByAdding toNm amt
                      InterpreterOverlay nextOpId mp <- get
                      tell [(nextOpId, Cmd fromNm toNm amt demoSig)]
                      put $ InterpreterOverlay (succ nextOpId) mp
                      let val = VLitF $ LText $ pack "success"
                      ref <- lift $ heapAllocate val -- this shoudld be unit, but whatever
                      return (val,ref)
              _bad -> throwInterpError $ PrimFailure $ "invalid account(s)" ++ show (fromNm, toNm)

    b -> error $ "deep invariant failure : bad args to transfer primop, the arguments were" ++ show b
applyPrimDemo (PrimopId "+") refs = applyRationalArithmetic (+) refs
applyPrimDemo (PrimopId "-") refs = applyRationalArithmetic (-) refs
applyPrimDemo (PrimopId "*") refs = applyRationalArithmetic (*) refs
applyPrimDemo (PrimopId "/") refs = applyRationalArithmetic (/) refs
applyPrimDemo a b = error $ "called command " ++ show a ++ " with " ++ show (length b) ++ " arguments"
