path: root/Omni/Agent/Telegram/Types.hs

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

-- | Telegram Bot Types - Data types and JSON parsing for Telegram API.
--
-- : out omni-agent-telegram-types
-- : dep aeson
module Omni.Agent.Telegram.Types
  ( -- * Configuration
    TelegramConfig (..),
    defaultTelegramConfig,
    isUserAllowed,

    -- * Message Types
    TelegramMessage (..),
    TelegramUpdate (..),
    TelegramDocument (..),
    TelegramPhoto (..),
    TelegramVoice (..),
    TelegramReplyMessage (..),
    BotAddedToGroup (..),
    ChatType (..),

    -- * Parsing
    parseUpdate,
    parseBotAddedToGroup,
    parseDocument,
    parseLargestPhoto,
    parsePhotoSize,
    parseVoice,
    parseReplyMessage,

    -- * Utilities
    isPdf,
    isSupportedVoiceFormat,
    isGroupChat,
    shouldRespondInGroup,

    -- * 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.Text as Text
import qualified Omni.Test as Test

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

test :: Test.Tree
test =
  Test.group
    "Omni.Agent.Telegram.Types"
    [ 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",
                  tgOpenRouterApiKey = "or-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 "key"
        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 "key"
        isUserAllowed cfg 12345 Test.@=? True,
      Test.unit "TelegramMessage JSON roundtrip" <| do
        let msg =
              TelegramMessage
                { tmUpdateId = 123,
                  tmChatId = 456,
                  tmChatType = Private,
                  tmUserId = 789,
                  tmUserFirstName = "Test",
                  tmUserLastName = Just "User",
                  tmText = "Hello bot",
                  tmDocument = Nothing,
                  tmPhoto = Nothing,
                  tmVoice = Nothing,
                  tmReplyTo = Nothing,
                  tmThreadId = 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 "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,
      Test.unit "isSupportedVoiceFormat accepts ogg" <| do
        let voice = TelegramVoice "id" 10 (Just "audio/ogg") Nothing
        isSupportedVoiceFormat voice Test.@=? True,
      Test.unit "isSupportedVoiceFormat accepts opus" <| do
        let voice = TelegramVoice "id" 10 (Just "audio/opus") Nothing
        isSupportedVoiceFormat voice Test.@=? True,
      Test.unit "isSupportedVoiceFormat defaults to True for unknown" <| do
        let voice = TelegramVoice "id" 10 Nothing Nothing
        isSupportedVoiceFormat voice Test.@=? True
    ]

data TelegramConfig = TelegramConfig
  { tgBotToken :: Text,
    tgPollingTimeout :: Int,
    tgApiBaseUrl :: Text,
    tgAllowedUserIds :: [Int],
    tgKagiApiKey :: Maybe Text,
    tgOpenRouterApiKey :: 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,
        "openrouter_api_key" .= tgOpenRouterApiKey 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")
        <*> (v .: "openrouter_api_key")

defaultTelegramConfig :: Text -> [Int] -> Maybe Text -> Text -> TelegramConfig
defaultTelegramConfig token allowedIds kagiKey openRouterKey =
  TelegramConfig
    { tgBotToken = token,
      tgPollingTimeout = 30,
      tgApiBaseUrl = "https://api.telegram.org",
      tgAllowedUserIds = allowedIds,
      tgKagiApiKey = kagiKey,
      tgOpenRouterApiKey = openRouterKey
    }

isUserAllowed :: TelegramConfig -> Int -> Bool
isUserAllowed cfg usrId =
  null (tgAllowedUserIds cfg) || usrId `elem` tgAllowedUserIds cfg

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")

data TelegramPhoto = TelegramPhoto
  { tpFileId :: Text,
    tpWidth :: Int,
    tpHeight :: Int,
    tpFileSize :: Maybe Int
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON TelegramPhoto where
  toJSON p =
    Aeson.object
      [ "file_id" .= tpFileId p,
        "width" .= tpWidth p,
        "height" .= tpHeight p,
        "file_size" .= tpFileSize p
      ]

instance Aeson.FromJSON TelegramPhoto where
  parseJSON =
    Aeson.withObject "TelegramPhoto" <| \v ->
      (TelegramPhoto </ (v .: "file_id"))
        <*> (v .: "width")
        <*> (v .: "height")
        <*> (v .:? "file_size")

data TelegramVoice = TelegramVoice
  { tvFileId :: Text,
    tvDuration :: Int,
    tvMimeType :: Maybe Text,
    tvFileSize :: Maybe Int
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON TelegramVoice where
  toJSON v =
    Aeson.object
      [ "file_id" .= tvFileId v,
        "duration" .= tvDuration v,
        "mime_type" .= tvMimeType v,
        "file_size" .= tvFileSize v
      ]

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

data TelegramReplyMessage = TelegramReplyMessage
  { trMessageId :: Int,
    trFromFirstName :: Maybe Text,
    trFromLastName :: Maybe Text,
    trText :: Text
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON TelegramReplyMessage where
  toJSON r =
    Aeson.object
      [ "message_id" .= trMessageId r,
        "from_first_name" .= trFromFirstName r,
        "from_last_name" .= trFromLastName r,
        "text" .= trText r
      ]

instance Aeson.FromJSON TelegramReplyMessage where
  parseJSON =
    Aeson.withObject "TelegramReplyMessage" <| \v ->
      (TelegramReplyMessage </ (v .: "message_id"))
        <*> (v .:? "from_first_name")
        <*> (v .:? "from_last_name")
        <*> (v .:? "text" .!= "")

data BotAddedToGroup = BotAddedToGroup
  { bagUpdateId :: Int,
    bagChatId :: Int,
    bagAddedByUserId :: Int,
    bagAddedByFirstName :: Text
  }
  deriving (Show, Eq, Generic)

data ChatType = Private | Group | Supergroup | Channel
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON ChatType where
  toJSON Private = Aeson.String "private"
  toJSON Group = Aeson.String "group"
  toJSON Supergroup = Aeson.String "supergroup"
  toJSON Channel = Aeson.String "channel"

instance Aeson.FromJSON ChatType where
  parseJSON = Aeson.withText "ChatType" parseChatType
    where
      parseChatType "private" = pure Private
      parseChatType "group" = pure Group
      parseChatType "supergroup" = pure Supergroup
      parseChatType "channel" = pure Channel
      parseChatType _ = pure Private

data TelegramMessage = TelegramMessage
  { tmUpdateId :: Int,
    tmChatId :: Int,
    tmChatType :: ChatType,
    tmThreadId :: Maybe Int,
    tmUserId :: Int,
    tmUserFirstName :: Text,
    tmUserLastName :: Maybe Text,
    tmText :: Text,
    tmDocument :: Maybe TelegramDocument,
    tmPhoto :: Maybe TelegramPhoto,
    tmVoice :: Maybe TelegramVoice,
    tmReplyTo :: Maybe TelegramReplyMessage
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON TelegramMessage where
  toJSON m =
    Aeson.object
      [ "update_id" .= tmUpdateId m,
        "chat_id" .= tmChatId m,
        "chat_type" .= tmChatType m,
        "thread_id" .= tmThreadId m,
        "user_id" .= tmUserId m,
        "user_first_name" .= tmUserFirstName m,
        "user_last_name" .= tmUserLastName m,
        "text" .= tmText m,
        "document" .= tmDocument m,
        "photo" .= tmPhoto m,
        "voice" .= tmVoice m,
        "reply_to" .= tmReplyTo m
      ]

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

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")

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
  let chatType = case KeyMap.lookup "type" chatObj of
        Just (Aeson.String "private") -> Private
        Just (Aeson.String "group") -> Group
        Just (Aeson.String "supergroup") -> Supergroup
        Just (Aeson.String "channel") -> Channel
        _ -> Private
  let threadId = case KeyMap.lookup "message_thread_id" msgObj 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 photo = case KeyMap.lookup "photo" msgObj of
        Just (Aeson.Array photos) -> parseLargestPhoto (toList photos)
        _ -> Nothing
  let voice = case KeyMap.lookup "voice" msgObj of
        Just (Aeson.Object voiceObj) -> parseVoice voiceObj
        _ -> Nothing
  let replyTo = case KeyMap.lookup "reply_to_message" msgObj of
        Just (Aeson.Object replyObj) -> parseReplyMessage replyObj
        _ -> Nothing
  let hasContent = not (Text.null text) || not (Text.null caption) || isJust document || isJust photo || isJust voice
  guard hasContent
  pure
    TelegramMessage
      { tmUpdateId = updateId,
        tmChatId = chatId,
        tmChatType = chatType,
        tmThreadId = threadId,
        tmUserId = userId,
        tmUserFirstName = firstName,
        tmUserLastName = lastName,
        tmText = if Text.null text then caption else text,
        tmDocument = document,
        tmPhoto = photo,
        tmVoice = voice,
        tmReplyTo = replyTo
      }

parseBotAddedToGroup :: Text -> Aeson.Value -> Maybe BotAddedToGroup
parseBotAddedToGroup botUsername 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
  let chatType = case KeyMap.lookup "type" chatObj of
        Just (Aeson.String t) -> t
        _ -> "private"
  guard (chatType == "group" || chatType == "supergroup")
  Aeson.Object fromObj <- KeyMap.lookup "from" msgObj
  addedByUserId <- case KeyMap.lookup "id" fromObj of
    Just (Aeson.Number n) -> Just (round n)
    _ -> Nothing
  addedByFirstName <- case KeyMap.lookup "first_name" fromObj of
    Just (Aeson.String s) -> Just s
    _ -> Nothing
  Aeson.Array newMembers <- KeyMap.lookup "new_chat_members" msgObj
  let botWasAdded = any (isBotUser botUsername) (toList newMembers)
  guard botWasAdded
  pure
    BotAddedToGroup
      { bagUpdateId = updateId,
        bagChatId = chatId,
        bagAddedByUserId = addedByUserId,
        bagAddedByFirstName = addedByFirstName
      }
  where
    isBotUser :: Text -> Aeson.Value -> Bool
    isBotUser username (Aeson.Object userObj) =
      case KeyMap.lookup "username" userObj of
        Just (Aeson.String u) -> Text.toLower u == Text.toLower username
        _ -> False
    isBotUser _ _ = False

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
      }

parseLargestPhoto :: [Aeson.Value] -> Maybe TelegramPhoto
parseLargestPhoto photos = do
  let parsed = mapMaybe parsePhotoSize photos
  case parsed of
    [] -> Nothing
    ps -> Just (maximumBy (comparing tpWidth) ps)

parsePhotoSize :: Aeson.Value -> Maybe TelegramPhoto
parsePhotoSize val = do
  Aeson.Object obj <- pure val
  fileId <- case KeyMap.lookup "file_id" obj of
    Just (Aeson.String s) -> Just s
    _ -> Nothing
  width <- case KeyMap.lookup "width" obj of
    Just (Aeson.Number n) -> Just (round n)
    _ -> Nothing
  height <- case KeyMap.lookup "height" obj of
    Just (Aeson.Number n) -> Just (round n)
    _ -> Nothing
  let fileSize = case KeyMap.lookup "file_size" obj of
        Just (Aeson.Number n) -> Just (round n)
        _ -> Nothing
  pure
    TelegramPhoto
      { tpFileId = fileId,
        tpWidth = width,
        tpHeight = height,
        tpFileSize = fileSize
      }

parseVoice :: Aeson.Object -> Maybe TelegramVoice
parseVoice obj = do
  fileId <- case KeyMap.lookup "file_id" obj of
    Just (Aeson.String s) -> Just s
    _ -> Nothing
  duration <- case KeyMap.lookup "duration" obj of
    Just (Aeson.Number n) -> Just (round n)
    _ -> Nothing
  let mimeType = case KeyMap.lookup "mime_type" obj of
        Just (Aeson.String s) -> Just s
        _ -> Nothing
      fileSize = case KeyMap.lookup "file_size" obj of
        Just (Aeson.Number n) -> Just (round n)
        _ -> Nothing
  pure
    TelegramVoice
      { tvFileId = fileId,
        tvDuration = duration,
        tvMimeType = mimeType,
        tvFileSize = fileSize
      }

parseReplyMessage :: Aeson.Object -> Maybe TelegramReplyMessage
parseReplyMessage obj = do
  messageId <- case KeyMap.lookup "message_id" obj of
    Just (Aeson.Number n) -> Just (round n)
    _ -> Nothing
  let fromFirstName = case KeyMap.lookup "from" obj of
        Just (Aeson.Object fromObj) -> case KeyMap.lookup "first_name" fromObj of
          Just (Aeson.String s) -> Just s
          _ -> Nothing
        _ -> Nothing
      fromLastName = case KeyMap.lookup "from" obj of
        Just (Aeson.Object fromObj) -> case KeyMap.lookup "last_name" fromObj of
          Just (Aeson.String s) -> Just s
          _ -> Nothing
        _ -> Nothing
      text = case KeyMap.lookup "text" obj of
        Just (Aeson.String s) -> s
        _ -> case KeyMap.lookup "caption" obj of
          Just (Aeson.String s) -> s
          _ -> ""
  pure
    TelegramReplyMessage
      { trMessageId = messageId,
        trFromFirstName = fromFirstName,
        trFromLastName = fromLastName,
        trText = text
      }

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

isSupportedVoiceFormat :: TelegramVoice -> Bool
isSupportedVoiceFormat voice =
  case tvMimeType voice of
    Just "audio/ogg" -> True
    Just "audio/opus" -> True
    Just "audio/x-opus+ogg" -> True
    Nothing -> True
    _ -> False

isGroupChat :: TelegramMessage -> Bool
isGroupChat msg = tmChatType msg `elem` [Group, Supergroup]

shouldRespondInGroup :: Text -> TelegramMessage -> Bool
shouldRespondInGroup botUsername msg
  | not (isGroupChat msg) = True
  | isMentioned = True
  | isReplyToBot = True
  | otherwise = False
  where
    msgText = Text.toLower (tmText msg)
    mention = "@" <> Text.toLower botUsername
    isMentioned = mention `Text.isInfixOf` msgText
    isReplyToBot = isJust (tmReplyTo msg)