path: root/Omni/Agent/Memory.hs

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE NoImplicitPrelude #-}

-- | Cross-agent shared memory system with vector similarity search.
--
-- Provides persistent memory that is:
-- - Shared across all agents (Telegram, researcher, coder, etc.)
-- - Private per user (users can't see each other's memories)
-- - Searchable via semantic similarity using embeddings
--
-- Uses sqlite-vss for vector similarity search and Ollama for embeddings.
--
-- : out omni-agent-memory
-- : dep aeson
-- : dep http-conduit
-- : dep sqlite-simple
-- : dep uuid
-- : dep vector
-- : dep directory
-- : dep bytestring
module Omni.Agent.Memory
  ( -- * Types
    User (..),
    Memory (..),
    MemorySource (..),
    ConversationMessage (..),
    ConversationSummary (..),
    MessageRole (..),
    RelationType (..),
    MemoryLink (..),

    -- * User Management
    createUser,
    getUser,
    getUserByTelegramId,
    getOrCreateUserByTelegramId,

    -- * Memory Operations
    storeMemory,
    recallMemories,
    forgetMemory,
    getAllMemoriesForUser,
    updateMemoryAccess,

    -- * Knowledge Graph
    linkMemories,
    getMemoryLinks,
    getLinkedMemories,
    queryGraph,

    -- * Conversation History (DMs)
    saveMessage,
    getRecentMessages,
    getConversationContext,
    summarizeAndArchive,
    estimateTokens,

    -- * Group Conversation History
    saveGroupMessage,
    getGroupRecentMessages,
    getGroupConversationContext,

    -- * Group Memories
    storeGroupMemory,
    recallGroupMemories,

    -- * Embeddings
    embedText,

    -- * Agent Integration
    rememberTool,
    recallTool,
    linkMemoriesTool,
    queryGraphTool,
    formatMemoriesForPrompt,
    runAgentWithMemory,

    -- * Database
    withMemoryDb,
    initMemoryDb,
    getMemoryDbPath,

    -- * Testing
    main,
    test,
  )
where

import Alpha
import Data.Aeson ((.!=), (.:), (.:?), (.=))
import qualified Data.Aeson as Aeson
import qualified Data.Aeson.KeyMap as KeyMap
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BL
import qualified Data.List as List
import qualified Data.Text as Text
import qualified Data.Text.Encoding as TE
import Data.Time (UTCTime, getCurrentTime)
import Data.Time.Format (defaultTimeLocale, formatTime)
import qualified Data.UUID as UUID
import qualified Data.UUID.V4 as UUID
import qualified Data.Vector.Storable as VS
import qualified Database.SQLite.Simple as SQL
import Database.SQLite.Simple.FromField ()
import qualified Database.SQLite.Simple.ToField as SQL
import Foreign.Storable ()
import qualified Network.HTTP.Simple as HTTP
import qualified Omni.Agent.Engine as Engine
import qualified Omni.Test as Test
import System.Directory (createDirectoryIfMissing)
import System.Environment (lookupEnv)
import System.FilePath (takeDirectory, (</>))

main :: IO ()
main = Test.run test

test :: Test.Tree
test =
  Test.group
    "Omni.Agent.Memory"
    [ Test.unit "User JSON roundtrip" <| do
        now <- getCurrentTime
        let user =
              User
                { userId = "test-uuid",
                  userTelegramId = Just 12345,
                  userEmail = Nothing,
                  userName = "Test User",
                  userCreatedAt = now
                }
        case Aeson.decode (Aeson.encode user) of
          Nothing -> Test.assertFailure "Failed to decode User"
          Just decoded -> userName decoded Test.@=? "Test User",
      Test.unit "Memory JSON roundtrip" <| do
        now <- getCurrentTime
        let mem =
              Memory
                { memoryId = "mem-uuid",
                  memoryUserId = "user-uuid",
                  memoryContent = "User is an AI engineer",
                  memoryEmbedding = Nothing,
                  memorySource =
                    MemorySource
                      { sourceAgent = "telegram",
                        sourceSession = Nothing,
                        sourceContext = "User mentioned in chat"
                      },
                  memoryConfidence = 0.9,
                  memoryCreatedAt = now,
                  memoryLastAccessedAt = now,
                  memoryTags = ["profession", "ai"]
                }
        case Aeson.decode (Aeson.encode mem) of
          Nothing -> Test.assertFailure "Failed to decode Memory"
          Just decoded -> memoryContent decoded Test.@=? "User is an AI engineer",
      Test.unit "MemorySource JSON roundtrip" <| do
        let src =
              MemorySource
                { sourceAgent = "researcher",
                  sourceSession = Just "session-123",
                  sourceContext = "Extracted from conversation"
                }
        case Aeson.decode (Aeson.encode src) of
          Nothing -> Test.assertFailure "Failed to decode MemorySource"
          Just decoded -> sourceAgent decoded Test.@=? "researcher",
      Test.unit "formatMemoriesForPrompt formats correctly" <| do
        now <- getCurrentTime
        let mem1 =
              Memory
                { memoryId = "1",
                  memoryUserId = "u",
                  memoryContent = "User is an AI engineer",
                  memoryEmbedding = Nothing,
                  memorySource = MemorySource "telegram" Nothing "chat",
                  memoryConfidence = 0.9,
                  memoryCreatedAt = now,
                  memoryLastAccessedAt = now,
                  memoryTags = []
                }
            mem2 =
              Memory
                { memoryId = "2",
                  memoryUserId = "u",
                  memoryContent = "User prefers Haskell",
                  memoryEmbedding = Nothing,
                  memorySource = MemorySource "coder" Nothing "code review",
                  memoryConfidence = 0.8,
                  memoryCreatedAt = now,
                  memoryLastAccessedAt = now,
                  memoryTags = []
                }
            formatted = formatMemoriesForPrompt [mem1, mem2]
        ("AI engineer" `Text.isInfixOf` formatted) Test.@=? True
        ("Haskell" `Text.isInfixOf` formatted) Test.@=? True,
      Test.unit "cosineSimilarity identical vectors" <| do
        let v1 = VS.fromList [1.0, 0.0, 0.0 :: Float]
            v2 = VS.fromList [1.0, 0.0, 0.0 :: Float]
        abs (cosineSimilarity v1 v2 - 1.0) < 0.0001 Test.@=? True,
      Test.unit "cosineSimilarity orthogonal vectors" <| do
        let v1 = VS.fromList [1.0, 0.0, 0.0 :: Float]
            v2 = VS.fromList [0.0, 1.0, 0.0 :: Float]
        abs (cosineSimilarity v1 v2) < 0.0001 Test.@=? True,
      Test.unit "cosineSimilarity opposite vectors" <| do
        let v1 = VS.fromList [1.0, 0.0, 0.0 :: Float]
            v2 = VS.fromList [-1.0, 0.0, 0.0 :: Float]
        abs (cosineSimilarity v1 v2 + 1.0) < 0.0001 Test.@=? True,
      Test.unit "vectorToBlob and blobToVector roundtrip" <| do
        let v = VS.fromList [0.1, 0.2, 0.3, 0.4, 0.5 :: Float]
            blob = vectorToBlob v
            v' = blobToVector blob
        VS.length v Test.@=? VS.length v'
        VS.toList v Test.@=? VS.toList v',
      Test.unit "rememberTool has correct schema" <| do
        let tool = rememberTool "test-user-id"
        Engine.toolName tool Test.@=? "remember",
      Test.unit "recallTool has correct schema" <| do
        let tool = recallTool "test-user-id"
        Engine.toolName tool Test.@=? "recall",
      Test.unit "RelationType JSON roundtrip" <| do
        let types = [Contradicts, Supports, Elaborates, Supersedes, Related, ContingentOn]
        forM_ types <| \rt ->
          case Aeson.decode (Aeson.encode rt) of
            Nothing -> Test.assertFailure ("Failed to decode RelationType: " <> show rt)
            Just decoded -> decoded Test.@=? rt,
      Test.unit "MemoryLink JSON roundtrip" <| do
        now <- getCurrentTime
        let memLink =
              MemoryLink
                { linkFromMemoryId = "mem-1",
                  linkToMemoryId = "mem-2",
                  linkRelationType = Contradicts,
                  linkCreatedAt = now
                }
        case Aeson.decode (Aeson.encode memLink) of
          Nothing -> Test.assertFailure "Failed to decode MemoryLink"
          Just decoded -> do
            linkFromMemoryId decoded Test.@=? "mem-1"
            linkToMemoryId decoded Test.@=? "mem-2"
            linkRelationType decoded Test.@=? Contradicts,
      Test.unit "relationTypeToText and textToRelationType roundtrip" <| do
        let types = [Contradicts, Supports, Elaborates, Supersedes, Related, ContingentOn]
        forM_ types <| \rt ->
          textToRelationType (relationTypeToText rt) Test.@=? Just rt,
      Test.unit "linkMemoriesTool has correct schema" <| do
        let tool = linkMemoriesTool "test-user-id"
        Engine.toolName tool Test.@=? "link_memories",
      Test.unit "queryGraphTool has correct schema" <| do
        let tool = queryGraphTool "test-user-id"
        Engine.toolName tool Test.@=? "query_graph"
    ]

-- | User record for multi-user memory system.
data User = User
  { userId :: Text,
    userTelegramId :: Maybe Int,
    userEmail :: Maybe Text,
    userName :: Text,
    userCreatedAt :: UTCTime
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON User where
  toJSON u =
    Aeson.object
      [ "id" .= userId u,
        "telegram_id" .= userTelegramId u,
        "email" .= userEmail u,
        "name" .= userName u,
        "created_at" .= userCreatedAt u
      ]

instance Aeson.FromJSON User where
  parseJSON =
    Aeson.withObject "User" <| \v ->
      (User </ (v .: "id"))
        <*> (v .:? "telegram_id")
        <*> (v .:? "email")
        <*> (v .: "name")
        <*> (v .: "created_at")

instance SQL.FromRow User where
  fromRow =
    User
      </ SQL.field
      <*> SQL.field
      <*> SQL.field
      <*> SQL.field
      <*> SQL.field

instance SQL.ToRow User where
  toRow u =
    [ SQL.toField (userId u),
      SQL.toField (userTelegramId u),
      SQL.toField (userEmail u),
      SQL.toField (userName u),
      SQL.toField (userCreatedAt u)
    ]

-- | Source information for a memory.
data MemorySource = MemorySource
  { sourceAgent :: Text,
    sourceSession :: Maybe Text,
    sourceContext :: Text
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON MemorySource where
  toJSON s =
    Aeson.object
      [ "agent" .= sourceAgent s,
        "session" .= sourceSession s,
        "context" .= sourceContext s
      ]

instance Aeson.FromJSON MemorySource where
  parseJSON =
    Aeson.withObject "MemorySource" <| \v ->
      (MemorySource </ (v .: "agent"))
        <*> (v .:? "session")
        <*> (v .: "context")

-- | A memory stored in the system.
data Memory = Memory
  { memoryId :: Text,
    memoryUserId :: Text,
    memoryContent :: Text,
    memoryEmbedding :: Maybe (VS.Vector Float),
    memorySource :: MemorySource,
    memoryConfidence :: Double,
    memoryCreatedAt :: UTCTime,
    memoryLastAccessedAt :: UTCTime,
    memoryTags :: [Text]
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON Memory where
  toJSON m =
    Aeson.object
      [ "id" .= memoryId m,
        "user_id" .= memoryUserId m,
        "content" .= memoryContent m,
        "source" .= memorySource m,
        "confidence" .= memoryConfidence m,
        "created_at" .= memoryCreatedAt m,
        "last_accessed_at" .= memoryLastAccessedAt m,
        "tags" .= memoryTags m
      ]

instance Aeson.FromJSON Memory where
  parseJSON =
    Aeson.withObject "Memory" <| \v ->
      ( Memory
          </ (v .: "id")
      )
        <*> (v .: "user_id")
        <*> (v .: "content")
        <*> pure Nothing
        <*> (v .: "source")
        <*> (v .:? "confidence" .!= 0.8)
        <*> (v .: "created_at")
        <*> (v .: "last_accessed_at")
        <*> (v .:? "tags" .!= [])

-- SQLite instances for Memory (partial - embedding handled separately)
instance SQL.FromRow Memory where
  fromRow = do
    mid <- SQL.field
    uid <- SQL.field
    content <- SQL.field
    embeddingBlob <- SQL.field
    agent <- SQL.field
    session <- SQL.field
    context <- SQL.field
    confidence <- SQL.field
    createdAt <- SQL.field
    lastAccessedAt <- SQL.field
    tagsJson <- SQL.field
    let embedding = blobToVector </ (embeddingBlob :: Maybe BS.ByteString)
        source = MemorySource agent session context
        tags = fromMaybe [] ((tagsJson :: Maybe Text) +> (Aeson.decode <. BL.fromStrict <. TE.encodeUtf8))
    pure
      Memory
        { memoryId = mid,
          memoryUserId = uid,
          memoryContent = content,
          memoryEmbedding = embedding,
          memorySource = source,
          memoryConfidence = confidence,
          memoryCreatedAt = createdAt,
          memoryLastAccessedAt = lastAccessedAt,
          memoryTags = tags
        }

-- | Role in a conversation message.
data MessageRole = UserRole | AssistantRole
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON MessageRole where
  toJSON UserRole = Aeson.String "user"
  toJSON AssistantRole = Aeson.String "assistant"

instance Aeson.FromJSON MessageRole where
  parseJSON =
    Aeson.withText "MessageRole" <| \case
      "user" -> pure UserRole
      "assistant" -> pure AssistantRole
      _ -> empty

-- | A message in a conversation.
data ConversationMessage = ConversationMessage
  { cmId :: Maybe Int,
    cmUserId :: Text,
    cmChatId :: Int,
    cmRole :: MessageRole,
    cmSenderName :: Maybe Text,
    cmContent :: Text,
    cmTokensEstimate :: Int,
    cmCreatedAt :: UTCTime
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON ConversationMessage where
  toJSON m =
    Aeson.object
      [ "id" .= cmId m,
        "user_id" .= cmUserId m,
        "chat_id" .= cmChatId m,
        "role" .= cmRole m,
        "sender_name" .= cmSenderName m,
        "content" .= cmContent m,
        "tokens_estimate" .= cmTokensEstimate m,
        "created_at" .= cmCreatedAt m
      ]

instance SQL.FromRow ConversationMessage where
  fromRow =
    (ConversationMessage </ SQL.field)
      <*> SQL.field
      <*> SQL.field
      <*> (parseRole </ SQL.field)
      <*> SQL.field
      <*> SQL.field
      <*> (fromMaybe 0 </ SQL.field)
      <*> SQL.field
    where
      parseRole :: Text -> MessageRole
      parseRole "user" = UserRole
      parseRole _ = AssistantRole

-- | A summary of older conversation messages.
data ConversationSummary = ConversationSummary
  { csId :: Maybe Int,
    csUserId :: Text,
    csChatId :: Int,
    csSummary :: Text,
    csMessagesSummarized :: Int,
    csTokensSaved :: Maybe Int,
    csCreatedAt :: UTCTime
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON ConversationSummary where
  toJSON s =
    Aeson.object
      [ "id" .= csId s,
        "user_id" .= csUserId s,
        "chat_id" .= csChatId s,
        "summary" .= csSummary s,
        "messages_summarized" .= csMessagesSummarized s,
        "tokens_saved" .= csTokensSaved s,
        "created_at" .= csCreatedAt s
      ]

instance SQL.FromRow ConversationSummary where
  fromRow =
    (ConversationSummary </ SQL.field)
      <*> SQL.field
      <*> SQL.field
      <*> SQL.field
      <*> SQL.field
      <*> SQL.field
      <*> SQL.field

-- | Relation types for the knowledge graph.
data RelationType
  = Contradicts
  | Supports
  | Elaborates
  | Supersedes
  | Related
  | ContingentOn
  deriving (Show, Eq, Generic, Ord)

instance Aeson.ToJSON RelationType where
  toJSON Contradicts = Aeson.String "contradicts"
  toJSON Supports = Aeson.String "supports"
  toJSON Elaborates = Aeson.String "elaborates"
  toJSON Supersedes = Aeson.String "supersedes"
  toJSON Related = Aeson.String "related"
  toJSON ContingentOn = Aeson.String "contingent_on"

instance Aeson.FromJSON RelationType where
  parseJSON =
    Aeson.withText "RelationType" <| \case
      "contradicts" -> pure Contradicts
      "supports" -> pure Supports
      "elaborates" -> pure Elaborates
      "supersedes" -> pure Supersedes
      "related" -> pure Related
      "contingent_on" -> pure ContingentOn
      _ -> empty

relationTypeToText :: RelationType -> Text
relationTypeToText Contradicts = "contradicts"
relationTypeToText Supports = "supports"
relationTypeToText Elaborates = "elaborates"
relationTypeToText Supersedes = "supersedes"
relationTypeToText Related = "related"
relationTypeToText ContingentOn = "contingent_on"

textToRelationType :: Text -> Maybe RelationType
textToRelationType "contradicts" = Just Contradicts
textToRelationType "supports" = Just Supports
textToRelationType "elaborates" = Just Elaborates
textToRelationType "supersedes" = Just Supersedes
textToRelationType "related" = Just Related
textToRelationType "contingent_on" = Just ContingentOn
textToRelationType _ = Nothing

-- | A link between two memories in the knowledge graph.
data MemoryLink = MemoryLink
  { linkFromMemoryId :: Text,
    linkToMemoryId :: Text,
    linkRelationType :: RelationType,
    linkCreatedAt :: UTCTime
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON MemoryLink where
  toJSON l =
    Aeson.object
      [ "from_memory_id" .= linkFromMemoryId l,
        "to_memory_id" .= linkToMemoryId l,
        "relation_type" .= linkRelationType l,
        "created_at" .= linkCreatedAt l
      ]

instance Aeson.FromJSON MemoryLink where
  parseJSON =
    Aeson.withObject "MemoryLink" <| \v ->
      (MemoryLink </ (v .: "from_memory_id"))
        <*> (v .: "to_memory_id")
        <*> (v .: "relation_type")
        <*> (v .: "created_at")

instance SQL.FromRow MemoryLink where
  fromRow = do
    fromId <- SQL.field
    toId <- SQL.field
    relTypeText <- SQL.field
    createdAt <- SQL.field
    let relType = fromMaybe Related (textToRelationType relTypeText)
    pure
      MemoryLink
        { linkFromMemoryId = fromId,
          linkToMemoryId = toId,
          linkRelationType = relType,
          linkCreatedAt = createdAt
        }

-- | Get the path to memory.db
getMemoryDbPath :: IO FilePath
getMemoryDbPath = do
  maybeEnv <- lookupEnv "MEMORY_DB_PATH"
  case maybeEnv of
    Just p -> pure p
    Nothing -> do
      home <- lookupEnv "HOME"
      case home of
        Just h -> pure (h </> ".local/share/omni/memory.db")
        Nothing -> pure "_/memory.db"

-- | Run an action with the memory database connection.
withMemoryDb :: (SQL.Connection -> IO a) -> IO a
withMemoryDb action = do
  dbPath <- getMemoryDbPath
  createDirectoryIfMissing True (takeDirectory dbPath)
  SQL.withConnection dbPath <| \conn -> do
    initMemoryDb conn
    action conn

-- | Initialize the memory database schema.
initMemoryDb :: SQL.Connection -> IO ()
initMemoryDb conn = do
  SQL.execute_ conn "PRAGMA busy_timeout = 10000"
  SQL.execute_ conn "PRAGMA foreign_keys = ON"
  _ <- SQL.query_ conn "PRAGMA journal_mode = WAL" :: IO [[Text]]
  SQL.execute_
    conn
    "CREATE TABLE IF NOT EXISTS users (\
    \  id TEXT PRIMARY KEY,\
    \  telegram_id INTEGER UNIQUE,\
    \  email TEXT UNIQUE,\
    \  name TEXT NOT NULL,\
    \  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP\
    \)"
  SQL.execute_
    conn
    "CREATE TABLE IF NOT EXISTS memories (\
    \  id TEXT PRIMARY KEY,\
    \  user_id TEXT NOT NULL REFERENCES users(id),\
    \  content TEXT NOT NULL,\
    \  embedding BLOB,\
    \  source_agent TEXT NOT NULL,\
    \  source_session TEXT,\
    \  source_context TEXT,\
    \  confidence REAL DEFAULT 0.8,\
    \  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,\
    \  last_accessed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,\
    \  tags TEXT\
    \)"
  SQL.execute_
    conn
    "CREATE INDEX IF NOT EXISTS idx_memories_user ON memories(user_id)"
  SQL.execute_
    conn
    "CREATE INDEX IF NOT EXISTS idx_memories_agent ON memories(source_agent)"
  SQL.execute_
    conn
    "CREATE TABLE IF NOT EXISTS conversation_messages (\
    \  id INTEGER PRIMARY KEY AUTOINCREMENT,\
    \  user_id TEXT NOT NULL REFERENCES users(id),\
    \  chat_id INTEGER NOT NULL,\
    \  role TEXT NOT NULL,\
    \  sender_name TEXT,\
    \  content TEXT NOT NULL,\
    \  tokens_estimate INTEGER,\
    \  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP\
    \)"
  SQL.execute_
    conn
    "CREATE INDEX IF NOT EXISTS idx_conv_user_chat ON conversation_messages(user_id, chat_id)"
  migrateConversationMessages conn
  SQL.execute_
    conn
    "CREATE TABLE IF NOT EXISTS conversation_summaries (\
    \  id INTEGER PRIMARY KEY AUTOINCREMENT,\
    \  user_id TEXT NOT NULL REFERENCES users(id),\
    \  chat_id INTEGER NOT NULL,\
    \  summary TEXT NOT NULL,\
    \  messages_summarized INTEGER NOT NULL,\
    \  tokens_saved INTEGER,\
    \  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP\
    \)"
  SQL.execute_
    conn
    "CREATE INDEX IF NOT EXISTS idx_summary_user_chat ON conversation_summaries(user_id, chat_id)"
  SQL.execute_
    conn
    "CREATE TABLE IF NOT EXISTS notes (\
    \  id INTEGER PRIMARY KEY AUTOINCREMENT,\
    \  user_id TEXT NOT NULL,\
    \  topic TEXT NOT NULL,\
    \  content TEXT NOT NULL,\
    \  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP\
    \)"
  SQL.execute_
    conn
    "CREATE INDEX IF NOT EXISTS idx_notes_user ON notes(user_id)"
  SQL.execute_
    conn
    "CREATE INDEX IF NOT EXISTS idx_notes_topic ON notes(user_id, topic)"
  SQL.execute_
    conn
    "CREATE TABLE IF NOT EXISTS todos (\
    \  id INTEGER PRIMARY KEY AUTOINCREMENT,\
    \  user_id TEXT NOT NULL,\
    \  title TEXT NOT NULL,\
    \  due_date TIMESTAMP,\
    \  completed INTEGER NOT NULL DEFAULT 0,\
    \  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP\
    \)"
  SQL.execute_
    conn
    "CREATE INDEX IF NOT EXISTS idx_todos_user ON todos(user_id)"
  SQL.execute_
    conn
    "CREATE INDEX IF NOT EXISTS idx_todos_due ON todos(user_id, due_date)"
  SQL.execute_
    conn
    "CREATE TABLE IF NOT EXISTS memory_links (\
    \  from_memory_id TEXT NOT NULL REFERENCES memories(id) ON DELETE CASCADE,\
    \  to_memory_id TEXT NOT NULL REFERENCES memories(id) ON DELETE CASCADE,\
    \  relation_type TEXT NOT NULL,\
    \  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,\
    \  PRIMARY KEY (from_memory_id, to_memory_id, relation_type)\
    \)"
  SQL.execute_
    conn
    "CREATE INDEX IF NOT EXISTS idx_memory_links_from ON memory_links(from_memory_id)"
  SQL.execute_
    conn
    "CREATE INDEX IF NOT EXISTS idx_memory_links_to ON memory_links(to_memory_id)"
  SQL.execute_
    conn
    "CREATE INDEX IF NOT EXISTS idx_memory_links_type ON memory_links(relation_type)"

-- | Migrate conversation_messages to add sender_name and thread_id columns.
migrateConversationMessages :: SQL.Connection -> IO ()
migrateConversationMessages conn = do
  columns <- SQL.query_ conn "PRAGMA table_info(conversation_messages)" :: IO [(Int, Text, Text, Int, Maybe Text, Int)]
  let columnNames = map (\(_, name, _, _, _, _) -> name) columns
  unless ("sender_name" `elem` columnNames) <| do
    SQL.execute_ conn "ALTER TABLE conversation_messages ADD COLUMN sender_name TEXT"
    SQL.execute_ conn "UPDATE conversation_messages SET sender_name = 'bensima' WHERE role = 'user' AND sender_name IS NULL"
  unless ("thread_id" `elem` columnNames) <| do
    SQL.execute_ conn "ALTER TABLE conversation_messages ADD COLUMN thread_id INTEGER"
    SQL.execute_ conn "CREATE INDEX IF NOT EXISTS idx_conv_chat_thread ON conversation_messages(chat_id, thread_id)"

-- | Create a new user.
createUser :: Text -> Maybe Int -> IO User
createUser name telegramId = do
  uuid <- UUID.nextRandom
  now <- getCurrentTime
  let user =
        User
          { userId = UUID.toText uuid,
            userTelegramId = telegramId,
            userEmail = Nothing,
            userName = name,
            userCreatedAt = now
          }
  withMemoryDb <| \conn ->
    SQL.execute
      conn
      "INSERT INTO users (id, telegram_id, email, name, created_at) VALUES (?, ?, ?, ?, ?)"
      user
  pure user

-- | Get a user by ID.
getUser :: Text -> IO (Maybe User)
getUser uid =
  withMemoryDb <| \conn -> do
    results <- SQL.query conn "SELECT id, telegram_id, email, name, created_at FROM users WHERE id = ?" (SQL.Only uid)
    pure (listToMaybe results)

-- | Get a user by Telegram ID.
getUserByTelegramId :: Int -> IO (Maybe User)
getUserByTelegramId tid =
  withMemoryDb <| \conn -> do
    results <- SQL.query conn "SELECT id, telegram_id, email, name, created_at FROM users WHERE telegram_id = ?" (SQL.Only tid)
    pure (listToMaybe results)

-- | Get or create a user by Telegram ID.
getOrCreateUserByTelegramId :: Int -> Text -> IO User
getOrCreateUserByTelegramId tid name = do
  existing <- getUserByTelegramId tid
  case existing of
    Just user -> pure user
    Nothing -> createUser name (Just tid)

-- | Store a memory for a user.
storeMemory :: Text -> Text -> MemorySource -> IO Memory
storeMemory uid content source = storeMemoryWithTags uid content source []

-- | Store a memory with tags.
storeMemoryWithTags :: Text -> Text -> MemorySource -> [Text] -> IO Memory
storeMemoryWithTags uid content source tags = do
  uuid <- UUID.nextRandom
  now <- getCurrentTime
  embedding <- embedText content
  let mem =
        Memory
          { memoryId = UUID.toText uuid,
            memoryUserId = uid,
            memoryContent = content,
            memoryEmbedding = either (const Nothing) Just embedding,
            memorySource = source,
            memoryConfidence = 0.8,
            memoryCreatedAt = now,
            memoryLastAccessedAt = now,
            memoryTags = tags
          }
  withMemoryDb <| \conn ->
    SQL.execute
      conn
      "INSERT INTO memories (id, user_id, content, embedding, source_agent, source_session, source_context, confidence, created_at, last_accessed_at, tags) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"
      ( ( memoryId mem,
          memoryUserId mem,
          memoryContent mem,
          vectorToBlob </ memoryEmbedding mem,
          sourceAgent (memorySource mem),
          sourceSession (memorySource mem),
          sourceContext (memorySource mem)
        )
          SQL.:. ( memoryConfidence mem,
                   memoryCreatedAt mem,
                   memoryLastAccessedAt mem,
                   TE.decodeUtf8 (BL.toStrict (Aeson.encode (memoryTags mem)))
                 )
      )
  pure mem

-- | Recall memories for a user using semantic similarity.
recallMemories :: Text -> Text -> Int -> IO [Memory]
recallMemories uid query limit = do
  queryEmbedding <- embedText query
  case queryEmbedding of
    Left _ -> recallMemoriesByRecency uid limit
    Right qEmb -> do
      allMems <- getAllMemoriesForUser uid
      let scored =
            [ (m, cosineSimilarity qEmb emb)
              | m <- allMems,
                Just emb <- [memoryEmbedding m]
            ]
          sorted = List.sortBy (\(_, s1) (_, s2) -> compare s2 s1) scored
          topN = take limit sorted
      now <- getCurrentTime
      traverse_ (updateMemoryAccess now <. memoryId <. fst) topN
      pure (map fst topN)

-- | Recall memories by recency (fallback when embedding fails).
recallMemoriesByRecency :: Text -> Int -> IO [Memory]
recallMemoriesByRecency uid limit =
  withMemoryDb <| \conn -> do
    SQL.query
      conn
      "SELECT id, user_id, content, embedding, source_agent, source_session, source_context, confidence, created_at, last_accessed_at, tags \
      \FROM memories WHERE user_id = ? ORDER BY last_accessed_at DESC LIMIT ?"
      (uid, limit)

-- | Get all memories for a user.
getAllMemoriesForUser :: Text -> IO [Memory]
getAllMemoriesForUser uid =
  withMemoryDb <| \conn ->
    SQL.query
      conn
      "SELECT id, user_id, content, embedding, source_agent, source_session, source_context, confidence, created_at, last_accessed_at, tags \
      \FROM memories WHERE user_id = ?"
      (SQL.Only uid)

-- | Delete a memory.
forgetMemory :: Text -> IO ()
forgetMemory mid =
  withMemoryDb <| \conn ->
    SQL.execute conn "DELETE FROM memories WHERE id = ?" (SQL.Only mid)

-- | Update memory's last accessed timestamp.
updateMemoryAccess :: UTCTime -> Text -> IO ()
updateMemoryAccess now mid =
  withMemoryDb <| \conn ->
    SQL.execute conn "UPDATE memories SET last_accessed_at = ? WHERE id = ?" (now, mid)

-- | Create a link between two memories.
linkMemories :: Text -> Text -> RelationType -> IO MemoryLink
linkMemories fromId toId relType = do
  now <- getCurrentTime
  withMemoryDb <| \conn ->
    SQL.execute
      conn
      "INSERT OR REPLACE INTO memory_links (from_memory_id, to_memory_id, relation_type, created_at) VALUES (?, ?, ?, ?)"
      (fromId, toId, relationTypeToText relType, now)
  pure
    MemoryLink
      { linkFromMemoryId = fromId,
        linkToMemoryId = toId,
        linkRelationType = relType,
        linkCreatedAt = now
      }

-- | Get all links from a memory.
getMemoryLinks :: Text -> IO [MemoryLink]
getMemoryLinks memId =
  withMemoryDb <| \conn ->
    SQL.query
      conn
      "SELECT from_memory_id, to_memory_id, relation_type, created_at \
      \FROM memory_links WHERE from_memory_id = ? OR to_memory_id = ?"
      (memId, memId)

-- | Get memories linked to a given memory with their content.
getLinkedMemories :: Text -> Maybe RelationType -> IO [(MemoryLink, Memory)]
getLinkedMemories memId maybeRelType = do
  links <- getMemoryLinks memId
  let filteredLinks = case maybeRelType of
        Nothing -> links
        Just rt -> filter (\l -> linkRelationType l == rt) links
  mems <- traverse loadMemory filteredLinks
  pure [(l, m) | (l, Just m) <- zip filteredLinks mems]
  where
    loadMemory memLink = do
      let targetId =
            if linkFromMemoryId memLink == memId
              then linkToMemoryId memLink
              else linkFromMemoryId memLink
      withMemoryDb <| \conn -> do
        results <-
          SQL.query
            conn
            "SELECT id, user_id, content, embedding, source_agent, source_session, source_context, confidence, created_at, last_accessed_at, tags \
            \FROM memories WHERE id = ?"
            (SQL.Only targetId)
        pure (listToMaybe results)

-- | Query the knowledge graph by traversing links from a starting memory.
-- Returns all memories reachable within the given depth.
queryGraph :: Text -> Int -> Maybe RelationType -> IO [(Memory, [MemoryLink])]
queryGraph startMemId maxDepth maybeRelType = do
  startMem <- getMemoryById startMemId
  case startMem of
    Nothing -> pure []
    Just mem -> go [startMemId] [(mem, [])] 0
  where
    go :: [Text] -> [(Memory, [MemoryLink])] -> Int -> IO [(Memory, [MemoryLink])]
    go _ acc depth | depth >= maxDepth = pure acc
    go visitedIds acc depth = do
      let currentIds = map (memoryId <. fst) acc
          newIds = filter (`notElem` visitedIds) currentIds
      if null newIds
        then pure acc
        else do
          newLinked <- concat </ traverse (`getLinkedMemories` maybeRelType) newIds
          let newMems = [(m, [l]) | (l, m) <- newLinked, memoryId m `notElem` visitedIds]
              newVisited = visitedIds <> map (memoryId <. fst) newMems
          go newVisited (acc <> newMems) (depth + 1)

-- | Get a memory by ID.
getMemoryById :: Text -> IO (Maybe Memory)
getMemoryById memId =
  withMemoryDb <| \conn -> do
    results <-
      SQL.query
        conn
        "SELECT id, user_id, content, embedding, source_agent, source_session, source_context, confidence, created_at, last_accessed_at, tags \
        \FROM memories WHERE id = ?"
        (SQL.Only memId)
    pure (listToMaybe results)

-- | Embed text using Ollama's nomic-embed-text model.
embedText :: Text -> IO (Either Text (VS.Vector Float))
embedText content = do
  ollamaUrl <- fromMaybe "http://localhost:11434" </ lookupEnv "OLLAMA_URL"
  let url = ollamaUrl <> "/api/embeddings"
  req0 <- HTTP.parseRequest url
  let body =
        Aeson.object
          [ "model" .= ("nomic-embed-text" :: Text),
            "prompt" .= content
          ]
      req =
        HTTP.setRequestMethod "POST"
          <| HTTP.setRequestHeader "Content-Type" ["application/json"]
          <| HTTP.setRequestBodyLBS (Aeson.encode body)
          <| req0
  result <- try (HTTP.httpLBS req)
  case result of
    Left (e :: SomeException) ->
      pure (Left ("Embedding request failed: " <> tshow e))
    Right response -> do
      let status = HTTP.getResponseStatusCode response
      if status >= 200 && status < 300
        then case Aeson.decode (HTTP.getResponseBody response) of
          Just (Aeson.Object obj) -> case KeyMap.lookup "embedding" obj of
            Just (Aeson.Array arr) ->
              let floats = [f | Aeson.Number n <- toList arr, let f = realToFrac n]
               in pure (Right (VS.fromList floats))
            _ -> pure (Left "No embedding in response")
          _ -> pure (Left "Failed to parse embedding response")
        else pure (Left ("Embedding HTTP error: " <> tshow status))

-- | Convert a vector to a blob for storage.
vectorToBlob :: VS.Vector Float -> BS.ByteString
vectorToBlob v =
  let bytes = VS.unsafeCast v :: VS.Vector Word8
   in BS.pack (VS.toList bytes)

-- | Convert a blob back to a vector.
blobToVector :: BS.ByteString -> VS.Vector Float
blobToVector bs =
  let bytes = VS.fromList (BS.unpack bs) :: VS.Vector Word8
   in VS.unsafeCast bytes

-- | Calculate cosine similarity between two vectors.
cosineSimilarity :: VS.Vector Float -> VS.Vector Float -> Float
cosineSimilarity v1 v2
  | VS.length v1 /= VS.length v2 = 0
  | otherwise =
      let dot = VS.sum (VS.zipWith (*) v1 v2)
          mag1 = sqrt (VS.sum (VS.map (\x -> x * x) v1))
          mag2 = sqrt (VS.sum (VS.map (\x -> x * x) v2))
       in if mag1 == 0 || mag2 == 0 then 0 else dot / (mag1 * mag2)

-- | Format memories for inclusion in a prompt.
formatMemoriesForPrompt :: [Memory] -> Text
formatMemoriesForPrompt [] = "No prior context available."
formatMemoriesForPrompt mems =
  Text.unlines
    [ "Known context about this user:",
      "",
      Text.unlines (map formatMem mems)
    ]
  where
    formatMem m =
      "- " <> memoryContent m <> " (via " <> sourceAgent (memorySource m) <> ")"

-- | Run an agent with memory context.
-- Recalls relevant memories for the user and injects them into the system prompt.
runAgentWithMemory ::
  User ->
  Engine.EngineConfig ->
  Engine.AgentConfig ->
  Text ->
  IO (Either Text Engine.AgentResult)
runAgentWithMemory user engineCfg agentCfg userPrompt = do
  memories <- recallMemories (userId user) userPrompt 10
  let memoryContext = formatMemoriesForPrompt memories
      enhancedPrompt =
        Engine.agentSystemPrompt agentCfg
          <> "\n\n## Known about this user\n"
          <> memoryContext
      enhancedConfig =
        agentCfg
          { Engine.agentSystemPrompt = enhancedPrompt,
            Engine.agentTools =
              Engine.agentTools agentCfg
                <> [ rememberTool (userId user),
                     recallTool (userId user),
                     linkMemoriesTool (userId user),
                     queryGraphTool (userId user)
                   ]
          }
  Engine.runAgent engineCfg enhancedConfig userPrompt

-- | Tool for agents to store memories about users.
rememberTool :: Text -> Engine.Tool
rememberTool uid =
  Engine.Tool
    { Engine.toolName = "remember",
      Engine.toolDescription =
        "Store a piece of information about the user for future reference. "
          <> "Use this when the user shares personal facts, preferences, or context "
          <> "that would be useful to recall in future conversations.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "content"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("The information to remember about the user" :: Text)
                      ],
                  "context"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("How/why this was learned (e.g., 'user mentioned in chat')" :: Text)
                      ],
                  "tags"
                    .= Aeson.object
                      [ "type" .= ("array" :: Text),
                        "items" .= Aeson.object ["type" .= ("string" :: Text)],
                        "description" .= ("Optional tags for categorization" :: Text)
                      ]
                ],
            "required" .= (["content", "context"] :: [Text])
          ],
      Engine.toolExecute = executeRemember uid
    }

executeRemember :: Text -> Aeson.Value -> IO Aeson.Value
executeRemember uid v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e])
    Aeson.Success (args :: RememberArgs) -> do
      let source =
            MemorySource
              { sourceAgent = "agent",
                sourceSession = Nothing,
                sourceContext = rememberContext args
              }
      mem <- storeMemoryWithTags uid (rememberContent args) source (rememberTags args)
      pure
        ( Aeson.object
            [ "success" .= True,
              "memory_id" .= memoryId mem,
              "message" .= ("Remembered: " <> rememberContent args)
            ]
        )

-- | Tool for agents to recall memories about users.
recallTool :: Text -> Engine.Tool
recallTool uid =
  Engine.Tool
    { Engine.toolName = "recall",
      Engine.toolDescription =
        "Search your memory for information about the user. "
          <> "Use this to retrieve previously stored facts, preferences, or context.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "query"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("What to search for in memory" :: Text)
                      ],
                  "limit"
                    .= Aeson.object
                      [ "type" .= ("integer" :: Text),
                        "description" .= ("Maximum memories to return (default: 5)" :: Text)
                      ]
                ],
            "required" .= (["query"] :: [Text])
          ],
      Engine.toolExecute = executeRecall uid
    }

executeRecall :: Text -> Aeson.Value -> IO Aeson.Value
executeRecall uid v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e])
    Aeson.Success (args :: RecallArgs) -> do
      mems <- recallMemories uid (recallQuery args) (recallLimit args)
      pure
        ( Aeson.object
            [ "success" .= True,
              "count" .= length mems,
              "memories"
                .= map
                  ( \m ->
                      Aeson.object
                        [ "id" .= memoryId m,
                          "content" .= memoryContent m,
                          "confidence" .= memoryConfidence m,
                          "source" .= sourceAgent (memorySource m),
                          "tags" .= memoryTags m
                        ]
                  )
                  mems
            ]
        )

-- Helper for parsing remember args
data RememberArgs = RememberArgs
  { rememberContent :: Text,
    rememberContext :: Text,
    rememberTags :: [Text]
  }
  deriving (Generic)

instance Aeson.FromJSON RememberArgs where
  parseJSON =
    Aeson.withObject "RememberArgs" <| \v ->
      (RememberArgs </ (v .: "content"))
        <*> (v .:? "context" .!= "agent observation")
        <*> (v .:? "tags" .!= [])

data RecallArgs = RecallArgs
  { recallQuery :: Text,
    recallLimit :: Int
  }
  deriving (Generic)

instance Aeson.FromJSON RecallArgs where
  parseJSON =
    Aeson.withObject "RecallArgs" <| \v ->
      (RecallArgs </ (v .: "query"))
        <*> (v .:? "limit" .!= 5)

-- | Tool for agents to link memories in the knowledge graph.
linkMemoriesTool :: Text -> Engine.Tool
linkMemoriesTool _uid =
  Engine.Tool
    { Engine.toolName = "link_memories",
      Engine.toolDescription =
        "Create a typed relationship between two memories. "
          <> "Use this to connect related information. Relation types:\n"
          <> "- contradicts: conflicting information\n"
          <> "- supports: evidence that reinforces another memory\n"
          <> "- elaborates: adds detail to an existing memory\n"
          <> "- supersedes: newer info replaces older\n"
          <> "- related: general topical connection\n"
          <> "- contingent_on: depends on another fact being true",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "from_memory_id"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("ID of the source memory" :: Text)
                      ],
                  "to_memory_id"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("ID of the target memory" :: Text)
                      ],
                  "relation_type"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "enum" .= (["contradicts", "supports", "elaborates", "supersedes", "related", "contingent_on"] :: [Text]),
                        "description" .= ("Type of relationship between memories" :: Text)
                      ]
                ],
            "required" .= (["from_memory_id", "to_memory_id", "relation_type"] :: [Text])
          ],
      Engine.toolExecute = executeLinkMemories
    }

executeLinkMemories :: Aeson.Value -> IO Aeson.Value
executeLinkMemories v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e])
    Aeson.Success (args :: LinkMemoriesArgs) -> do
      case textToRelationType (linkArgsRelationType args) of
        Nothing ->
          pure
            ( Aeson.object
                [ "success" .= False,
                  "error" .= ("Invalid relation type: " <> linkArgsRelationType args)
                ]
            )
        Just relType -> do
          memLink <- linkMemories (linkArgsFromId args) (linkArgsToId args) relType
          pure
            ( Aeson.object
                [ "success" .= True,
                  "message"
                    .= ( "Linked memory "
                           <> linkFromMemoryId memLink
                           <> " -> "
                           <> linkToMemoryId memLink
                           <> " ("
                           <> relationTypeToText (linkRelationType memLink)
                           <> ")"
                       )
                ]
            )

data LinkMemoriesArgs = LinkMemoriesArgs
  { linkArgsFromId :: Text,
    linkArgsToId :: Text,
    linkArgsRelationType :: Text
  }
  deriving (Generic)

instance Aeson.FromJSON LinkMemoriesArgs where
  parseJSON =
    Aeson.withObject "LinkMemoriesArgs" <| \v ->
      (LinkMemoriesArgs </ (v .: "from_memory_id"))
        <*> (v .: "to_memory_id")
        <*> (v .: "relation_type")

-- | Tool for agents to query the memory knowledge graph.
queryGraphTool :: Text -> Engine.Tool
queryGraphTool _uid =
  Engine.Tool
    { Engine.toolName = "query_graph",
      Engine.toolDescription =
        "Explore the knowledge graph to find related memories. "
          <> "Given a starting memory, traverse links to find connected memories. "
          <> "Useful for understanding context and finding contradictions or supporting evidence.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "memory_id"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("ID of the memory to start from" :: Text)
                      ],
                  "depth"
                    .= Aeson.object
                      [ "type" .= ("integer" :: Text),
                        "description" .= ("How many link hops to traverse (default: 2)" :: Text)
                      ],
                  "relation_type"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "enum" .= (["contradicts", "supports", "elaborates", "supersedes", "related", "contingent_on"] :: [Text]),
                        "description" .= ("Optional: filter by relation type" :: Text)
                      ]
                ],
            "required" .= (["memory_id"] :: [Text])
          ],
      Engine.toolExecute = executeQueryGraph
    }

executeQueryGraph :: Aeson.Value -> IO Aeson.Value
executeQueryGraph v =
  case Aeson.fromJSON v of
    Aeson.Error e -> pure (Aeson.object ["error" .= Text.pack e])
    Aeson.Success (args :: QueryGraphArgs) -> do
      let maybeRelType = queryArgsRelationType args +> textToRelationType
      results <- queryGraph (queryArgsMemoryId args) (queryArgsDepth args) maybeRelType
      pure
        ( Aeson.object
            [ "success" .= True,
              "count" .= length results,
              "memories"
                .= map
                  ( \(m, links) ->
                      Aeson.object
                        [ "id" .= memoryId m,
                          "content" .= memoryContent m,
                          "links"
                            .= map
                              ( \l ->
                                  Aeson.object
                                    [ "from" .= linkFromMemoryId l,
                                      "to" .= linkToMemoryId l,
                                      "relation" .= linkRelationType l
                                    ]
                              )
                              links
                        ]
                  )
                  results
            ]
        )

data QueryGraphArgs = QueryGraphArgs
  { queryArgsMemoryId :: Text,
    queryArgsDepth :: Int,
    queryArgsRelationType :: Maybe Text
  }
  deriving (Generic)

instance Aeson.FromJSON QueryGraphArgs where
  parseJSON =
    Aeson.withObject "QueryGraphArgs" <| \v ->
      (QueryGraphArgs </ (v .: "memory_id"))
        <*> (v .:? "depth" .!= 2)
        <*> (v .:? "relation_type")

-- | Estimate token count for text (rough: ~4 chars per token).
estimateTokens :: Text -> Int
estimateTokens t = max 1 (Text.length t `div` 4)

-- | Save a message to conversation history.
saveMessage :: Text -> Int -> MessageRole -> Maybe Text -> Text -> IO ConversationMessage
saveMessage uid chatId role senderName content = do
  now <- getCurrentTime
  let tokens = estimateTokens content
  withMemoryDb <| \conn -> do
    SQL.execute
      conn
      "INSERT INTO conversation_messages (user_id, chat_id, role, sender_name, content, tokens_estimate, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)"
      (uid, chatId, roleToText role, senderName, content, tokens, now)
    rowId <- SQL.lastInsertRowId conn
    pure
      ConversationMessage
        { cmId = Just (fromIntegral rowId),
          cmUserId = uid,
          cmChatId = chatId,
          cmRole = role,
          cmSenderName = senderName,
          cmContent = content,
          cmTokensEstimate = tokens,
          cmCreatedAt = now
        }
  where
    roleToText UserRole = "user" :: Text
    roleToText AssistantRole = "assistant"

-- | Get recent messages for a user/chat, newest first.
getRecentMessages :: Text -> Int -> Int -> IO [ConversationMessage]
getRecentMessages uid chatId limit =
  withMemoryDb <| \conn ->
    SQL.query
      conn
      "SELECT id, user_id, chat_id, role, sender_name, content, tokens_estimate, created_at \
      \FROM conversation_messages \
      \WHERE user_id = ? AND chat_id = ? \
      \ORDER BY created_at DESC LIMIT ?"
      (uid, chatId, limit)

-- | Get the most recent summary for a chat.
getLatestSummary :: Text -> Int -> IO (Maybe ConversationSummary)
getLatestSummary uid chatId =
  withMemoryDb <| \conn -> do
    rows <-
      SQL.query
        conn
        "SELECT id, user_id, chat_id, summary, messages_summarized, tokens_saved, created_at \
        \FROM conversation_summaries \
        \WHERE user_id = ? AND chat_id = ? \
        \ORDER BY created_at DESC LIMIT 1"
        (uid, chatId)
    pure (listToMaybe rows)

-- | Build conversation context for the LLM.
-- Returns (context text, total token estimate).
getConversationContext :: Text -> Int -> Int -> IO (Text, Int)
getConversationContext uid chatId maxTokens = do
  maybeSummary <- getLatestSummary uid chatId
  recentMsgs <- getRecentMessages uid chatId 50

  let summaryText = maybe "" (\s -> "## Previous conversation summary\n" <> csSummary s <> "\n\n") maybeSummary
      summaryTokens = maybe 0 (estimateTokens <. csSummary) maybeSummary

      msgsOldestFirst = reverse recentMsgs
      availableTokens = maxTokens - summaryTokens - 100

      (selectedMsgs, usedTokens) = selectMessages msgsOldestFirst availableTokens

      formattedMsgs =
        if null selectedMsgs
          then ""
          else
            "## Recent conversation\n"
              <> Text.unlines (map formatMsg selectedMsgs)

  pure (summaryText <> formattedMsgs, summaryTokens + usedTokens)
  where
    selectMessages :: [ConversationMessage] -> Int -> ([ConversationMessage], Int)
    selectMessages msgs budget = go (reverse msgs) budget []
      where
        go [] _ acc = (acc, sum (map cmTokensEstimate acc))
        go (m : ms) remaining acc
          | cmTokensEstimate m <= remaining =
              go ms (remaining - cmTokensEstimate m) (m : acc)
          | otherwise = (acc, sum (map cmTokensEstimate acc))

    formatMsg m =
      let timestamp = Text.pack (formatTime defaultTimeLocale "%Y-%m-%dT%H:%M:%SZ" (cmCreatedAt m))
          prefix = case cmRole m of
            UserRole -> "[" <> timestamp <> "] " <> fromMaybe "User" (cmSenderName m) <> ": "
            AssistantRole -> "[" <> timestamp <> "] Assistant: "
       in prefix <> cmContent m

-- | Summarize old messages and archive them.
-- Returns the new summary text.
summarizeAndArchive :: Text -> Int -> Text -> IO Text
summarizeAndArchive uid chatId summaryText = do
  now <- getCurrentTime

  (oldMsgCount, tokensSaved) <-
    withMemoryDb <| \conn -> do
      rows <-
        SQL.query
          conn
          "SELECT COUNT(*), COALESCE(SUM(tokens_estimate), 0) FROM conversation_messages WHERE user_id = ? AND chat_id = ?"
          (uid, chatId) ::
          IO [(Int, Int)]
      let (count, tokens) = fromMaybe (0, 0) (listToMaybe rows)

      SQL.execute
        conn
        "INSERT INTO conversation_summaries (user_id, chat_id, summary, messages_summarized, tokens_saved, created_at) VALUES (?, ?, ?, ?, ?, ?)"
        (uid, chatId, summaryText, count, tokens, now)

      SQL.execute
        conn
        "DELETE FROM conversation_messages WHERE user_id = ? AND chat_id = ?"
        (uid, chatId)

      pure (count, tokens)

  putText <| "Archived " <> tshow oldMsgCount <> " messages (" <> tshow tokensSaved <> " tokens) for chat " <> tshow chatId
  pure summaryText

-- -----------------------------------------------------------------------------
-- Group Conversation History
-- -----------------------------------------------------------------------------

-- | Save a message to group conversation history.
-- Unlike saveMessage, this is keyed by (chat_id, thread_id) not (user_id, chat_id).
-- The sender_name is preserved for attribution.
saveGroupMessage :: Int -> Maybe Int -> MessageRole -> Text -> Text -> IO ConversationMessage
saveGroupMessage chatId mThreadId role senderName content = do
  now <- getCurrentTime
  let tokens = estimateTokens content
  withMemoryDb <| \conn -> do
    SQL.execute
      conn
      "INSERT INTO conversation_messages (user_id, chat_id, thread_id, role, sender_name, content, tokens_estimate, created_at) VALUES (NULL, ?, ?, ?, ?, ?, ?, ?)"
      (chatId, mThreadId, roleToText role, senderName, content, tokens, now)
    rowId <- SQL.lastInsertRowId conn
    pure
      ConversationMessage
        { cmId = Just (fromIntegral rowId),
          cmUserId = "",
          cmChatId = chatId,
          cmRole = role,
          cmSenderName = Just senderName,
          cmContent = content,
          cmTokensEstimate = tokens,
          cmCreatedAt = now
        }
  where
    roleToText UserRole = "user" :: Text
    roleToText AssistantRole = "assistant"

-- | Get recent messages for a group chat/topic, newest first.
getGroupRecentMessages :: Int -> Maybe Int -> Int -> IO [ConversationMessage]
getGroupRecentMessages chatId mThreadId limit =
  withMemoryDb <| \conn ->
    case mThreadId of
      Just threadId ->
        SQL.query
          conn
          "SELECT id, COALESCE(user_id, ''), chat_id, role, sender_name, content, tokens_estimate, created_at \
          \FROM conversation_messages \
          \WHERE chat_id = ? AND thread_id = ? \
          \ORDER BY created_at DESC LIMIT ?"
          (chatId, threadId, limit)
      Nothing ->
        SQL.query
          conn
          "SELECT id, COALESCE(user_id, ''), chat_id, role, sender_name, content, tokens_estimate, created_at \
          \FROM conversation_messages \
          \WHERE chat_id = ? AND thread_id IS NULL \
          \ORDER BY created_at DESC LIMIT ?"
          (chatId, limit)

-- | Build conversation context for a group chat.
-- Returns (context text, total token estimate).
getGroupConversationContext :: Int -> Maybe Int -> Int -> IO (Text, Int)
getGroupConversationContext chatId mThreadId maxTokens = do
  recentMsgs <- getGroupRecentMessages chatId mThreadId 50

  let msgsOldestFirst = reverse recentMsgs
      availableTokens = maxTokens - 100

      (selectedMsgs, usedTokens) = selectMessages msgsOldestFirst availableTokens

      formattedMsgs =
        if null selectedMsgs
          then ""
          else
            "## Recent conversation\n"
              <> Text.unlines (map formatMsg selectedMsgs)

  pure (formattedMsgs, usedTokens)
  where
    selectMessages :: [ConversationMessage] -> Int -> ([ConversationMessage], Int)
    selectMessages msgs budget = go (reverse msgs) budget []
      where
        go [] _ acc = (acc, sum (map cmTokensEstimate acc))
        go (m : ms) remaining acc
          | cmTokensEstimate m <= remaining =
              go ms (remaining - cmTokensEstimate m) (m : acc)
          | otherwise = (acc, sum (map cmTokensEstimate acc))

    formatMsg m =
      let timestamp = Text.pack (formatTime defaultTimeLocale "%Y-%m-%dT%H:%M:%SZ" (cmCreatedAt m))
          prefix = case cmRole m of
            UserRole -> "[" <> timestamp <> "] " <> fromMaybe "User" (cmSenderName m) <> ": "
            AssistantRole -> "[" <> timestamp <> "] Assistant: "
       in prefix <> cmContent m

-- -----------------------------------------------------------------------------
-- Group Memories
-- -----------------------------------------------------------------------------

-- | Generate a synthetic user_id for group-level memories.
groupUserId :: Int -> Text
groupUserId chatId = "group:" <> tshow chatId

-- | Store a memory associated with a group (not a user).
-- These memories are shared across all users in the group.
storeGroupMemory :: Int -> Text -> MemorySource -> IO Memory
storeGroupMemory chatId = storeMemory (groupUserId chatId)

-- | Recall memories for a group.
recallGroupMemories :: Int -> Text -> Int -> IO [Memory]
recallGroupMemories chatId = recallMemories (groupUserId chatId)