{-# LANGUAGE ScopedTypeVariables, StandaloneDeriving, DeriveGeneric,
    TupleSections, RecordWildCards, InstanceSigs #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}

-- |
-- Running TH splices
--
module GHCi.TH
  ( startTH
  , runModFinalizerRefs
  , runTH
  , GHCiQException(..)
  ) where

import GHCi.Message
import GHCi.RemoteTypes
import GHC.Serialized

import Control.Exception
import qualified Control.Monad.Fail as Fail
import Data.Binary
import Data.Binary.Put
import Data.ByteString (ByteString)
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as LB
import Data.Data
import Data.Dynamic
import Data.IORef
import Data.Map (Map)
import qualified Data.Map as M
import Data.Maybe
import GHC.Desugar
import qualified Language.Haskell.TH        as TH
import qualified Language.Haskell.TH.Syntax as TH
import Unsafe.Coerce

initQState :: Pipe -> QState
initQState p = QState M.empty Nothing p

newtype GHCiQ a = GHCiQ { runGHCiQ :: QState -> IO (a, QState) }

data GHCiQException = GHCiQException QState String
  deriving (Show, Typeable)

instance Exception GHCiQException

instance Functor GHCiQ where
  fmap f (GHCiQ s) = GHCiQ $ fmap (\(x,s') -> (f x,s')) . s

instance Applicative GHCiQ where
  f <*> a = GHCiQ $ \s ->
    do (f',s')  <- runGHCiQ f s
       (a',s'') <- runGHCiQ a s'
       return (f' a', s'')
  pure x = GHCiQ (\s -> return (x,s))

instance Monad GHCiQ where
  m >>= f = GHCiQ $ \s ->
    do (m', s')  <- runGHCiQ m s
       (a,  s'') <- runGHCiQ (f m') s'
       return (a, s'')

  fail = Fail.fail

instance Fail.MonadFail GHCiQ where
  fail err  = GHCiQ $ \s -> throwIO (GHCiQException s err)

getState :: GHCiQ QState
getState = GHCiQ $ \s -> return (s,s)

noLoc :: TH.Loc
noLoc = TH.Loc "<no file>" "<no package>" "<no module>" (0,0) (0,0)

ghcCmd :: Binary a => Message (THResult a) -> GHCiQ a
ghcCmd m = GHCiQ $ \s -> do
  r <- remoteCall (qsPipe s) m
  case r of
    THException str -> throwIO (GHCiQException s str)
    THComplete res -> return (res, s)

instance TH.Quasi GHCiQ where
  qNewName str = ghcCmd (NewName str)
  qReport isError msg = ghcCmd (Report isError msg)

  -- See Note [TH recover with -fexternal-interpreter] in TcSplice
  qRecover (GHCiQ h) (GHCiQ a) = GHCiQ $ \s -> (do
    remoteCall (qsPipe s) StartRecover
    (r, s') <- a s
    remoteCall (qsPipe s) (EndRecover False)
    return (r,s'))
      `catch`
       \GHCiQException{} -> remoteCall (qsPipe s) (EndRecover True) >> h s
  qLookupName isType occ = ghcCmd (LookupName isType occ)
  qReify name = ghcCmd (Reify name)
  qReifyFixity name = ghcCmd (ReifyFixity name)
  qReifyInstances name tys = ghcCmd (ReifyInstances name tys)
  qReifyRoles name = ghcCmd (ReifyRoles name)

  -- To reify annotations, we send GHC the AnnLookup and also the
  -- TypeRep of the thing we're looking for, to avoid needing to
  -- serialize irrelevant annotations.
  qReifyAnnotations :: forall a . Data a => TH.AnnLookup -> GHCiQ [a]
  qReifyAnnotations lookup =
    map (deserializeWithData . B.unpack) <$>
      ghcCmd (ReifyAnnotations lookup typerep)
    where typerep = typeOf (undefined :: a)

  qReifyModule m = ghcCmd (ReifyModule m)
  qReifyConStrictness name = ghcCmd (ReifyConStrictness name)
  qLocation = fromMaybe noLoc . qsLocation <$> getState
  qRunIO m = GHCiQ $ \s -> fmap (,s) m
  qAddDependentFile file = ghcCmd (AddDependentFile file)
  qAddTopDecls decls = ghcCmd (AddTopDecls decls)
  qAddModFinalizer fin = GHCiQ (\s -> mkRemoteRef fin >>= return . (, s)) >>=
                         ghcCmd . AddModFinalizer
  qGetQ = GHCiQ $ \s ->
    let lookup :: forall a. Typeable a => Map TypeRep Dynamic -> Maybe a
        lookup m = fromDynamic =<< M.lookup (typeOf (undefined::a)) m
    in return (lookup (qsMap s), s)
  qPutQ k = GHCiQ $ \s ->
    return ((), s { qsMap = M.insert (typeOf k) (toDyn k) (qsMap s) })
  qIsExtEnabled x = ghcCmd (IsExtEnabled x)
  qExtsEnabled = ghcCmd ExtsEnabled

startTH :: IO (RemoteRef (IORef QState))
startTH = do
  r <- newIORef (initQState (error "startTH: no pipe"))
  mkRemoteRef r

-- | Runs the mod finalizers.
--
-- The references must be created on the caller process.
runModFinalizerRefs :: Pipe -> RemoteRef (IORef QState)
                    -> [RemoteRef (TH.Q ())]
                    -> IO ()
runModFinalizerRefs pipe rstate qrefs = do
  qs <- mapM localRef qrefs
  qstateref <- localRef rstate
  qstate <- readIORef qstateref
  _ <- runGHCiQ (TH.runQ $ sequence_ qs) qstate { qsPipe = pipe }
  return ()

runTH
  :: Pipe -> RemoteRef (IORef QState) -> HValueRef
  -> THResultType
  -> Maybe TH.Loc
  -> IO ByteString
runTH pipe rstate rhv ty mb_loc = do
  hv <- localRef rhv
  case ty of
    THExp -> runTHQ pipe rstate mb_loc (unsafeCoerce hv :: TH.Q TH.Exp)
    THPat -> runTHQ pipe rstate mb_loc (unsafeCoerce hv :: TH.Q TH.Pat)
    THType -> runTHQ pipe rstate mb_loc (unsafeCoerce hv :: TH.Q TH.Type)
    THDec -> runTHQ pipe rstate mb_loc (unsafeCoerce hv :: TH.Q [TH.Dec])
    THAnnWrapper -> do
      hv <- unsafeCoerce <$> localRef rhv
      case hv :: AnnotationWrapper of
        AnnotationWrapper thing -> return $!
          LB.toStrict (runPut (put (toSerialized serializeWithData thing)))

runTHQ
  :: Binary a => Pipe -> RemoteRef (IORef QState) -> Maybe TH.Loc -> TH.Q a
  -> IO ByteString
runTHQ pipe@Pipe{..} rstate mb_loc ghciq = do
  qstateref <- localRef rstate
  qstate <- readIORef qstateref
  let st = qstate { qsLocation = mb_loc, qsPipe = pipe }
  (r,new_state) <- runGHCiQ (TH.runQ ghciq) st
  writeIORef qstateref new_state
  return $! LB.toStrict (runPut (put r))