Worker.hs

module Kioku.Cli.Commands.Worker
  ( WorkerOptions (..),
    workerOptionsParser,
    runWorker,
  )
where

import Baikai.Embedding (EmbeddingModel)
import Control.Concurrent (threadDelay)
import Control.Concurrent.Async (race)
import Control.Exception (SomeException, displayException, try)
import Data.Text qualified as Text
import Data.Time (getCurrentTime)
import Data.UUID qualified as UUID
import Effectful (Eff, IOE, (:>))
import Keiro.Timer qualified as Timer
import Kioku.AI.Config (AIFeature (MemoryEmbedding))
import Kioku.AI.Runtime (AIRuntime)
import Kioku.Api.Access (MemorySpaceId, memorySpaceIdText, mkMemorySpaceId)
import Kioku.App (AppEffects, AppEnv, runAppIO, withNoopAppEnv)
import Kioku.Cli.AIConfig (aiConfigOption, loadAIRuntime)
import Kioku.Cli.Context (cliContextProvider)
import Kioku.Distill.L1 (FindMergeCandidates, recallCandidates)
import Kioku.Distill.Runtime (newDistillRuntime)
import Kioku.Distill.Timer.Deferred
import Kioku.Distill.Timer.Outcome (FireOutcome (..))
import Kioku.Distill.Timer.Worker (drainKiokuTimers, runKiokuTimerWorkerOnce)
import Kioku.Memory.Embedding (resolveEmbeddingConfig)
import Kioku.Memory.Embedding.Worker
  ( EmbeddingBackfillScope (..),
    backfillMissingEmbeddings,
    mkEmbeddingWorkerEnv,
    runEmbeddingWorkerHost,
  )
import Kioku.Recall.Capability (VectorCapability (..), detectVectorCapability)
import Kiroku.Store.Connection (KirokuStore, defaultConnectionSettings, withStore)
import Kiroku.Store.Effect (Store)
import Options.Applicative
import System.Environment (lookupEnv)
import System.Exit (ExitCode (..), exitWith)
import System.IO (hPutStrLn, stderr)

-- | The two one-shot modes are unrelated — an embedding backfill and firing one distillation
-- timer — so there is no combined meaning to define. As two 'switch'es they were silently
-- ordered: @--backfill --timers-once@ checked @timersOnce@ first and ignored @--backfill@
-- without a word. As a sum parsed from mutually exclusive alternatives, passing both is a
-- parse error.
--
-- @--space@ belongs to the backfill and to nothing else, so it is a field of that constructor
-- rather than a top-level option. Its default is 'BackfillEverySpace': a worker serves every
-- space in its database, and defaulting to one — @KIOKU_MEMORY_SPACE@, say — would let an
-- operator run a backfill, see a count, and never learn that the other spaces are still
-- unsearchable.
data WorkerOptions
  = WorkerConfigured !FilePath !WorkerOptions
  | WorkerContinuous
  | WorkerBackfill !EmbeddingBackfillScope
  | WorkerDeferredList
  | WorkerDeferredResume !Timer.TimerId
  | WorkerTimersOnce
  deriving stock (Eq, Show)

workerOptionsParser :: Parser WorkerOptions
workerOptionsParser = (\config mode -> maybe mode (`WorkerConfigured` mode) config) <$> aiConfigOption <*> workerModeParser

workerModeParser :: Parser WorkerOptions
workerModeParser =
  hsubparser
    (command "deferred" (info deferredParser (progDesc "List or resume authorized deferred timers")))
    <|> ( flag'
            WorkerBackfill
            ( long "backfill"
                <> help "Run one embedding backfill pass and exit (conflicts with --timers-once)"
            )
            <*> backfillScopeParser
        )
    <|> flag'
      WorkerTimersOnce
      ( long "timers-once"
          <> help "Claim and fire at most one due kioku distillation timer, then exit (conflicts with --backfill)"
      )
    <|> pure WorkerContinuous

deferredParser :: Parser WorkerOptions
deferredParser =
  hsubparser
    ( command "list" (info (pure WorkerDeferredList) (progDesc "List deferred distillation work"))
        <> command
          "resume"
          ( info
              (WorkerDeferredResume <$> argument timerIdReader (metavar "TIMER_ID"))
              (progDesc "Resume one deferred timer using foreground AI capabilities")
          )
    )
  where
    timerIdReader = eitherReader $ \raw ->
      maybe (Left "TIMER_ID must be a UUID") (Right . Timer.TimerId) (UUID.fromString raw)

backfillScopeParser :: Parser EmbeddingBackfillScope
backfillScopeParser =
  maybe BackfillEverySpace BackfillOneSpace
    <$> optional
      ( option
          (eitherReader parseMemorySpace)
          ( long "space"
              <> metavar "MEMORY_SPACE_ID"
              <> help "Backfill only this memory space (default: every space in the database)"
          )
      )

parseMemorySpace :: String -> Either String MemorySpaceId
parseMemorySpace raw =
  case mkMemorySpaceId (Text.pack raw) of
    Left err -> Left (Text.unpack err)
    Right space -> Right space

runWorker :: WorkerOptions -> IO ()
runWorker opts = case opts of
  WorkerConfigured path mode -> runConfiguredWorker (Just path) mode
  _ -> runConfiguredWorker Nothing opts

runConfiguredWorker :: Maybe FilePath -> WorkerOptions -> IO ()
runConfiguredWorker path opts = do
  ai <- loadAIRuntime (case opts of WorkerDeferredResume _ -> True; _ -> False) path
  case opts of
    WorkerBackfill _ -> either (dieWorker . show) (const (pure ())) (resolveEmbeddingConfig ai MemoryEmbedding)
    _ -> pure ()
  connStr <- requireEnv "PG_CONNECTION_STRING"
  let settings = defaultConnectionSettings (Text.pack connStr)
  withStore settings $ \st ->
    withNoopAppEnv settings \env -> case opts of
      WorkerDeferredList -> runDeferredList env
      WorkerDeferredResume tid -> runDeferredResume ai env tid
      WorkerTimersOnce -> runTimerOnce ai env
      WorkerBackfill scope -> case resolveEmbeddingConfig ai MemoryEmbedding of
        Left err -> dieWorker (show err)
        Right model -> withCapability env model $ \capability -> runBackfill ai env capability scope
      WorkerContinuous -> case resolveEmbeddingConfig ai MemoryEmbedding of
        Left _ -> do
          putStrLn "Memory embeddings disabled by AI policy; running timer worker only."
          runTimerLoop ai env VectorExtensionUnavailable
        Right model -> withCapability env model $ \capability -> runContinuousWorker ai env st capability model
      WorkerConfigured nested mode -> runConfiguredWorker (Just nested) mode

withCapability :: AppEnv -> EmbeddingModel -> (VectorCapability -> IO a) -> IO a
withCapability env _config k = do
  result <- runAppIO env (detectVectorCapability 1536)
  case result of
    Left storeErr -> ioError (userError ("kioku worker store error: " <> show storeErr))
    Right capability -> k capability

-- | Merge candidates come from hybrid recall over the atom's own text, not from
-- a priority-ordered scan prefix: a duplicate ranked below the scan window was
-- invisible to the consolidator and got re-stored forever. Recall degrades to
-- FTS-only without pgvector and to keyword-only when the embedding call fails,
-- so no capability gating is needed here.
mergeCandidateFinder ::
  (IOE :> es, Store :> es) =>
  AIRuntime ->
  VectorCapability ->
  FindMergeCandidates es
mergeCandidateFinder ai capability =
  recallCandidates ai capability mergeCandidateLimit

mergeCandidateLimit :: Int
mergeCandidateLimit = 8

runBackfill :: AIRuntime -> AppEnv -> VectorCapability -> EmbeddingBackfillScope -> IO ()
runBackfill ai env capability scope = do
  -- Refuse before any event is touched: a backfill under a mismatched dimension count would
  -- embed every memory in the store and fail the ::vector cast on every single one.
  case capability of
    VectorDimensionMismatch configured actual ->
      dieWorker (dimensionMismatchMessage configured actual)
    _ -> pure ()
  embeddingEnv <- either (ioError . userError . show) pure (mkEmbeddingWorkerEnv ai)
  result <- runAppIO env (backfillMissingEmbeddings capability embeddingEnv scope)
  case result of
    Left storeErr -> ioError (userError ("kioku worker backfill store error: " <> show storeErr))
    Right count ->
      putStrLn ("Backfilled " <> show count <> " memory embeddings " <> backfillScopeLabel scope <> ".")

backfillScopeLabel :: EmbeddingBackfillScope -> String
backfillScopeLabel = \case
  BackfillEverySpace -> "across every memory space"
  BackfillOneSpace space -> "in memory space " <> Text.unpack (memorySpaceIdText space)

-- | Run both pipelines under supervision.
--
-- The timer loop used to run on a bare 'forkIO', which produced two silent
-- deaths. A store error aborted the loop's single error scope and killed only
-- the forked thread, leaving a process that looked alive while all distillation
-- had stopped. In the other direction, an embedding halt made shibuya exit its
-- processor gracefully, 'waitApp' return, and the whole process exit 0 — taking
-- the timer loop with it.
--
-- 'race' makes both directions loud: whichever pipeline stops first ends the
-- race, and the process exits non-zero with a reason so a supervisor restarts
-- it. Neither side is expected to return at all.
runContinuousWorker :: AIRuntime -> AppEnv -> KirokuStore -> VectorCapability -> EmbeddingModel -> IO ()
runContinuousWorker ai env store capability _config = do
  contexts <- cliContextProvider @(Eff AppEffects)
  case capability of
    VectorAvailable -> do
      startupBackfill ai env capability
      outcome <-
        try @SomeException $
          race
            (runTimerLoop ai env capability)
            (runAppIO env (runEmbeddingWorkerHost store contexts capability ai))
      case outcome of
        -- A halted processor can tear its own machinery down hard enough to
        -- surface as an exception rather than a clean return (shibuya's halt path
        -- can leave a thread blocked in STM). Either way the pipeline is gone;
        -- what matters is that the operator gets a reason and a non-zero exit
        -- instead of a process that looks alive.
        Left err ->
          dieWorker ("worker pipeline crashed: " <> displayException err)
        Right (Left ()) ->
          dieWorker "timer loop stopped unexpectedly"
        Right (Right (Left storeErr)) ->
          dieWorker ("embedding worker stopped with store error: " <> show storeErr)
        Right (Right (Right ())) ->
          -- The handler already printed the halt reason at decision time.
          dieWorker "embedding worker stopped (processor halted or subscription ended)"
    VectorExtensionUnavailable -> do
      putStrLn "pgvector is not available; recall will run FTS-only; running kioku timer worker only."
      runTimerLoop ai env capability
    VectorColumnsUnavailable missing -> do
      putStrLn ("pgvector columns are missing (" <> Text.unpack (Text.intercalate ", " missing) <> "); running kioku timer worker only.")
      runTimerLoop ai env capability
    -- Loud, but not fatal. Every embedding write would fail on the ::vector cast, so there
    -- is no point starting the embedding host — but distillation timers have nothing to do
    -- with embeddings, and killing the whole worker would stop them too.
    VectorDimensionMismatch configured actual -> do
      hPutStrLn stderr ("kioku worker: " <> dimensionMismatchMessage configured actual <> "; running kioku timer worker only.")
      runTimerLoop ai env capability

-- | A dimension mismatch would otherwise be discovered one failed event at a time, forever.
dimensionMismatchMessage :: Int -> Int -> String
dimensionMismatchMessage configured actual =
  "embedding dimension mismatch: KIOKU_EMBEDDING_DIMENSIONS="
    <> show configured
    <> " but kioku.memories.embedding is vector("
    <> show actual
    <> "); fix the env var or migrate the column"

-- | Recover embeddings lost to an outage that outlasted the retry window.
-- Idempotent, so it is safe on every start. A failure here is only a warning:
-- if the database is down, the loops' own retry and exit behavior is the honest
-- place for that to surface, not a special case at startup.
--
-- It covers every space, not @KIOKU_MEMORY_SPACE@: this process is about to subscribe to every
-- space's memory events, so recovering only one space's would leave the others' recall degraded
-- with nothing to say so.
startupBackfill :: AIRuntime -> AppEnv -> VectorCapability -> IO ()
startupBackfill ai env capability = do
  embeddingEnv <- either (ioError . userError . show) pure (mkEmbeddingWorkerEnv ai)
  result <-
    runAppIO
      env
      ( backfillMissingEmbeddings
          capability
          embeddingEnv
          BackfillEverySpace
      )
  case result of
    Left storeErr ->
      hPutStrLn stderr ("kioku worker: startup backfill failed: " <> show storeErr)
    Right count ->
      putStrLn ("Startup backfill: embedded " <> show count <> " missing memory embeddings.")

dieWorker :: String -> IO ()
dieWorker msg = do
  hPutStrLn stderr ("kioku worker: " <> msg <> "; exiting")
  exitWith (ExitFailure 1)

runTimerOnce :: AIRuntime -> AppEnv -> IO ()
runTimerOnce ai env = do
  let rt = newDistillRuntime ai Nothing
  contexts <- cliContextProvider @(Eff AppEffects)
  now <- getCurrentTime
  result <- runAppIO env do
    capability <- detectVectorCapability 1536
    runKiokuTimerWorkerOnce Nothing contexts rt (mergeCandidateFinder ai capability) now
  case result of
    Left storeErr -> ioError (userError ("kioku timer worker store error: " <> show storeErr))
    Right Nothing -> putStrLn "No due kioku distillation timers."
    Right (Just _) -> putStrLn "Processed one due kioku distillation timer."

-- | Drain due timers, sleep, repeat — forever.
--
-- Each pass gets its own 'runAppIO', and therefore its own store-error scope.
-- That is the whole point: the old loop lived inside a single 'runAppIO', so the
-- first transient store error aborted the @forever@ and killed the loop. Here a
-- store error is logged and retried with capped exponential backoff (5s doubling
-- to 60s, reset on success), because a database outage should not require an
-- operator to restart anything — and restarting would not have helped.
--
-- This never returns normally, so 'race' seeing it finish genuinely means
-- something impossible happened. Non-store exceptions propagate to 'race', which
-- is equally loud.
runTimerLoop :: AIRuntime -> AppEnv -> VectorCapability -> IO ()
runTimerLoop ai env capability = do
  let rt = newDistillRuntime ai Nothing
  contexts <- cliContextProvider @(Eff AppEffects)
  putStrLn "kioku timer worker started."
  let go failures = do
        result <- runAppIO env (drainKiokuTimers Nothing contexts rt (mergeCandidateFinder ai capability))
        case result of
          Left storeErr -> do
            hPutStrLn stderr ("kioku timer worker: store error (will retry): " <> show storeErr)
            threadDelay (storeErrorBackoffMicros failures)
            go (failures + 1)
          Right _processed -> do
            threadDelay defaultTimerPollMicros
            go 0
  go 0

-- | 5s doubling per consecutive failure, capped at 60s.
storeErrorBackoffMicros :: Int -> Int
storeErrorBackoffMicros failures =
  min (60 * 1000 * 1000) (5 * 1000 * 1000 * (2 ^ min 8 (max 0 failures)))

-- | With draining, the poll interval no longer caps throughput: a burst of due
-- timers is processed in one pass rather than one per interval.
defaultTimerPollMicros :: Int
defaultTimerPollMicros = 5 * 1000 * 1000

requireEnv :: String -> IO String
requireEnv name = do
  found <- lookupEnv name
  case found of
    Just envValue -> pure envValue
    Nothing -> ioError (userError (name <> " is not set"))

runDeferredList :: AppEnv -> IO ()
runDeferredList env = do
  contexts <- cliContextProvider @(Eff AppEffects)
  let go cursor = do
        result <- runAppIO env (listDeferredTimers contexts (Timer.DeadTimerPageRequest 100 cursor))
        case result of
          Left err -> dieWorker (show err)
          Right (Left err) -> dieWorker (show err)
          Right (Right page) -> do
            mapM_ render page.entries
            maybe (pure ()) (go . Just) page.nextAfterTimerId
      render entry = do
        let Timer.TimerId uuid = entry.timer.timerId
        putStrLn
          ( UUID.toString uuid
              <> " space="
              <> Text.unpack (memorySpaceIdText entry.memorySpace)
              <> " features="
              <> show entry.features
              <> " attempts="
              <> show entry.timer.attempts
              <> " reason="
              <> Text.unpack entry.reason
          )
        putStrLn ("  Resume: kioku worker deferred resume " <> UUID.toString uuid <> " --ai-config FILE")
  go Nothing

runDeferredResume :: AIRuntime -> AppEnv -> Timer.TimerId -> IO ()
runDeferredResume ai env tid = do
  contexts <- cliContextProvider @(Eff AppEffects)
  result <- runAppIO env do
    capability <- detectVectorCapability 1536
    resumeDeferredTimer contexts (newDistillRuntime ai Nothing) (mergeCandidateFinder ai capability) tid
  case result of
    Right (DeferredFinished (FireCompleted _)) -> putStrLn "Completed the original deferred timer."
    Right outcome -> dieWorker (show outcome <> "; unfinished work remains parked; inspect configuration, authorization, and attempt ceiling before retrying")
    Left err -> dieWorker (show err)