{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -Wno-incomplete-patterns #-}

module Managed.Probe.ToProbe
  ( toProbe
  , ToProbe(..)
  ) where

import Control.Monad.Catch (MonadThrow)
import Data.Managed
import Data.Typeable (Proxy(..), TypeRep, Typeable, typeOf)
import Managed.Exception (badNumberOfArgs, noParseArg, throwM)
import Managed.Probe.Internal.Params (paramsCnt)

-- | Converts any suitable function to a 'Probe'
toProbe ::
     forall e fn. (Typeable fn, ToProbe fn e)
  => fn
  -> Probe e
toProbe :: fn -> Probe e
toProbe fn
x =
  let t :: TypeRep
t = fn -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf fn
x
   in Probe :: forall e. ([In e] -> IO (Out e)) -> TypeRep -> Probe e
Probe {typeRep :: TypeRep
typeRep = TypeRep
t, call :: [In e] -> IO (Out e)
call = TypeRep -> ([In e] -> IO (Out e)) -> [In e] -> IO (Out e)
forall a b. TypeRep -> ([a] -> IO b) -> [a] -> IO b
checkArgs TypeRep
t (Proxy e -> fn -> [In e] -> IO (Out e)
forall fn e. ToProbe fn e => Proxy e -> fn -> [In e] -> IO (Out e)
apply (Proxy e
forall k (t :: k). Proxy t
Proxy @e) fn
x)}

-- | Class of functions that can be converted to a Probe
class ToProbe fn e
  where
  -- | Read arguments from a list, apply them to a function, and encode the result
  apply :: Proxy e -> fn -> [In e] -> IO (Out e)

instance {-# OVERLAPPABLE #-} (Encode a e) => ToProbe a e where
  apply :: Proxy e -> a -> [In e] -> IO (Out e)
apply Proxy e
_ a
c [] = Out e -> IO (Out e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Out e -> IO (Out e)) -> Out e -> IO (Out e)
forall a b. (a -> b) -> a -> b
$ (a -> Out e
forall t rep. Encode t rep => t -> Out rep
encode @a @e) a
c

instance {-# OVERLAPPING #-} (Encode a e) => ToProbe (IO a) e where
  apply :: Proxy e -> IO a -> [In e] -> IO (Out e)
apply Proxy e
_ IO a
c [] = (Encode a e => a -> Out e
forall t rep. Encode t rep => t -> Out rep
encode @a @e) (a -> Out e) -> IO a -> IO (Out e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO a
c

instance {-# OVERLAPPING #-} (Decode a e, ToProbe b e) =>
                             ToProbe (a -> b) e where
  apply :: Proxy e -> (a -> b) -> [In e] -> IO (Out e)
apply Proxy e
_ a -> b
f (In e
x:[In e]
xs) = do
    a
r <- Proxy e -> In e -> IO a
forall a e (m :: * -> *).
(MonadThrow m, Decode a e) =>
Proxy e -> In e -> m a
decodeSingle (Proxy e
forall k (t :: k). Proxy t
Proxy @e) In e
x
    Proxy e -> b -> [In e] -> IO (Out e)
forall fn e. ToProbe fn e => Proxy e -> fn -> [In e] -> IO (Out e)
apply (Proxy e
forall k (t :: k). Proxy t
Proxy @e) (a -> b
f a
r) [In e]
xs

-- Helper functions
decodeSingle ::
     forall a e m. (MonadThrow m, Decode a e)
  => Proxy e
  -> In e
  -> m a
decodeSingle :: Proxy e -> In e -> m a
decodeSingle Proxy e
_ In e
x'
  | (Just a
x) <- (In e -> Maybe a
forall t rep. Decode t rep => In rep -> Maybe t
decode @a @e) In e
x' = a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
x
  | Bool
otherwise = AgentException -> m a
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM AgentException
noParseArg

checkArgs :: TypeRep -> ([a] -> IO b) -> [a] -> IO b
checkArgs :: TypeRep -> ([a] -> IO b) -> [a] -> IO b
checkArgs TypeRep
t = Int -> ([a] -> IO b) -> [a] -> IO b
forall a b. Int -> ([a] -> IO b) -> [a] -> IO b
withArgs (TypeRep -> Int
paramsCnt TypeRep
t)

withArgs :: Int -> ([a] -> IO b) -> [a] -> IO b
withArgs :: Int -> ([a] -> IO b) -> [a] -> IO b
withArgs Int
n [a] -> IO b
f [a]
args
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
args = [a] -> IO b
f [a]
args
  | Bool
otherwise = AgentException -> IO b
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (AgentException -> IO b) -> AgentException -> IO b
forall a b. (a -> b) -> a -> b
$ Int -> [a] -> AgentException
forall a. Int -> [a] -> AgentException
badNumberOfArgs Int
n [a]
args