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