{-# LANGUAGE DataKinds #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeOperators #-}

module MyLib (defaultMain, test) where

import AI
import Control.Exception.Annotated.UnliftIO (Exception, SomeException, catch, throw)
import Control.Lens (over, re, strict, view, (&), (.~), (^.))
import Control.Monad (forM_, forever)
import Data.Aeson (decode, encode)
import Data.Either (fromRight)
import Data.Maybe (fromMaybe)
import Data.Set qualified as S
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.IO qualified as T
import Data.Text.Strict.Lens (utf8)
import Embeddings (indexEntries, search)
import Katip qualified as K
import M
import Network.Wai.Handler.Warp (defaultSettings, runSettings, setPort)
import Options.Generic (Generic, ParseRecord (..), getRecord, unDefValue, unHelpful, type (<!>), type (<?>))
import Pipes (Pipe, for, lift, runEffect, (>->))
import Pipes.Prelude qualified as P
import Server (Health (..), app)
import System.Directory (createDirectoryIfMissing, getHomeDirectory)
import System.Environment (lookupEnv)
import System.FilePath (takeDirectory)
import Text.Pandoc (PandocError, def, readHtml, runPure, writeMarkdown)
import Text.Printf (printf)
import UnliftIO (MonadUnliftIO, liftIO)
import UnliftIO.Async (mapConcurrently_)
import UnliftIO.Concurrent (threadDelay)
import UnliftIO.Exception (displayException, isAsyncException)
import UnliftIO.IORef (IORef, newIORef, writeIORef)
import Wallabag
import Pipes.Concurrent (performGC)
import qualified Version

data Options
  = Quick {verbose :: Bool}
  | Full {verbose :: Bool}
  | Version
  | Index
      { verbose :: Bool,
        rebuild :: Bool <?> "Rebuild index from scratch" <!> "False",
        embeddingsDb :: Maybe FilePath <?> "Path to embeddings database"
      }
  | Search
      { verbose :: Bool,
        query :: Text <?> "Search query",
        limit :: Int <?> "Maximum results" <!> "10",
        embeddingsDb :: Maybe FilePath <?> "Path to embeddings database"
      }
  | Daemon
      { verbose :: Bool,
        credentials :: FilePath <?> "Path to the credentials.json",
        port :: Int <?> "Port to run the diagnostics" <!> "8094",
        embeddingsDb :: Maybe FilePath <?> "Path to embeddings database"
      }
  deriving (Generic, Show)

instance ParseRecord Options

tooShort :: Entry -> Bool
tooShort e = length (T.lines (view content e)) < 5

missingTags :: Entry -> Bool
missingTags Entry {_entryTags = []} = True
missingTags Entry {} = False

htmlToMarkdown :: Text -> Either PandocError Text
htmlToMarkdown input = runPure $ do
  doc <- readHtml def input
  writeMarkdown def doc

processableEntries :: (Monad m) => Pipe Entry Entry m ()
processableEntries = P.filter (not . tooShort) >-> P.filter missingTags

-- Convert html into markdown
markdownEntry :: (Monad m) => Pipe Entry Entry m ()
markdownEntry = P.map (over content (\x -> fromRight x $ htmlToMarkdown x))

run :: K.Severity -> Scan -> IO ()
run sev scan = initialS Nothing sev $ \s -> do
  runM s $ runScan scan

runScan :: Scan -> M ()
runScan scan = do
  Pipes.runEffect $
    Pipes.for
      (sourceEntries scan >-> markdownEntry >-> processableEntries)
      (lift . doWork)
  liftIO performGC

runDaemon :: Int -> FilePath -> K.Severity -> IO ()
runDaemon diagnosticsPort credsPath sev = initialS (Just credsPath) sev $ \s -> runM s $ do
  let serverSettings = setPort diagnosticsPort defaultSettings
  health <- newIORef OK

  mapConcurrently_
    id
    [ backgroundRunner hourly $ worker health Unarchived,
      backgroundRunner weekly $ worker health Everything,
      worker health Unarchived, -- Do a quick scan on startup
      liftIO $ runSettings serverSettings (app s health)
    ]
  K.logFM K.InfoS "finished"
  where
    worker health scan = handle (reportError health) $ do
      runScan scan
      writeIORef health OK
    weekly = 7 * 86400 * 1_000_000
    hourly = 60 * 60 * 1_000_000
    backgroundRunner delay action = forever $ do
      threadDelay delay
      action
    reportError :: (K.KatipContext m, K.Katip m, MonadUnliftIO m) => IORef Health -> SomeException -> m ()
    reportError health e
      | isAsyncException e = throw e
      | otherwise = do
          writeIORef health Degraded
          K.logFM K.ErrorS $ K.ls $ displayException e
    handle :: (Exception e, MonadUnliftIO m) => (e -> m a) -> m a -> m a
    handle = flip catch

test :: IO ()
test = run K.DebugS Unarchived

suggestTags :: Text -> M [Text]
suggestTags txt = do
  c <- view creds
  let prompt =
        T.unlines
          [ "Assign tags for the given article from this set of tags: " <> encode hardLabels ^. strict . utf8,
            "Respond in JSON. Respond only in JSON, nothing more",
            "Example response: [\"Politics\"]",
            "",
            "The article contents:",
            txt
          ]
  unparsed <- completions c prompt
  pure $ maybe [] parseTags unparsed
  where
    parseTags :: Text -> [Text]
    parseTags = fromMaybe [] . decode . view (re utf8 . re strict) . T.dropWhileEnd (/= ']') . T.dropWhile (/= '[')

suggestTagsForEntry :: Entry -> M Entry
suggestTagsForEntry e = do
  newTags <- suggestTags (e ^. content)
  pure $ e & tags .~ newTags

doWork :: Entry -> M ()
doWork e = do
  entryWithTags <- suggestTagsForEntry e
  applyTags entryWithTags

hardLabels :: S.Set Text
hardLabels =
  S.fromList
    [ "Politics",
      "Writing",
      "Nature",
      "Camping",
      "Hiking",
      "Economics",
      "Climate",
      "Software",
      "Programming",
      "Maths",
      "Animals",
      "Health",
      "Sexuality",
      "Relationships",
      "Media and television",
      "Social media",
      "Gaming",
      "Plants",
      "Gadgets",
      "AI",
      "Emergency",
      "War",
      "Open Source"
    ]

resolveEmbeddingsDb :: Maybe FilePath -> IO FilePath
resolveEmbeddingsDb (Just fp) = pure fp
resolveEmbeddingsDb Nothing = do
  env <- lookupEnv "WALLABAG_EMBEDDINGS_DB"
  case env of
    Just fp -> pure fp
    Nothing -> do
      home <- getHomeDirectory
      let dir = home <> "/.wallabag-tagger"
      createDirectoryIfMissing True dir
      pure $ dir <> "/embeddings.db"

runWithDb :: Maybe FilePath -> K.Severity -> (FilePath -> M ()) -> IO ()
runWithDb mbDb sev f = initialS Nothing sev $ \s -> runM s $ do
  dbPath <- liftIO $ resolveEmbeddingsDb mbDb
  liftIO $ createDirectoryIfMissing True (takeDirectory dbPath)
  f dbPath

defaultMain :: IO ()
defaultMain = do
  opts <- getRecord "Wallabag tagger"
  case opts of
    Quick {verbose} -> run (severity verbose) Unarchived
    Full {verbose} -> run (severity verbose) Everything
    Version -> print Version.compileTimeVersion
    Index {verbose, rebuild, embeddingsDb} ->
      runWithDb (unHelpful embeddingsDb) (severity verbose) $ \db ->
        indexEntries db (unHelpful $ unDefValue rebuild)
    Search {verbose, query, limit, embeddingsDb} ->
      runWithDb (unHelpful embeddingsDb) (severity verbose) $ \db -> do
        entries <- search db (unHelpful query) (unHelpful $ unDefValue limit)
        liftIO $ forM_ entries $ \(score, _eid, u, t) -> do
          T.putStrLn $ T.pack $ printf "%.4f  %s" (score :: Double) (T.unpack t)
          T.putStrLn $ T.pack $ printf "        %s" (T.unpack u)
    Daemon {verbose, credentials, port, embeddingsDb} ->
      runDaemon (unHelpful $ unDefValue port) (unHelpful credentials) (severity verbose)
  where
    severity True = K.DebugS
    severity _ = K.InfoS