ReadModel.hs

module Kioku.Memory.ReadModel
  ( memoryInlineProjection,
    MemoryRow (..),
    MemoryByIdQuery (..),
    MemoriesByNamespaceQuery (..),
    MemoriesByScopeQuery (..),
    MemoriesBySessionQuery (..),
    MemoriesByTypeQuery (..),
    MemorySupersessionChainQuery (..),
    memoryByIdReadModel,
    memoriesByNamespaceReadModel,
    memoriesByNamespaceRowsReadModel,
    memoriesByScopeReadModel,
    memoriesByScopeRowsReadModel,
    memoriesBySessionReadModel,
    memoriesBySessionRowsReadModel,
    memoriesByTypeReadModel,
    memoriesByTypeRowsReadModel,
    memorySupersessionChainReadModel,
  )
where

import Contravariant.Extras (contrazip2, contrazip3)
import Data.Aeson qualified as Aeson
import Data.ByteString.Lazy qualified as BL
import Data.Functor.Contravariant ((>$<))
import Data.Generics.Labels ()
import Data.Int (Int32)
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Text.Encoding qualified as TE
import Hasql.Decoders qualified as D
import Hasql.Encoders qualified as E
import Hasql.Statement (Statement, preparable)
import Hasql.Transaction qualified as Tx
import Keiro.Projection (InlineProjection (..))
import Keiro.ReadModel (ConsistencyMode (..), ReadModel (..), StrongScope (..))
import Kioku.Api.Scope (scopeFromColumns, scopeKindText, scopeNamespaceText, scopeRefText)
import Kioku.Api.Types (MemoryRecord (..), confidenceToText, memoryTypeToText)
import Kioku.Id (idText)
import Kioku.Memory.Domain
import Kioku.Prelude
import Kiroku.Store.Types (RecordedEvent)

data MemoryRow = MemoryRow
  { memoryId :: !Text,
    agentId :: !Text,
    sessionId :: !(Maybe Text),
    namespace :: !Text,
    scopeKind :: !(Maybe Text),
    scopeRef :: !(Maybe Text),
    memoryType :: !Text,
    content :: !Text,
    priority :: !Int,
    confidence :: !Text,
    tags :: !(Set Text),
    status :: !Text,
    supersededBy :: !(Maybe Text),
    supersedes :: !(Maybe Text),
    createdAt :: !UTCTime,
    updatedAt :: !UTCTime
  }
  deriving stock (Generic, Eq, Show)

newtype MemoryByIdQuery = MemoryByIdQuery Text

newtype MemoriesByNamespaceQuery = MemoriesByNamespaceQuery Text

data MemoriesByScopeQuery = MemoriesByScopeQuery Text (Maybe Text) (Maybe Text)

newtype MemoriesBySessionQuery = MemoriesBySessionQuery Text

data MemoriesByTypeQuery = MemoriesByTypeQuery Text Text

newtype MemorySupersessionChainQuery = MemorySupersessionChainQuery Text

memoryInlineProjection :: InlineProjection MemoryEvent
memoryInlineProjection =
  InlineProjection
    { name = "kioku-memory-inline",
      apply = applyMemoryEvent
    }

applyMemoryEvent :: MemoryEvent -> RecordedEvent -> Tx.Transaction ()
applyMemoryEvent event _recorded =
  case event of
    MemoryRecorded d -> Tx.statement (recordedRow d) upsertMemoryStmt
    MemorySuperseded d ->
      Tx.statement
        (idText d.memoryId, idText d.supersededBy, d.supersededAt)
        updateMemorySupersededStmt
    MemoryArchived d ->
      Tx.statement (idText d.memoryId, d.archivedAt) updateMemoryArchivedStmt
    MemoryTagsUpdated d ->
      Tx.statement (idText d.memoryId, d.tags, d.updatedAt) updateMemoryTagsStmt
    MemoryConfidenceUpdated d ->
      Tx.statement (idText d.memoryId, confidenceToText d.confidence, d.updatedAt) updateMemoryConfidenceStmt
    MemoryMerged d ->
      Tx.statement (idText d.memoryId, idText d.mergedInto, d.mergedAt) updateMemoryMergedStmt

recordedRow :: MemoryRecordedData -> MemoryRow
recordedRow d =
  MemoryRow
    { memoryId = idText d.memoryId,
      agentId = d.agentId,
      sessionId = idText <$> d.sessionId,
      namespace = scopeNamespaceText d.scope,
      scopeKind = scopeKindText d.scope,
      scopeRef = scopeRefText d.scope,
      memoryType = memoryTypeToText d.memoryType,
      content = d.content,
      priority = d.priority,
      confidence = confidenceToText d.confidence,
      tags = d.tags,
      status = "active",
      supersededBy = Nothing,
      supersedes = idText <$> d.supersedes,
      createdAt = d.recordedAt,
      updatedAt = d.recordedAt
    }

memoryByIdReadModel :: ReadModel MemoryByIdQuery (Maybe MemoryRow)
memoryByIdReadModel =
  ReadModel
    { name = "kioku-memory-by-id",
      schema = "kiroku",
      tableName = "kioku_memories",
      subscriptionName = "kioku-memory-inline",
      version = 1,
      shapeHash = "kioku-memory-v1",
      defaultConsistency = Eventual,
      strongScope = EntireLog,
      query = \(MemoryByIdQuery mid) -> Tx.statement mid selectMemoryByIdStmt
    }

memoriesByNamespaceReadModel :: ReadModel MemoriesByNamespaceQuery [MemoryRecord]
memoriesByNamespaceReadModel =
  ReadModel
    { name = "kioku-memories-by-namespace",
      schema = "kiroku",
      tableName = "kioku_memories",
      subscriptionName = "kioku-memory-inline",
      version = 1,
      shapeHash = "kioku-memory-v1",
      defaultConsistency = Eventual,
      strongScope = EntireLog,
      query = \(MemoriesByNamespaceQuery ns) -> Tx.statement ns selectActiveByNamespaceStmt
    }

memoriesByNamespaceRowsReadModel :: ReadModel MemoriesByNamespaceQuery [MemoryRow]
memoriesByNamespaceRowsReadModel =
  ReadModel
    { name = "kioku-memory-rows-by-namespace",
      schema = "kiroku",
      tableName = "kioku_memories",
      subscriptionName = "kioku-memory-inline",
      version = 1,
      shapeHash = "kioku-memory-v1",
      defaultConsistency = Eventual,
      strongScope = EntireLog,
      query = \(MemoriesByNamespaceQuery ns) -> Tx.statement ns selectActiveByNamespaceRowsStmt
    }

memoriesByScopeReadModel :: ReadModel MemoriesByScopeQuery [MemoryRecord]
memoriesByScopeReadModel =
  ReadModel
    { name = "kioku-memories-by-scope",
      schema = "kiroku",
      tableName = "kioku_memories",
      subscriptionName = "kioku-memory-inline",
      version = 1,
      shapeHash = "kioku-memory-v1",
      defaultConsistency = Eventual,
      strongScope = EntireLog,
      query = \(MemoriesByScopeQuery ns sk sr) -> Tx.statement (ns, sk, sr) selectActiveByScopeStmt
    }

memoriesByScopeRowsReadModel :: ReadModel MemoriesByScopeQuery [MemoryRow]
memoriesByScopeRowsReadModel =
  ReadModel
    { name = "kioku-memory-rows-by-scope",
      schema = "kiroku",
      tableName = "kioku_memories",
      subscriptionName = "kioku-memory-inline",
      version = 1,
      shapeHash = "kioku-memory-v1",
      defaultConsistency = Eventual,
      strongScope = EntireLog,
      query = \(MemoriesByScopeQuery ns sk sr) -> Tx.statement (ns, sk, sr) selectActiveByScopeRowsStmt
    }

memoriesBySessionReadModel :: ReadModel MemoriesBySessionQuery [MemoryRecord]
memoriesBySessionReadModel =
  ReadModel
    { name = "kioku-memories-by-session",
      schema = "kiroku",
      tableName = "kioku_memories",
      subscriptionName = "kioku-memory-inline",
      version = 1,
      shapeHash = "kioku-memory-v1",
      defaultConsistency = Eventual,
      strongScope = EntireLog,
      query = \(MemoriesBySessionQuery sid) -> Tx.statement sid selectBySessionStmt
    }

memoriesBySessionRowsReadModel :: ReadModel MemoriesBySessionQuery [MemoryRow]
memoriesBySessionRowsReadModel =
  ReadModel
    { name = "kioku-memory-rows-by-session",
      schema = "kiroku",
      tableName = "kioku_memories",
      subscriptionName = "kioku-memory-inline",
      version = 1,
      shapeHash = "kioku-memory-v1",
      defaultConsistency = Eventual,
      strongScope = EntireLog,
      query = \(MemoriesBySessionQuery sid) -> Tx.statement sid selectBySessionRowsStmt
    }

memoriesByTypeReadModel :: ReadModel MemoriesByTypeQuery [MemoryRecord]
memoriesByTypeReadModel =
  ReadModel
    { name = "kioku-memories-by-type",
      schema = "kiroku",
      tableName = "kioku_memories",
      subscriptionName = "kioku-memory-inline",
      version = 1,
      shapeHash = "kioku-memory-v1",
      defaultConsistency = Eventual,
      strongScope = EntireLog,
      query = \(MemoriesByTypeQuery ns mt) -> Tx.statement (ns, mt) selectByTypeStmt
    }

memoriesByTypeRowsReadModel :: ReadModel MemoriesByTypeQuery [MemoryRow]
memoriesByTypeRowsReadModel =
  ReadModel
    { name = "kioku-memory-rows-by-type",
      schema = "kiroku",
      tableName = "kioku_memories",
      subscriptionName = "kioku-memory-inline",
      version = 1,
      shapeHash = "kioku-memory-v1",
      defaultConsistency = Eventual,
      strongScope = EntireLog,
      query = \(MemoriesByTypeQuery ns mt) -> Tx.statement (ns, mt) selectByTypeRowsStmt
    }

memorySupersessionChainReadModel :: ReadModel MemorySupersessionChainQuery [MemoryRow]
memorySupersessionChainReadModel =
  ReadModel
    { name = "kioku-memory-supersession-chain",
      schema = "kiroku",
      tableName = "kioku_memories",
      subscriptionName = "kioku-memory-inline",
      version = 1,
      shapeHash = "kioku-memory-v1",
      defaultConsistency = Eventual,
      strongScope = EntireLog,
      query = \(MemorySupersessionChainQuery mid) -> Tx.statement mid selectSupersessionChainStmt
    }

memoryRowDecoder :: D.Row MemoryRow
memoryRowDecoder =
  MemoryRow
    <$> D.column (D.nonNullable D.text)
    <*> D.column (D.nonNullable D.text)
    <*> D.column (D.nullable D.text)
    <*> D.column (D.nonNullable D.text)
    <*> D.column (D.nullable D.text)
    <*> D.column (D.nullable D.text)
    <*> D.column (D.nonNullable D.text)
    <*> D.column (D.nonNullable D.text)
    <*> (fromIntegral @Int32 @Int <$> D.column (D.nonNullable D.int4))
    <*> D.column (D.nonNullable D.text)
    <*> (decodeTags <$> D.column (D.nonNullable D.text))
    <*> D.column (D.nonNullable D.text)
    <*> D.column (D.nullable D.text)
    <*> D.column (D.nullable D.text)
    <*> D.column (D.nonNullable D.timestamptz)
    <*> D.column (D.nonNullable D.timestamptz)

memoryRecordDecoder :: D.Row MemoryRecord
memoryRecordDecoder =
  toRecord <$> memoryRowDecoder

toRecord :: MemoryRow -> MemoryRecord
toRecord row =
  MemoryRecord
    { memoryId = row.memoryId,
      agentId = row.agentId,
      sessionId = row.sessionId,
      scope = scopeFromColumns row.namespace row.scopeKind row.scopeRef,
      memoryType = row.memoryType,
      content = row.content,
      priority = row.priority,
      confidence = row.confidence,
      tags = row.tags,
      status = row.status,
      createdAt = row.createdAt
    }

encodeTags :: Set Text -> Text
encodeTags = TE.decodeUtf8 . BL.toStrict . Aeson.encode

decodeTags :: Text -> Set Text
decodeTags =
  fromMaybe Set.empty . Aeson.decode . BL.fromStrict . TE.encodeUtf8

memoryRowColumns :: Text
memoryRowColumns =
  "memory_id, agent_id, session_id, namespace, scope_kind, scope_ref, memory_type, content, priority, confidence, tags::text, status, superseded_by, supersedes, created_at, updated_at"

qualifiedMemoryRowColumns :: Text -> Text
qualifiedMemoryRowColumns prefix =
  prefix
    <> ".memory_id, "
    <> prefix
    <> ".agent_id, "
    <> prefix
    <> ".session_id, "
    <> prefix
    <> ".namespace, "
    <> prefix
    <> ".scope_kind, "
    <> prefix
    <> ".scope_ref, "
    <> prefix
    <> ".memory_type, "
    <> prefix
    <> ".content, "
    <> prefix
    <> ".priority, "
    <> prefix
    <> ".confidence, "
    <> prefix
    <> ".tags::text, "
    <> prefix
    <> ".status, "
    <> prefix
    <> ".superseded_by, "
    <> prefix
    <> ".supersedes, "
    <> prefix
    <> ".created_at, "
    <> prefix
    <> ".updated_at"

selectMemoryByIdStmt :: Statement Text (Maybe MemoryRow)
selectMemoryByIdStmt =
  preparable
    ( "SELECT "
        <> memoryRowColumns
        <> " FROM kioku_memories WHERE memory_id = $1"
    )
    (E.param (E.nonNullable E.text))
    (D.rowMaybe memoryRowDecoder)

selectActiveByNamespaceStmt :: Statement Text [MemoryRecord]
selectActiveByNamespaceStmt =
  preparable
    ( "SELECT "
        <> memoryRowColumns
        <> " FROM kioku_memories WHERE status = 'active' AND namespace = $1 ORDER BY priority ASC, created_at DESC"
    )
    (E.param (E.nonNullable E.text))
    (D.rowList memoryRecordDecoder)

selectActiveByNamespaceRowsStmt :: Statement Text [MemoryRow]
selectActiveByNamespaceRowsStmt =
  preparable
    ( "SELECT "
        <> memoryRowColumns
        <> " FROM kioku_memories WHERE status = 'active' AND namespace = $1 ORDER BY priority ASC, created_at DESC"
    )
    (E.param (E.nonNullable E.text))
    (D.rowList memoryRowDecoder)

-- | Active memories carrying __exactly__ the given scope.
--
-- Recall searches namespace-wide for a global scope; scoped reads are exact-scope. The
-- predicate below /requires/ @scope_kind@ and @scope_ref@ to be NULL for a global scope,
-- where 'Kioku.Recall.selectFtsCandidatesStmt' would drop the scope filter entirely and
-- return the whole namespace. Both behaviours are intentional; see docs/user/recall.md.
selectActiveByScopeStmt :: Statement (Text, Maybe Text, Maybe Text) [MemoryRecord]
selectActiveByScopeStmt =
  preparable
    ( "SELECT "
        <> memoryRowColumns
        <> " FROM kioku_memories WHERE status = 'active' AND namespace = $1 AND ((scope_kind = $2 AND scope_ref = $3) OR ($2 IS NULL AND scope_kind IS NULL AND $3 IS NULL AND scope_ref IS NULL)) ORDER BY priority ASC, created_at DESC"
    )
    ( contrazip3
        (E.param (E.nonNullable E.text))
        (E.param (E.nullable E.text))
        (E.param (E.nullable E.text))
    )
    (D.rowList memoryRecordDecoder)

selectActiveByScopeRowsStmt :: Statement (Text, Maybe Text, Maybe Text) [MemoryRow]
selectActiveByScopeRowsStmt =
  preparable
    ( "SELECT "
        <> memoryRowColumns
        <> " FROM kioku_memories WHERE status = 'active' AND namespace = $1 AND ((scope_kind = $2 AND scope_ref = $3) OR ($2 IS NULL AND scope_kind IS NULL AND $3 IS NULL AND scope_ref IS NULL)) ORDER BY priority ASC, created_at DESC"
    )
    ( contrazip3
        (E.param (E.nonNullable E.text))
        (E.param (E.nullable E.text))
        (E.param (E.nullable E.text))
    )
    (D.rowList memoryRowDecoder)

selectBySessionStmt :: Statement Text [MemoryRecord]
selectBySessionStmt =
  preparable
    ( "SELECT "
        <> memoryRowColumns
        <> " FROM kioku_memories WHERE session_id = $1 ORDER BY created_at DESC"
    )
    (E.param (E.nonNullable E.text))
    (D.rowList memoryRecordDecoder)

selectBySessionRowsStmt :: Statement Text [MemoryRow]
selectBySessionRowsStmt =
  preparable
    ( "SELECT "
        <> memoryRowColumns
        <> " FROM kioku_memories WHERE session_id = $1 ORDER BY created_at DESC"
    )
    (E.param (E.nonNullable E.text))
    (D.rowList memoryRowDecoder)

selectByTypeStmt :: Statement (Text, Text) [MemoryRecord]
selectByTypeStmt =
  preparable
    ( "SELECT "
        <> memoryRowColumns
        <> " FROM kioku_memories WHERE status = 'active' AND namespace = $1 AND memory_type = $2 ORDER BY priority ASC, created_at DESC"
    )
    ( contrazip2
        (E.param (E.nonNullable E.text))
        (E.param (E.nonNullable E.text))
    )
    (D.rowList memoryRecordDecoder)

selectByTypeRowsStmt :: Statement (Text, Text) [MemoryRow]
selectByTypeRowsStmt =
  preparable
    ( "SELECT "
        <> memoryRowColumns
        <> " FROM kioku_memories WHERE status = 'active' AND namespace = $1 AND memory_type = $2 ORDER BY priority ASC, created_at DESC"
    )
    ( contrazip2
        (E.param (E.nonNullable E.text))
        (E.param (E.nonNullable E.text))
    )
    (D.rowList memoryRowDecoder)

selectSupersessionChainStmt :: Statement Text [MemoryRow]
selectSupersessionChainStmt =
  preparable
    ( "WITH RECURSIVE chain AS ("
        <> "SELECT "
        <> memoryRowColumns
        <> " FROM kioku_memories WHERE memory_id = $1 "
        <> "UNION "
        <> "SELECT "
        <> qualifiedMemoryRowColumns "m"
        <> " FROM kioku_memories m JOIN chain c ON "
        <> "m.memory_id = c.supersedes "
        <> "OR m.supersedes = c.memory_id "
        <> "OR m.memory_id = c.superseded_by "
        <> "OR m.superseded_by = c.memory_id"
        <> ") SELECT "
        <> memoryRowColumns
        <> " FROM chain ORDER BY created_at ASC, memory_id ASC"
    )
    (E.param (E.nonNullable E.text))
    (D.rowList memoryRowDecoder)

upsertMemoryStmt :: Statement MemoryRow ()
upsertMemoryStmt =
  preparable
    """
    INSERT INTO kioku_memories
      (memory_id, agent_id, session_id, namespace, scope_kind, scope_ref, memory_type, content,
       priority, confidence, tags, status, superseded_by, supersedes, created_at, updated_at)
    VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11::jsonb, $12, $13, $14, $15, $16)
    ON CONFLICT (memory_id) DO UPDATE SET
      agent_id = EXCLUDED.agent_id,
      session_id = EXCLUDED.session_id,
      namespace = EXCLUDED.namespace,
      scope_kind = EXCLUDED.scope_kind,
      scope_ref = EXCLUDED.scope_ref,
      memory_type = EXCLUDED.memory_type,
      content = EXCLUDED.content,
      priority = EXCLUDED.priority,
      confidence = EXCLUDED.confidence,
      tags = EXCLUDED.tags,
      status = EXCLUDED.status,
      superseded_by = EXCLUDED.superseded_by,
      supersedes = EXCLUDED.supersedes,
      updated_at = EXCLUDED.updated_at
    """
    memoryRowEncoder
    D.noResult

memoryRowEncoder :: E.Params MemoryRow
memoryRowEncoder =
  ((\row -> row.memoryId) >$< E.param (E.nonNullable E.text))
    <> ((\row -> row.agentId) >$< E.param (E.nonNullable E.text))
    <> ((\row -> row.sessionId) >$< E.param (E.nullable E.text))
    <> ((\row -> row.namespace) >$< E.param (E.nonNullable E.text))
    <> ((\row -> row.scopeKind) >$< E.param (E.nullable E.text))
    <> ((\row -> row.scopeRef) >$< E.param (E.nullable E.text))
    <> ((\row -> row.memoryType) >$< E.param (E.nonNullable E.text))
    <> ((\row -> row.content) >$< E.param (E.nonNullable E.text))
    <> ((fromIntegral @Int @Int32 . \row -> row.priority) >$< E.param (E.nonNullable E.int4))
    <> ((\row -> row.confidence) >$< E.param (E.nonNullable E.text))
    <> ((encodeTags . (\row -> row.tags)) >$< E.param (E.nonNullable E.text))
    <> ((\row -> row.status) >$< E.param (E.nonNullable E.text))
    <> ((\row -> row.supersededBy) >$< E.param (E.nullable E.text))
    <> ((\row -> row.supersedes) >$< E.param (E.nullable E.text))
    <> ((\row -> row.createdAt) >$< E.param (E.nonNullable E.timestamptz))
    <> ((\row -> row.updatedAt) >$< E.param (E.nonNullable E.timestamptz))

updateMemorySupersededStmt :: Statement (Text, Text, UTCTime) ()
updateMemorySupersededStmt =
  preparable
    "UPDATE kioku_memories SET status = 'superseded', superseded_by = $2, updated_at = $3 WHERE memory_id = $1"
    ( contrazip3
        (E.param (E.nonNullable E.text))
        (E.param (E.nonNullable E.text))
        (E.param (E.nonNullable E.timestamptz))
    )
    D.noResult

updateMemoryArchivedStmt :: Statement (Text, UTCTime) ()
updateMemoryArchivedStmt =
  preparable
    "UPDATE kioku_memories SET status = 'archived', updated_at = $2 WHERE memory_id = $1"
    (contrazip2 (E.param (E.nonNullable E.text)) (E.param (E.nonNullable E.timestamptz)))
    D.noResult

updateMemoryTagsStmt :: Statement (Text, Set Text, UTCTime) ()
updateMemoryTagsStmt =
  preparable
    "UPDATE kioku_memories SET tags = $2::jsonb, updated_at = $3 WHERE memory_id = $1"
    ( contrazip3
        (E.param (E.nonNullable E.text))
        (encodeTags >$< E.param (E.nonNullable E.text))
        (E.param (E.nonNullable E.timestamptz))
    )
    D.noResult

updateMemoryConfidenceStmt :: Statement (Text, Text, UTCTime) ()
updateMemoryConfidenceStmt =
  preparable
    "UPDATE kioku_memories SET confidence = $2, updated_at = $3 WHERE memory_id = $1"
    ( contrazip3
        (E.param (E.nonNullable E.text))
        (E.param (E.nonNullable E.text))
        (E.param (E.nonNullable E.timestamptz))
    )
    D.noResult

updateMemoryMergedStmt :: Statement (Text, Text, UTCTime) ()
updateMemoryMergedStmt =
  preparable
    "UPDATE kioku_memories SET status = 'merged', superseded_by = $2, updated_at = $3 WHERE memory_id = $1"
    ( contrazip3
        (E.param (E.nonNullable E.text))
        (E.param (E.nonNullable E.text))
        (E.param (E.nonNullable E.timestamptz))
    )
    D.noResult