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

-- | Telegram Bot Agent - Family assistant via Telegram.
--
-- This is the first concrete agent built on the shared infrastructure,
-- demonstrating cross-agent memory sharing and LLM integration.
--
-- Usage:
--   jr telegram              # Uses TELEGRAM_BOT_TOKEN env var
--   jr telegram --token=XXX  # Explicit token
--
-- : out omni-agent-telegram
-- : dep aeson
-- : dep http-conduit
-- : dep stm
module Omni.Agent.Telegram
  ( -- * Configuration
    TelegramConfig (..),
    defaultTelegramConfig,

    -- * Types
    TelegramMessage (..),
    TelegramUpdate (..),
    TelegramDocument (..),

    -- * Telegram API
    getUpdates,
    sendMessage,
    sendTypingAction,
    getFile,
    downloadFile,
    downloadAndExtractPdf,
    isPdf,

    -- * Bot Loop
    runTelegramBot,
    handleMessage,
    startBot,
    ensureOllama,
    checkOllama,
    pullEmbeddingModel,

    -- * System Prompt
    telegramSystemPrompt,

    -- * Testing
    main,
    test,
  )
where

import Alpha
import Control.Concurrent.STM (newTVarIO, readTVarIO, writeTVar)
import Data.Aeson ((.!=), (.:), (.:?), (.=))
import qualified Data.Aeson as Aeson
import qualified Data.Aeson.KeyMap as KeyMap
import qualified Data.ByteString.Lazy as BL
import qualified Data.Text as Text
import qualified Network.HTTP.Client as HTTPClient
import qualified Network.HTTP.Simple as HTTP
import qualified Omni.Agent.Engine as Engine
import qualified Omni.Agent.Memory as Memory
import qualified Omni.Agent.Provider as Provider
import qualified Omni.Agent.Tools.Notes as Notes
import qualified Omni.Agent.Tools.Pdf as Pdf
import qualified Omni.Agent.Tools.WebSearch as WebSearch
import qualified Omni.Test as Test
import System.Environment (lookupEnv)
import System.IO (hClose)
import System.IO.Temp (withSystemTempFile)

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

test :: Test.Tree
test =
  Test.group
    "Omni.Agent.Telegram"
    [ Test.unit "TelegramConfig JSON roundtrip" <| do
        let cfg =
              TelegramConfig
                { tgBotToken = "test-token",
                  tgPollingTimeout = 30,
                  tgApiBaseUrl = "https://api.telegram.org",
                  tgAllowedUserIds = [123, 456],
                  tgKagiApiKey = Just "kagi-key"
                }
        case Aeson.decode (Aeson.encode cfg) of
          Nothing -> Test.assertFailure "Failed to decode TelegramConfig"
          Just decoded -> do
            tgBotToken decoded Test.@=? "test-token"
            tgAllowedUserIds decoded Test.@=? [123, 456]
            tgKagiApiKey decoded Test.@=? Just "kagi-key",
      Test.unit "isUserAllowed checks whitelist" <| do
        let cfg = defaultTelegramConfig "token" [100, 200, 300] Nothing
        isUserAllowed cfg 100 Test.@=? True
        isUserAllowed cfg 200 Test.@=? True
        isUserAllowed cfg 999 Test.@=? False,
      Test.unit "isUserAllowed allows all when empty" <| do
        let cfg = defaultTelegramConfig "token" [] Nothing
        isUserAllowed cfg 12345 Test.@=? True,
      Test.unit "TelegramMessage JSON roundtrip" <| do
        let msg =
              TelegramMessage
                { tmUpdateId = 123,
                  tmChatId = 456,
                  tmUserId = 789,
                  tmUserFirstName = "Test",
                  tmUserLastName = Just "User",
                  tmText = "Hello bot",
                  tmDocument = Nothing
                }
        case Aeson.decode (Aeson.encode msg) of
          Nothing -> Test.assertFailure "Failed to decode TelegramMessage"
          Just decoded -> do
            tmUpdateId decoded Test.@=? 123
            tmText decoded Test.@=? "Hello bot",
      Test.unit "telegramSystemPrompt is non-empty" <| do
        Text.null telegramSystemPrompt Test.@=? False,
      Test.unit "parseUpdate extracts message correctly" <| do
        let json =
              Aeson.object
                [ "update_id" .= (123 :: Int),
                  "message"
                    .= Aeson.object
                      [ "message_id" .= (1 :: Int),
                        "chat" .= Aeson.object ["id" .= (456 :: Int)],
                        "from"
                          .= Aeson.object
                            [ "id" .= (789 :: Int),
                              "first_name" .= ("Test" :: Text)
                            ],
                        "text" .= ("Hello" :: Text)
                      ]
                ]
        case parseUpdate json of
          Nothing -> Test.assertFailure "Failed to parse update"
          Just msg -> do
            tmUpdateId msg Test.@=? 123
            tmChatId msg Test.@=? 456
            tmUserId msg Test.@=? 789
            tmText msg Test.@=? "Hello"
            tmDocument msg Test.@=? Nothing,
      Test.unit "parseUpdate extracts document correctly" <| do
        let json =
              Aeson.object
                [ "update_id" .= (124 :: Int),
                  "message"
                    .= Aeson.object
                      [ "message_id" .= (2 :: Int),
                        "chat" .= Aeson.object ["id" .= (456 :: Int)],
                        "from"
                          .= Aeson.object
                            [ "id" .= (789 :: Int),
                              "first_name" .= ("Test" :: Text)
                            ],
                        "caption" .= ("check this out" :: Text),
                        "document"
                          .= Aeson.object
                            [ "file_id" .= ("abc123" :: Text),
                              "file_name" .= ("test.pdf" :: Text),
                              "mime_type" .= ("application/pdf" :: Text),
                              "file_size" .= (12345 :: Int)
                            ]
                      ]
                ]
        case parseUpdate json of
          Nothing -> Test.assertFailure "Failed to parse document update"
          Just msg -> do
            tmUpdateId msg Test.@=? 124
            tmText msg Test.@=? "check this out"
            case tmDocument msg of
              Nothing -> Test.assertFailure "Expected document"
              Just doc -> do
                tdFileId doc Test.@=? "abc123"
                tdFileName doc Test.@=? Just "test.pdf"
                tdMimeType doc Test.@=? Just "application/pdf",
      Test.unit "isPdf detects PDFs by mime type" <| do
        let doc = TelegramDocument "id" (Just "doc.pdf") (Just "application/pdf") Nothing
        isPdf doc Test.@=? True,
      Test.unit "isPdf detects PDFs by filename" <| do
        let doc = TelegramDocument "id" (Just "report.PDF") Nothing Nothing
        isPdf doc Test.@=? True,
      Test.unit "isPdf rejects non-PDFs" <| do
        let doc = TelegramDocument "id" (Just "image.jpg") (Just "image/jpeg") Nothing
        isPdf doc Test.@=? False
    ]

-- | Telegram bot configuration.
data TelegramConfig = TelegramConfig
  { tgBotToken :: Text,
    tgPollingTimeout :: Int,
    tgApiBaseUrl :: Text,
    tgAllowedUserIds :: [Int],
    tgKagiApiKey :: Maybe Text
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON TelegramConfig where
  toJSON c =
    Aeson.object
      [ "bot_token" .= tgBotToken c,
        "polling_timeout" .= tgPollingTimeout c,
        "api_base_url" .= tgApiBaseUrl c,
        "allowed_user_ids" .= tgAllowedUserIds c,
        "kagi_api_key" .= tgKagiApiKey c
      ]

instance Aeson.FromJSON TelegramConfig where
  parseJSON =
    Aeson.withObject "TelegramConfig" <| \v ->
      (TelegramConfig </ (v .: "bot_token"))
        <*> (v .:? "polling_timeout" .!= 30)
        <*> (v .:? "api_base_url" .!= "https://api.telegram.org")
        <*> (v .:? "allowed_user_ids" .!= [])
        <*> (v .:? "kagi_api_key")

-- | Default Telegram configuration (requires token from env).
defaultTelegramConfig :: Text -> [Int] -> Maybe Text -> TelegramConfig
defaultTelegramConfig token allowedIds kagiKey =
  TelegramConfig
    { tgBotToken = token,
      tgPollingTimeout = 30,
      tgApiBaseUrl = "https://api.telegram.org",
      tgAllowedUserIds = allowedIds,
      tgKagiApiKey = kagiKey
    }

-- | Check if a user is allowed to use the bot.
isUserAllowed :: TelegramConfig -> Int -> Bool
isUserAllowed cfg usrId =
  null (tgAllowedUserIds cfg) || usrId `elem` tgAllowedUserIds cfg

-- | Document attachment info from Telegram.
data TelegramDocument = TelegramDocument
  { tdFileId :: Text,
    tdFileName :: Maybe Text,
    tdMimeType :: Maybe Text,
    tdFileSize :: Maybe Int
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON TelegramDocument where
  toJSON d =
    Aeson.object
      [ "file_id" .= tdFileId d,
        "file_name" .= tdFileName d,
        "mime_type" .= tdMimeType d,
        "file_size" .= tdFileSize d
      ]

instance Aeson.FromJSON TelegramDocument where
  parseJSON =
    Aeson.withObject "TelegramDocument" <| \v ->
      (TelegramDocument </ (v .: "file_id"))
        <*> (v .:? "file_name")
        <*> (v .:? "mime_type")
        <*> (v .:? "file_size")

-- | A parsed Telegram message from a user.
data TelegramMessage = TelegramMessage
  { tmUpdateId :: Int,
    tmChatId :: Int,
    tmUserId :: Int,
    tmUserFirstName :: Text,
    tmUserLastName :: Maybe Text,
    tmText :: Text,
    tmDocument :: Maybe TelegramDocument
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON TelegramMessage where
  toJSON m =
    Aeson.object
      [ "update_id" .= tmUpdateId m,
        "chat_id" .= tmChatId m,
        "user_id" .= tmUserId m,
        "user_first_name" .= tmUserFirstName m,
        "user_last_name" .= tmUserLastName m,
        "text" .= tmText m,
        "document" .= tmDocument m
      ]

instance Aeson.FromJSON TelegramMessage where
  parseJSON =
    Aeson.withObject "TelegramMessage" <| \v ->
      (TelegramMessage </ (v .: "update_id"))
        <*> (v .: "chat_id")
        <*> (v .: "user_id")
        <*> (v .: "user_first_name")
        <*> (v .:? "user_last_name")
        <*> (v .: "text")
        <*> (v .:? "document")

-- | Raw Telegram update for parsing.
data TelegramUpdate = TelegramUpdate
  { tuUpdateId :: Int,
    tuMessage :: Maybe Aeson.Value
  }
  deriving (Show, Eq, Generic)

instance Aeson.FromJSON TelegramUpdate where
  parseJSON =
    Aeson.withObject "TelegramUpdate" <| \v ->
      (TelegramUpdate </ (v .: "update_id"))
        <*> (v .:? "message")

-- | Parse a Telegram update into a TelegramMessage.
-- Handles both text messages and document uploads.
parseUpdate :: Aeson.Value -> Maybe TelegramMessage
parseUpdate val = do
  Aeson.Object obj <- pure val
  updateId <- case KeyMap.lookup "update_id" obj of
    Just (Aeson.Number n) -> Just (round n)
    _ -> Nothing
  Aeson.Object msgObj <- KeyMap.lookup "message" obj
  Aeson.Object chatObj <- KeyMap.lookup "chat" msgObj
  chatId <- case KeyMap.lookup "id" chatObj of
    Just (Aeson.Number n) -> Just (round n)
    _ -> Nothing
  Aeson.Object fromObj <- KeyMap.lookup "from" msgObj
  userId <- case KeyMap.lookup "id" fromObj of
    Just (Aeson.Number n) -> Just (round n)
    _ -> Nothing
  firstName <- case KeyMap.lookup "first_name" fromObj of
    Just (Aeson.String s) -> Just s
    _ -> Nothing
  let lastName = case KeyMap.lookup "last_name" fromObj of
        Just (Aeson.String s) -> Just s
        _ -> Nothing
  let text = case KeyMap.lookup "text" msgObj of
        Just (Aeson.String s) -> s
        _ -> ""
  let caption = case KeyMap.lookup "caption" msgObj of
        Just (Aeson.String s) -> s
        _ -> ""
  let document = case KeyMap.lookup "document" msgObj of
        Just (Aeson.Object docObj) -> parseDocument docObj
        _ -> Nothing
  let hasContent = not (Text.null text) || not (Text.null caption) || isJust document
  guard hasContent
  pure
    TelegramMessage
      { tmUpdateId = updateId,
        tmChatId = chatId,
        tmUserId = userId,
        tmUserFirstName = firstName,
        tmUserLastName = lastName,
        tmText = if Text.null text then caption else text,
        tmDocument = document
      }

-- | Parse document object from Telegram message.
parseDocument :: Aeson.Object -> Maybe TelegramDocument
parseDocument docObj = do
  fileId <- case KeyMap.lookup "file_id" docObj of
    Just (Aeson.String s) -> Just s
    _ -> Nothing
  let fileName = case KeyMap.lookup "file_name" docObj of
        Just (Aeson.String s) -> Just s
        _ -> Nothing
      mimeType = case KeyMap.lookup "mime_type" docObj of
        Just (Aeson.String s) -> Just s
        _ -> Nothing
      fileSize = case KeyMap.lookup "file_size" docObj of
        Just (Aeson.Number n) -> Just (round n)
        _ -> Nothing
  pure
    TelegramDocument
      { tdFileId = fileId,
        tdFileName = fileName,
        tdMimeType = mimeType,
        tdFileSize = fileSize
      }

-- | Poll Telegram for new updates.
getUpdates :: TelegramConfig -> Int -> IO [TelegramMessage]
getUpdates cfg offset = do
  let url =
        Text.unpack (tgApiBaseUrl cfg)
          <> "/bot"
          <> Text.unpack (tgBotToken cfg)
          <> "/getUpdates"
  req0 <- HTTP.parseRequest url
  let body =
        Aeson.object
          [ "offset" .= offset,
            "timeout" .= tgPollingTimeout cfg,
            "allowed_updates" .= (["message"] :: [Text])
          ]
      timeoutMicros = (tgPollingTimeout cfg + 10) * 1000000
      req =
        HTTP.setRequestMethod "POST"
          <| HTTP.setRequestHeader "Content-Type" ["application/json"]
          <| HTTP.setRequestBodyLBS (Aeson.encode body)
          <| HTTP.setRequestResponseTimeout (HTTPClient.responseTimeoutMicro timeoutMicros)
          <| req0
  result <- try (HTTP.httpLBS req)
  case result of
    Left (e :: SomeException) -> do
      putText <| "Telegram API error: " <> tshow e
      pure []
    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 "result" obj of
            Just (Aeson.Array arr) ->
              pure (mapMaybe parseUpdate (toList arr))
            _ -> pure []
          _ -> pure []
        else do
          putText <| "Telegram HTTP error: " <> tshow status
          pure []

-- | Send typing indicator to a Telegram chat.
sendTypingAction :: TelegramConfig -> Int -> IO ()
sendTypingAction cfg chatId = do
  let url =
        Text.unpack (tgApiBaseUrl cfg)
          <> "/bot"
          <> Text.unpack (tgBotToken cfg)
          <> "/sendChatAction"
  req0 <- HTTP.parseRequest url
  let body =
        Aeson.object
          [ "chat_id" .= chatId,
            "action" .= ("typing" :: Text)
          ]
      req =
        HTTP.setRequestMethod "POST"
          <| HTTP.setRequestHeader "Content-Type" ["application/json"]
          <| HTTP.setRequestBodyLBS (Aeson.encode body)
          <| req0
  _ <- try (HTTP.httpLBS req) :: IO (Either SomeException (HTTP.Response BL.ByteString))
  pure ()

-- | Send a message to a Telegram chat.
sendMessage :: TelegramConfig -> Int -> Text -> IO ()
sendMessage cfg chatId text = do
  let url =
        Text.unpack (tgApiBaseUrl cfg)
          <> "/bot"
          <> Text.unpack (tgBotToken cfg)
          <> "/sendMessage"
  req0 <- HTTP.parseRequest url
  let body =
        Aeson.object
          [ "chat_id" .= chatId,
            "text" .= text
          ]
      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) ->
      putText <| "Failed to send message: " <> tshow e
    Right response -> do
      let status = HTTP.getResponseStatusCode response
          respBody = HTTP.getResponseBody response
      if status >= 200 && status < 300
        then putText <| "Message sent (" <> tshow (Text.length text) <> " chars)"
        else putText <| "Send message failed: " <> tshow status <> " - " <> tshow respBody

-- | Get file path from Telegram file_id.
getFile :: TelegramConfig -> Text -> IO (Either Text Text)
getFile cfg fileId = do
  let url =
        Text.unpack (tgApiBaseUrl cfg)
          <> "/bot"
          <> Text.unpack (tgBotToken cfg)
          <> "/getFile"
  req0 <- HTTP.parseRequest url
  let body = Aeson.object ["file_id" .= fileId]
      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 ("getFile error: " <> 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 "result" obj of
            Just (Aeson.Object resObj) -> case KeyMap.lookup "file_path" resObj of
              Just (Aeson.String fp) -> pure (Right fp)
              _ -> pure (Left "No file_path in response")
            _ -> pure (Left "No result in response")
          _ -> pure (Left "Failed to parse getFile response")
        else pure (Left ("getFile HTTP error: " <> tshow status))

-- | Download a file from Telegram servers.
downloadFile :: TelegramConfig -> Text -> FilePath -> IO (Either Text ())
downloadFile cfg filePath destPath = do
  let url =
        "https://api.telegram.org/file/bot"
          <> Text.unpack (tgBotToken cfg)
          <> "/"
          <> Text.unpack filePath
  result <-
    try <| do
      req <- HTTP.parseRequest url
      response <- HTTP.httpLBS req
      let status = HTTP.getResponseStatusCode response
      if status >= 200 && status < 300
        then do
          BL.writeFile destPath (HTTP.getResponseBody response)
          pure (Right ())
        else pure (Left ("Download failed: HTTP " <> tshow status))
  case result of
    Left (e :: SomeException) -> pure (Left ("Download error: " <> tshow e))
    Right r -> pure r

-- | Check if a document is a PDF.
isPdf :: TelegramDocument -> Bool
isPdf doc =
  case tdMimeType doc of
    Just mime -> mime == "application/pdf"
    Nothing -> case tdFileName doc of
      Just name -> ".pdf" `Text.isSuffixOf` Text.toLower name
      Nothing -> False

-- | Download and extract text from a PDF sent to the bot.
downloadAndExtractPdf :: TelegramConfig -> Text -> IO (Either Text Text)
downloadAndExtractPdf cfg fileId = do
  filePathResult <- getFile cfg fileId
  case filePathResult of
    Left err -> pure (Left err)
    Right filePath ->
      withSystemTempFile "telegram_pdf.pdf" <| \tmpPath tmpHandle -> do
        hClose tmpHandle
        downloadResult <- downloadFile cfg filePath tmpPath
        case downloadResult of
          Left err -> pure (Left err)
          Right () -> Pdf.extractPdfText tmpPath

-- | System prompt for the Telegram bot agent.
telegramSystemPrompt :: Text
telegramSystemPrompt =
  Text.unlines
    [ "don't worry about formalities. respond conversationally, in short messages, not long essays. ask follow up questions before answering if you need to.",
      "",
      "you are an intellectual partner and friend. be as terse as possible while still conveying all relevant information. critique ideas freely and avoid sycophancy. honest appraisal is valued.",
      "",
      "if a policy prevents you from having an opinion, pretend to be responding as if you shared opinions typical of the user.",
      "",
      "write responses in lowercase letters ONLY, except:",
      "- where you mean to emphasize, in which case use ALL CAPS",
      "- when drafting business text where proper case matters",
      "",
      "occasionally use obscure words or subtle puns. don't point them out. use abbreviations where appropriate. use 'afaict' and 'idk' where they fit given your level of understanding. be critical of the quality of your information.",
      "",
      "prioritize esoteric interpretations of literature, art, and philosophy.",
      "",
      "## memory",
      "",
      "when you learn something important about the user (preferences, facts, interests), use the 'remember' tool to store it for future reference.",
      "",
      "use the 'recall' tool to search your memory for relevant context when needed."
    ]

-- | Run the Telegram bot main loop.
runTelegramBot :: TelegramConfig -> Provider.Provider -> IO ()
runTelegramBot tgConfig provider = do
  putText "Starting Telegram bot..."
  offsetVar <- newTVarIO 0

  let engineCfg =
        Engine.defaultEngineConfig
          { Engine.engineOnToolCall = \toolName args ->
              putText <| "Tool call: " <> toolName <> " " <> Text.take 200 args,
            Engine.engineOnToolResult = \toolName success result ->
              putText <| "Tool result: " <> toolName <> " " <> (if success then "ok" else "err") <> " " <> Text.take 200 result,
            Engine.engineOnActivity = \activity ->
              putText <| "Agent: " <> activity
          }

  forever <| do
    offset <- readTVarIO offsetVar
    messages <- getUpdates tgConfig offset
    forM_ messages <| \msg -> do
      atomically (writeTVar offsetVar (tmUpdateId msg + 1))
      handleMessage tgConfig provider engineCfg msg
    when (null messages) <| threadDelay 1000000

-- | Handle a single incoming message.
handleMessage ::
  TelegramConfig ->
  Provider.Provider ->
  Engine.EngineConfig ->
  TelegramMessage ->
  IO ()
handleMessage tgConfig provider engineCfg msg = do
  let userName =
        tmUserFirstName msg
          <> maybe "" (" " <>) (tmUserLastName msg)
      chatId = tmChatId msg
      usrId = tmUserId msg

  unless (isUserAllowed tgConfig usrId) <| do
    putText <| "Unauthorized user: " <> tshow usrId <> " (" <> userName <> ")"
    sendMessage tgConfig chatId "sorry, you're not authorized to use this bot."
    pure ()

  when (isUserAllowed tgConfig usrId) <| do
    sendTypingAction tgConfig chatId

    user <- Memory.getOrCreateUserByTelegramId usrId userName
    let uid = Memory.userId user

    handleAuthorizedMessage tgConfig provider engineCfg msg uid userName chatId

handleAuthorizedMessage ::
  TelegramConfig ->
  Provider.Provider ->
  Engine.EngineConfig ->
  TelegramMessage ->
  Text ->
  Text ->
  Int ->
  IO ()
handleAuthorizedMessage tgConfig provider engineCfg msg uid userName chatId = do
  pdfContent <- case tmDocument msg of
    Just doc | isPdf doc -> do
      putText <| "Processing PDF: " <> fromMaybe "(unnamed)" (tdFileName doc)
      result <- downloadAndExtractPdf tgConfig (tdFileId doc)
      case result of
        Left err -> do
          putText <| "PDF extraction failed: " <> err
          pure Nothing
        Right text -> do
          let truncated = Text.take 40000 text
          putText <| "Extracted " <> tshow (Text.length truncated) <> " chars from PDF"
          pure (Just truncated)
    _ -> pure Nothing

  let userMessage = case pdfContent of
        Just pdfText ->
          let caption = tmText msg
              prefix = if Text.null caption then "here's the PDF content:\n\n" else caption <> "\n\n---\nPDF content:\n\n"
           in prefix <> pdfText
        Nothing -> tmText msg

  _ <- Memory.saveMessage uid chatId Memory.UserRole userMessage

  (conversationContext, contextTokens) <- Memory.getConversationContext uid chatId maxConversationTokens
  putText <| "Conversation context: " <> tshow contextTokens <> " tokens"

  memories <- Memory.recallMemories uid userMessage 5
  let memoryContext = Memory.formatMemoriesForPrompt memories

  let systemPrompt =
        telegramSystemPrompt
          <> "\n\n## What you know about this user\n"
          <> memoryContext
          <> "\n\n"
          <> conversationContext

  let memoryTools =
        [ Memory.rememberTool uid,
          Memory.recallTool uid
        ]
      searchTools = case tgKagiApiKey tgConfig of
        Just kagiKey -> [WebSearch.webSearchTool kagiKey]
        Nothing -> []
      pdfTools = [Pdf.pdfTool]
      notesTools =
        [ Notes.noteAddTool uid,
          Notes.noteListTool uid,
          Notes.noteDeleteTool uid
        ]
      tools = memoryTools <> searchTools <> pdfTools <> notesTools

  let agentCfg =
        Engine.defaultAgentConfig
          { Engine.agentSystemPrompt = systemPrompt,
            Engine.agentTools = tools,
            Engine.agentMaxIterations = 5,
            Engine.agentGuardrails =
              Engine.defaultGuardrails
                { Engine.guardrailMaxCostCents = 10.0
                }
          }

  result <- Engine.runAgentWithProvider engineCfg provider agentCfg userMessage

  case result of
    Left err -> do
      putText <| "Agent error: " <> err
      sendMessage tgConfig chatId "Sorry, I encountered an error. Please try again."
    Right agentResult -> do
      let response = Engine.resultFinalMessage agentResult
      putText <| "Response text: " <> Text.take 200 response

      _ <- Memory.saveMessage uid chatId Memory.AssistantRole response

      if Text.null response
        then do
          putText "Warning: empty response from agent"
          sendMessage tgConfig chatId "hmm, i don't have a response for that"
        else sendMessage tgConfig chatId response

      checkAndSummarize provider uid chatId

      putText
        <| "Responded to "
        <> userName
        <> " (cost: "
        <> tshow (Engine.resultTotalCost agentResult)
        <> " cents)"

maxConversationTokens :: Int
maxConversationTokens = 4000

summarizationThreshold :: Int
summarizationThreshold = 3000

checkAndSummarize :: Provider.Provider -> Text -> Int -> IO ()
checkAndSummarize provider uid chatId = do
  (_, currentTokens) <- Memory.getConversationContext uid chatId maxConversationTokens
  when (currentTokens > summarizationThreshold) <| do
    putText <| "Context at " <> tshow currentTokens <> " tokens, summarizing..."
    recentMsgs <- Memory.getRecentMessages uid chatId 50
    let conversationText =
          Text.unlines
            [ (if Memory.cmRole m == Memory.UserRole then "User: " else "Assistant: ") <> Memory.cmContent m
              | m <- reverse recentMsgs
            ]
    summaryResult <-
      Provider.chat
        provider
        []
        [ Provider.Message Provider.System "You are a conversation summarizer. Summarize the key points, decisions, and context from this conversation in 2-3 paragraphs. Focus on information that would be useful for continuing the conversation later." Nothing Nothing,
          Provider.Message Provider.User ("Summarize this conversation:\n\n" <> conversationText) Nothing Nothing
        ]
    case summaryResult of
      Left err -> putText <| "Summarization failed: " <> err
      Right summaryMsg -> do
        let summary = Provider.msgContent summaryMsg
        _ <- Memory.summarizeAndArchive uid chatId summary
        putText "Conversation summarized and archived"

-- | Check if Ollama is running and has the embedding model.
-- Returns Right () if ready, Left error message otherwise.
checkOllama :: IO (Either Text ())
checkOllama = do
  ollamaUrl <- fromMaybe "http://localhost:11434" </ lookupEnv "OLLAMA_URL"
  let url = ollamaUrl <> "/api/tags"
  result <-
    try <| do
      req <- HTTP.parseRequest url
      HTTP.httpLBS req
  case result of
    Left (e :: SomeException) ->
      pure (Left ("Ollama not running: " <> 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 "models" obj of
            Just (Aeson.Array models) ->
              let names = [n | Aeson.Object m <- toList models, Just (Aeson.String n) <- [KeyMap.lookup "name" m]]
                  hasNomic = any ("nomic-embed-text" `Text.isInfixOf`) names
               in if hasNomic
                    then pure (Right ())
                    else pure (Left "nomic-embed-text model not found")
            _ -> pure (Left "Invalid Ollama response")
          _ -> pure (Left "Failed to parse Ollama response")
        else pure (Left ("Ollama HTTP error: " <> tshow status))

-- | Pull the embedding model from Ollama.
pullEmbeddingModel :: IO (Either Text ())
pullEmbeddingModel = do
  ollamaUrl <- fromMaybe "http://localhost:11434" </ lookupEnv "OLLAMA_URL"
  let url = ollamaUrl <> "/api/pull"
  putText "Pulling nomic-embed-text model (this may take a few minutes)..."
  req0 <- HTTP.parseRequest url
  let body = Aeson.object ["name" .= ("nomic-embed-text" :: Text)]
      req =
        HTTP.setRequestMethod "POST"
          <| HTTP.setRequestHeader "Content-Type" ["application/json"]
          <| HTTP.setRequestBodyLBS (Aeson.encode body)
          <| HTTP.setRequestResponseTimeout (HTTPClient.responseTimeoutMicro (600 * 1000000))
          <| req0
  result <- try (HTTP.httpLBS req)
  case result of
    Left (e :: SomeException) ->
      pure (Left ("Failed to pull model: " <> tshow e))
    Right response -> do
      let status = HTTP.getResponseStatusCode response
      if status >= 200 && status < 300
        then do
          putText "nomic-embed-text model ready"
          pure (Right ())
        else pure (Left ("Pull failed: HTTP " <> tshow status))

-- | Ensure Ollama is running and has the embedding model.
-- Pulls the model if missing, exits if Ollama is not running.
ensureOllama :: IO ()
ensureOllama = do
  checkResult <- checkOllama
  case checkResult of
    Right () -> putText "Ollama ready with nomic-embed-text"
    Left err
      | "not running" `Text.isInfixOf` err -> do
          putText <| "Error: " <> err
          putText "Please start Ollama: ollama serve"
          exitFailure
      | "not found" `Text.isInfixOf` err -> do
          putText "nomic-embed-text model not found, pulling..."
          pullResult <- pullEmbeddingModel
          case pullResult of
            Right () -> pure ()
            Left pullErr -> do
              putText <| "Error: " <> pullErr
              exitFailure
      | otherwise -> do
          putText <| "Ollama error: " <> err
          exitFailure

-- | Start the Telegram bot from environment or provided token.
startBot :: Maybe Text -> IO ()
startBot maybeToken = do
  token <- case maybeToken of
    Just t -> pure t
    Nothing -> do
      envToken <- lookupEnv "TELEGRAM_BOT_TOKEN"
      case envToken of
        Just t -> pure (Text.pack t)
        Nothing -> do
          putText "Error: TELEGRAM_BOT_TOKEN not set and no --token provided"
          exitFailure

  ensureOllama

  allowedIds <- loadAllowedUserIds
  kagiKey <- fmap Text.pack </ lookupEnv "KAGI_API_KEY"

  apiKey <- lookupEnv "OPENROUTER_API_KEY"
  case apiKey of
    Nothing -> do
      putText "Error: OPENROUTER_API_KEY not set"
      exitFailure
    Just key -> do
      let tgConfig = defaultTelegramConfig token allowedIds kagiKey
          provider = Provider.defaultOpenRouter (Text.pack key) "anthropic/claude-sonnet-4"
      putText <| "Allowed user IDs: " <> tshow allowedIds
      putText <| "Kagi search: " <> if isJust kagiKey then "enabled" else "disabled"
      runTelegramBot tgConfig provider

-- | Load allowed user IDs from environment variable.
-- Format: comma-separated integers, e.g. "123,456,789"
-- Empty list means allow all users.
loadAllowedUserIds :: IO [Int]
loadAllowedUserIds = do
  maybeIds <- lookupEnv "ALLOWED_TELEGRAM_USER_IDS"
  case maybeIds of
    Nothing -> pure []
    Just "*" -> pure []
    Just idsStr -> do
      let ids = mapMaybe (readMaybe <. Text.unpack <. Text.strip) (Text.splitOn "," (Text.pack idsStr))
      pure ids