{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TemplateHaskell #-}

module Emanote.Model.Stork.Index
  ( IndexVar,
    newIndex,
    clearStorkIndex,
    readOrBuildStorkIndex,
    File (File),
    Input (Input),
  )
where

import Control.Monad.Logger (MonadLoggerIO)
import Data.Text qualified as T
import Data.Time (NominalDiffTime, diffUTCTime, getCurrentTime)
import Emanote.Prelude (log, logD, logW)
import Numeric (showGFloat)
import Relude
import System.Process.ByteString (readProcessWithExitCode)
import System.Which (staticWhich)
import Toml (TomlCodec, encode, list, string, text, (.=))

-- | In-memory Stork index tracked in a @TVar@
newtype IndexVar = IndexVar (TVar (Maybe LByteString))

newIndex :: MonadIO m => m IndexVar
newIndex :: forall (m :: Type -> Type). MonadIO m => m IndexVar
newIndex =
  TVar (Maybe LByteString) -> IndexVar
IndexVar (TVar (Maybe LByteString) -> IndexVar)
-> m (TVar (Maybe LByteString)) -> m IndexVar
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe LByteString -> m (TVar (Maybe LByteString))
forall (m :: Type -> Type) a. MonadIO m => a -> m (TVar a)
newTVarIO Maybe LByteString
forall a. Monoid a => a
mempty

clearStorkIndex :: (MonadIO m) => IndexVar -> m ()
clearStorkIndex :: forall (m :: Type -> Type). MonadIO m => IndexVar -> m ()
clearStorkIndex (IndexVar TVar (Maybe LByteString)
var) = STM () -> m ()
forall (m :: Type -> Type) a. MonadIO m => STM a -> m a
atomically (STM () -> m ()) -> STM () -> m ()
forall a b. (a -> b) -> a -> b
$ TVar (Maybe LByteString) -> Maybe LByteString -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Maybe LByteString)
var Maybe LByteString
forall a. Monoid a => a
mempty

readOrBuildStorkIndex :: (MonadIO m, MonadLoggerIO m) => IndexVar -> Input -> m LByteString
readOrBuildStorkIndex :: forall (m :: Type -> Type).
(MonadIO m, MonadLoggerIO m) =>
IndexVar -> Input -> m LByteString
readOrBuildStorkIndex (IndexVar TVar (Maybe LByteString)
indexVar) Input
input = do
  TVar (Maybe LByteString) -> m (Maybe LByteString)
forall (m :: Type -> Type) a. MonadIO m => TVar a -> m a
readTVarIO TVar (Maybe LByteString)
indexVar m (Maybe LByteString)
-> (Maybe LByteString -> m LByteString) -> m LByteString
forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just LByteString
index -> do
      Text -> m ()
forall (m :: Type -> Type). MonadLogger m => Text -> m ()
logD Text
"STORK: Returning cached search index"
      LByteString -> m LByteString
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure LByteString
index
    Maybe LByteString
Nothing -> do
      -- TODO: What if there are concurrent reads? We probably need a lock.
      -- And we want to encapsulate this whole thing.
      Text -> m ()
forall (m :: Type -> Type). MonadLogger m => Text -> m ()
logW Text
"STORK: Generating search index (this may be expensive)"
      (Double
diff, !LByteString
index) <- m LByteString -> m (Double, LByteString)
forall (m :: Type -> Type) b. MonadIO m => m b -> m (Double, b)
timeIt (m LByteString -> m (Double, LByteString))
-> m LByteString -> m (Double, LByteString)
forall a b. (a -> b) -> a -> b
$ Input -> m LByteString
forall (m :: Type -> Type). MonadIO m => Input -> m LByteString
runStork Input
input
      Text -> m ()
forall (m :: Type -> Type). MonadLogger m => Text -> m ()
log (Text -> m ()) -> Text -> m ()
forall a b. (a -> b) -> a -> b
$ String -> Text
forall a. ToText a => a -> Text
toText (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ String
"STORK: Done generating search index in " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Maybe Int -> Double -> String -> String
forall a. RealFloat a => Maybe Int -> a -> String -> String
showGFloat (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
2) Double
diff String
"" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" seconds"
      STM () -> m ()
forall (m :: Type -> Type) a. MonadIO m => STM a -> m a
atomically (STM () -> m ()) -> STM () -> m ()
forall a b. (a -> b) -> a -> b
$ TVar (Maybe LByteString)
-> (Maybe LByteString -> Maybe LByteString) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (Maybe LByteString)
indexVar ((Maybe LByteString -> Maybe LByteString) -> STM ())
-> (Maybe LByteString -> Maybe LByteString) -> STM ()
forall a b. (a -> b) -> a -> b
$ \Maybe LByteString
_ -> LByteString -> Maybe LByteString
forall a. a -> Maybe a
Just LByteString
index
      LByteString -> m LByteString
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure LByteString
index
  where
    timeIt :: MonadIO m => m b -> m (Double, b)
    timeIt :: forall (m :: Type -> Type) b. MonadIO m => m b -> m (Double, b)
timeIt m b
m = do
      UTCTime
t0 <- IO UTCTime -> m UTCTime
forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO IO UTCTime
getCurrentTime
      !b
x <- m b
m
      UTCTime
t1 <- IO UTCTime -> m UTCTime
forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO IO UTCTime
getCurrentTime
      let NominalDiffTime
diff :: NominalDiffTime = UTCTime -> UTCTime -> NominalDiffTime
diffUTCTime UTCTime
t1 UTCTime
t0
      (Double, b) -> m (Double, b)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (NominalDiffTime -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac NominalDiffTime
diff, b
x)

storkBin :: FilePath
storkBin :: String
storkBin = $(staticWhich "stork")

runStork :: MonadIO m => Input -> m LByteString
runStork :: forall (m :: Type -> Type). MonadIO m => Input -> m LByteString
runStork Input
input = do
  let storkToml :: Text
storkToml = Text -> Text
handleTomlandBug (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ TomlCodec Input -> Input -> Text
forall a. TomlCodec a -> a -> Text
Toml.encode TomlCodec Input
inputCodec Input
input
  (ExitCode
_, !ByteString
index, ByteString
_) <-
    IO (ExitCode, ByteString, ByteString)
-> m (ExitCode, ByteString, ByteString)
forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO (IO (ExitCode, ByteString, ByteString)
 -> m (ExitCode, ByteString, ByteString))
-> IO (ExitCode, ByteString, ByteString)
-> m (ExitCode, ByteString, ByteString)
forall a b. (a -> b) -> a -> b
$
      String
-> [String] -> ByteString -> IO (ExitCode, ByteString, ByteString)
readProcessWithExitCode
        String
storkBin
        -- NOTE: Cannot use "--output -" due to bug in Rust or Stork:
        -- https://github.com/jameslittle230/stork/issues/262
        [String
"build", String
"-t", String
"--input", String
"-", String
"--output", String
"/dev/stdout"]
        (Text -> ByteString
forall a b. ConvertUtf8 a b => a -> b
encodeUtf8 Text
storkToml)
  LByteString -> m LByteString
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (LByteString -> m LByteString) -> LByteString -> m LByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> LByteString
forall l s. LazyStrict l s => s -> l
toLazy ByteString
index
  where
    handleTomlandBug :: Text -> Text
handleTomlandBug =
      -- HACK: Deal with tomland's bug.
      -- https://github.com/EmaApps/emanote/issues/336
      -- https://github.com/kowainik/tomland/issues/408
      --
      -- This could be problematic if the user literally uses \\U in their note
      -- title (but why would they?)
      Text -> Text -> Text -> Text
T.replace Text
"\\\\U" Text
"\\U"

newtype Input = Input
  { Input -> [File]
inputFiles :: [File]
  }
  deriving stock (Input -> Input -> Bool
(Input -> Input -> Bool) -> (Input -> Input -> Bool) -> Eq Input
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Input -> Input -> Bool
$c/= :: Input -> Input -> Bool
== :: Input -> Input -> Bool
$c== :: Input -> Input -> Bool
Eq, Int -> Input -> String -> String
[Input] -> String -> String
Input -> String
(Int -> Input -> String -> String)
-> (Input -> String) -> ([Input] -> String -> String) -> Show Input
forall a.
(Int -> a -> String -> String)
-> (a -> String) -> ([a] -> String -> String) -> Show a
showList :: [Input] -> String -> String
$cshowList :: [Input] -> String -> String
show :: Input -> String
$cshow :: Input -> String
showsPrec :: Int -> Input -> String -> String
$cshowsPrec :: Int -> Input -> String -> String
Show)

data File = File
  { File -> String
filePath :: FilePath,
    File -> Text
fileUrl :: Text,
    File -> Text
fileTitle :: Text
  }
  deriving stock (File -> File -> Bool
(File -> File -> Bool) -> (File -> File -> Bool) -> Eq File
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: File -> File -> Bool
$c/= :: File -> File -> Bool
== :: File -> File -> Bool
$c== :: File -> File -> Bool
Eq, Int -> File -> String -> String
[File] -> String -> String
File -> String
(Int -> File -> String -> String)
-> (File -> String) -> ([File] -> String -> String) -> Show File
forall a.
(Int -> a -> String -> String)
-> (a -> String) -> ([a] -> String -> String) -> Show a
showList :: [File] -> String -> String
$cshowList :: [File] -> String -> String
show :: File -> String
$cshow :: File -> String
showsPrec :: Int -> File -> String -> String
$cshowsPrec :: Int -> File -> String -> String
Show)

fileCodec :: TomlCodec File
fileCodec :: TomlCodec File
fileCodec =
  String -> Text -> Text -> File
File
    (String -> Text -> Text -> File)
-> Codec File String -> Codec File (Text -> Text -> File)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Key -> TomlCodec String
Toml.string Key
"path" TomlCodec String -> (File -> String) -> Codec File String
forall field a object.
Codec field a -> (object -> field) -> Codec object a
.= File -> String
filePath
    Codec File (Text -> Text -> File)
-> Codec File Text -> Codec File (Text -> File)
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Key -> TomlCodec Text
Toml.text Key
"url" TomlCodec Text -> (File -> Text) -> Codec File Text
forall field a object.
Codec field a -> (object -> field) -> Codec object a
.= File -> Text
fileUrl
    Codec File (Text -> File) -> Codec File Text -> TomlCodec File
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Key -> TomlCodec Text
Toml.text Key
"title" TomlCodec Text -> (File -> Text) -> Codec File Text
forall field a object.
Codec field a -> (object -> field) -> Codec object a
.= File -> Text
fileTitle

inputCodec :: TomlCodec Input
inputCodec :: TomlCodec Input
inputCodec =
  [File] -> Input
Input
    ([File] -> Input) -> Codec Input [File] -> TomlCodec Input
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> TomlCodec File -> Key -> TomlCodec [File]
forall a. TomlCodec a -> Key -> TomlCodec [a]
Toml.list TomlCodec File
fileCodec Key
"input.files" TomlCodec [File] -> (Input -> [File]) -> Codec Input [File]
forall field a object.
Codec field a -> (object -> field) -> Codec object a
.= Input -> [File]
inputFiles