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

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

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

-- | 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 forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: Type -> Type) a. MonadIO m => a -> m (TVar a)
newTVarIO 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) = forall (m :: Type -> Type) a. MonadIO m => STM a -> m a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TVar a -> a -> STM ()
writeTVar TVar (Maybe LByteString)
var forall a. Monoid a => a
mempty

readOrBuildStorkIndex :: (MonadIO m, MonadLoggerIO m) => IndexVar -> Config -> m LByteString
readOrBuildStorkIndex :: forall (m :: Type -> Type).
(MonadIO m, MonadLoggerIO m) =>
IndexVar -> Config -> m LByteString
readOrBuildStorkIndex (IndexVar TVar (Maybe LByteString)
indexVar) Config
config = do
  forall (m :: Type -> Type) a. MonadIO m => TVar a -> m a
readTVarIO TVar (Maybe LByteString)
indexVar forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just LByteString
index -> do
      forall (m :: Type -> Type). MonadLogger m => Text -> m ()
logD Text
"STORK: Returning cached search index"
      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.
      forall (m :: Type -> Type). MonadLogger m => Text -> m ()
logW Text
"STORK: Generating search index (this may be expensive)"
      (Double
diff, !LByteString
index) <- forall (m :: Type -> Type) b. MonadIO m => m b -> m (Double, b)
timeIt forall a b. (a -> b) -> a -> b
$ forall (m :: Type -> Type). MonadIO m => Config -> m LByteString
runStork Config
config
      forall (m :: Type -> Type). MonadLogger m => Text -> m ()
log forall a b. (a -> b) -> a -> b
$ forall a. ToText a => a -> Text
toText forall a b. (a -> b) -> a -> b
$ String
"STORK: Done generating search index in " forall a. Semigroup a => a -> a -> a
<> forall a. RealFloat a => Maybe Int -> a -> ShowS
showGFloat (forall a. a -> Maybe a
Just Int
2) Double
diff String
"" forall a. Semigroup a => a -> a -> a
<> String
" seconds"
      forall (m :: Type -> Type) a. MonadIO m => STM a -> m a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (Maybe LByteString)
indexVar forall a b. (a -> b) -> a -> b
$ \Maybe LByteString
_ -> forall a. a -> Maybe a
Just LByteString
index
      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 <- forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO IO UTCTime
getCurrentTime
      !b
x <- m b
m
      UTCTime
t1 <- 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
      forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (forall a b. (Real a, Fractional b) => a -> b
realToFrac NominalDiffTime
diff, b
x)

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

runStork :: MonadIO m => Config -> m LByteString
runStork :: forall (m :: Type -> Type). MonadIO m => Config -> m LByteString
runStork Config
config = do
  let storkToml :: Text
storkToml = Text -> Text
handleTomlandBug forall a b. (a -> b) -> a -> b
$ forall a. TomlCodec a -> a -> Text
Toml.encode TomlCodec Config
configCodec Config
config
  (ExitCode
_, !ByteString
index, ByteString
_) <-
    forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO 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"]
        (forall a b. ConvertUtf8 a b => a -> b
encodeUtf8 Text
storkToml)
  forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ 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/srid/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 Config = Config
  { Config -> Input
configInput :: Input
  }
  deriving stock (Config -> Config -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Config -> Config -> Bool
$c/= :: Config -> Config -> Bool
== :: Config -> Config -> Bool
$c== :: Config -> Config -> Bool
Eq, Int -> Config -> ShowS
[Config] -> ShowS
Config -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Config] -> ShowS
$cshowList :: [Config] -> ShowS
show :: Config -> String
$cshow :: Config -> String
showsPrec :: Int -> Config -> ShowS
$cshowsPrec :: Int -> Config -> ShowS
Show)

data Input = Input
  { Input -> [File]
inputFiles :: [File]
  , Input -> Handling
inputFrontmatterHandling :: Handling
  }
  deriving stock (Input -> Input -> Bool
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 -> ShowS
[Input] -> ShowS
Input -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Input] -> ShowS
$cshowList :: [Input] -> ShowS
show :: Input -> String
$cshow :: Input -> String
showsPrec :: Int -> Input -> ShowS
$cshowsPrec :: Int -> Input -> ShowS
Show)

data File = File
  { File -> String
filePath :: FilePath
  , File -> Text
fileUrl :: Text
  , File -> Text
fileTitle :: Text
  }
  deriving stock (File -> File -> Bool
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 -> ShowS
[File] -> ShowS
File -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [File] -> ShowS
$cshowList :: [File] -> ShowS
show :: File -> String
$cshow :: File -> String
showsPrec :: Int -> File -> ShowS
$cshowsPrec :: Int -> File -> ShowS
Show)

data Handling
  = Handling_Ignore
  | Handling_Omit
  | Handling_Parse
  deriving stock (Handling -> Handling -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Handling -> Handling -> Bool
$c/= :: Handling -> Handling -> Bool
== :: Handling -> Handling -> Bool
$c== :: Handling -> Handling -> Bool
Eq, Int -> Handling -> ShowS
[Handling] -> ShowS
Handling -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Handling] -> ShowS
$cshowList :: [Handling] -> ShowS
show :: Handling -> String
$cshow :: Handling -> String
showsPrec :: Int -> Handling -> ShowS
$cshowsPrec :: Int -> Handling -> ShowS
Show, forall x. Rep Handling x -> Handling
forall x. Handling -> Rep Handling x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Handling x -> Handling
$cfrom :: forall x. Handling -> Rep Handling x
Generic)
  deriving
    (Value -> Parser [Handling]
Value -> Parser Handling
forall a.
(Value -> Parser a) -> (Value -> Parser [a]) -> FromJSON a
parseJSONList :: Value -> Parser [Handling]
$cparseJSONList :: Value -> Parser [Handling]
parseJSON :: Value -> Parser Handling
$cparseJSON :: Value -> Parser Handling
FromJSON)
    via CustomJSON
          '[ ConstructorTagModifier '[StripPrefix "Handling_", CamelToSnake]
           ]
          Handling

instance Default Handling where
  def :: Handling
def = Handling
Handling_Omit

configCodec :: TomlCodec Config
configCodec :: TomlCodec Config
configCodec =
  Input -> Config
Config
    forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. TomlCodec a -> Key -> TomlCodec a
Toml.table TomlCodec Input
inputCodec Key
"input"
      forall field a object.
Codec field a -> (object -> field) -> Codec object a
.= Config -> Input
configInput
  where
    inputCodec :: TomlCodec Input
    inputCodec :: TomlCodec Input
inputCodec =
      [File] -> Handling -> Input
Input
        forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. TomlCodec a -> Key -> TomlCodec [a]
Toml.list TomlCodec File
fileCodec Key
"files"
          forall field a object.
Codec field a -> (object -> field) -> Codec object a
.= Input -> [File]
inputFiles
        forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> forall b a. Coercible @Type a b => TomlCodec a -> TomlCodec b
Toml.diwrap (Key -> TomlCodec Handling
handlingCodec Key
"frontmatter_handling")
          forall field a object.
Codec field a -> (object -> field) -> Codec object a
.= Input -> Handling
inputFrontmatterHandling
    fileCodec :: TomlCodec File
    fileCodec :: TomlCodec File
fileCodec =
      String -> Text -> Text -> File
File
        forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Key -> TomlCodec String
Toml.string Key
"path"
          forall field a object.
Codec field a -> (object -> field) -> Codec object a
.= File -> String
filePath
        forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Key -> TomlCodec Text
Toml.text Key
"url"
          forall field a object.
Codec field a -> (object -> field) -> Codec object a
.= File -> Text
fileUrl
        forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Key -> TomlCodec Text
Toml.text Key
"title"
          forall field a object.
Codec field a -> (object -> field) -> Codec object a
.= File -> Text
fileTitle
    handlingCodec :: Toml.Key -> TomlCodec Handling
    handlingCodec :: Key -> TomlCodec Handling
handlingCodec = forall a.
(a -> Text) -> (Text -> Either Text a) -> Key -> TomlCodec a
textBy Handling -> Text
showHandling Text -> Either Text Handling
parseHandling
      where
        showHandling :: Handling -> Text
        showHandling :: Handling -> Text
showHandling Handling
handling = case Handling
handling of
          Handling
Handling_Ignore -> Text
"Ignore"
          Handling
Handling_Omit -> Text
"Omit"
          Handling
Handling_Parse -> Text
"Parse"
        parseHandling :: Text -> Either Text Handling
        parseHandling :: Text -> Either Text Handling
parseHandling Text
handling = case Text
handling of
          Text
"Ignore" -> forall a b. b -> Either a b
Right Handling
Handling_Ignore
          Text
"Omit" -> forall a b. b -> Either a b
Right Handling
Handling_Omit
          Text
"Parse" -> forall a b. b -> Either a b
Right Handling
Handling_Parse
          Text
other -> forall a b. a -> Either a b
Left forall a b. (a -> b) -> a -> b
$ Text
"Unsupported value for frontmatter handling: " forall a. Semigroup a => a -> a -> a
<> Text
other