summaryrefslogtreecommitdiff
path: root/Omni/Agent/Telegram.hs
blob: f1c71e6083d7e3036a7714f95df3c7cf13cc37e9 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
{-# 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 Data.Time (getCurrentTime, utcToLocalTime)
import Data.Time.Format (defaultTimeLocale, formatTime)
import Data.Time.LocalTime (getCurrentTimeZone)
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.Calendar as Calendar
import qualified Omni.Agent.Tools.Notes as Notes
import qualified Omni.Agent.Tools.Pdf as Pdf
import qualified Omni.Agent.Tools.Todos as Todos
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.",
      "",
      "## important",
      "",
      "ALWAYS include a text response to the user after using tools. never end your turn with only tool calls."
    ]

-- | 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 (Just userName) 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

  now <- getCurrentTime
  tz <- getCurrentTimeZone
  let localTime = utcToLocalTime tz now
      timeStr = Text.pack (formatTime defaultTimeLocale "%A, %B %d, %Y at %H:%M" localTime)

  let systemPrompt =
        telegramSystemPrompt
          <> "\n\n## Current Date and Time\n"
          <> timeStr
          <> "\n\n## Current User\n"
          <> "You are talking to: "
          <> userName
          <> "\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
        ]
      calendarTools =
        [ Calendar.calendarListTool,
          Calendar.calendarAddTool,
          Calendar.calendarSearchTool
        ]
      todoTools =
        [ Todos.todoAddTool uid,
          Todos.todoListTool uid,
          Todos.todoCompleteTool uid,
          Todos.todoDeleteTool uid
        ]
      tools = memoryTools <> searchTools <> pdfTools <> notesTools <> calendarTools <> todoTools

  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 Nothing 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