{-# 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 (..),

    -- * Telegram API
    getUpdates,
    sendMessage,

    -- * Bot Loop
    runTelegramBot,
    handleMessage,
    startBot,

    -- * 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.Text as Text
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.Test as Test
import System.Environment (lookupEnv)

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"
                }
        case Aeson.decode (Aeson.encode cfg) of
          Nothing -> Test.assertFailure "Failed to decode TelegramConfig"
          Just decoded -> tgBotToken decoded Test.@=? "test-token",
      Test.unit "TelegramMessage JSON roundtrip" <| do
        let msg =
              TelegramMessage
                { tmUpdateId = 123,
                  tmChatId = 456,
                  tmUserId = 789,
                  tmUserFirstName = "Test",
                  tmUserLastName = Just "User",
                  tmText = "Hello bot"
                }
        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"
    ]

-- | Telegram bot configuration.
data TelegramConfig = TelegramConfig
  { tgBotToken :: Text,
    tgPollingTimeout :: Int,
    tgApiBaseUrl :: 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
      ]

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

-- | Default Telegram configuration (requires token from env).
defaultTelegramConfig :: Text -> TelegramConfig
defaultTelegramConfig token =
  TelegramConfig
    { tgBotToken = token,
      tgPollingTimeout = 30,
      tgApiBaseUrl = "https://api.telegram.org"
    }

-- | A parsed Telegram message from a user.
data TelegramMessage = TelegramMessage
  { tmUpdateId :: Int,
    tmChatId :: Int,
    tmUserId :: Int,
    tmUserFirstName :: Text,
    tmUserLastName :: Maybe Text,
    tmText :: Text
  }
  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
      ]

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

-- | 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.
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
  text <- case KeyMap.lookup "text" msgObj of
    Just (Aeson.String s) -> Just s
    _ -> Nothing
  pure
    TelegramMessage
      { tmUpdateId = updateId,
        tmChatId = chatId,
        tmUserId = userId,
        tmUserFirstName = firstName,
        tmUserLastName = lastName,
        tmText = text
      }

-- | 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])
          ]
      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) -> 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 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,
            "parse_mode" .= ("Markdown" :: 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
      unless (status >= 200 && status < 300)
        <| putText
        <| "Send message failed: "
        <> tshow status

-- | System prompt for the Telegram bot agent.
telegramSystemPrompt :: Text
telegramSystemPrompt =
  Text.unlines
    [ "You are a helpful family assistant on Telegram. You help with questions,",
      "remember important information about family members, and provide friendly assistance.",
      "",
      "When you learn something important about the user (preferences, facts about them,",
      "their interests, family details), use the 'remember' tool to store it for future reference.",
      "",
      "Be concise in responses - Telegram is a chat interface, not a document.",
      "Keep responses under 200 words unless the user asks for detail.",
      "Be friendly and helpful. This is a family bot, keep content appropriate.",
      "",
      "If the user asks something you don't know, be honest about it.",
      "You can use the 'recall' tool to search your memory for relevant information."
    ]

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

  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)

  user <- Memory.getOrCreateUserByTelegramId (tmUserId msg) userName

  memories <- Memory.recallMemories (Memory.userId user) (tmText msg) 5
  let memoryContext = Memory.formatMemoriesForPrompt memories

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

  let tools =
        [ Memory.rememberTool (Memory.userId user),
          Memory.recallTool (Memory.userId user)
        ]

  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 (tmText msg)

  case result of
    Left err -> do
      putText <| "Agent error: " <> err
      sendMessage tgConfig (tmChatId msg) "Sorry, I encountered an error. Please try again."
    Right agentResult -> do
      let response = Engine.resultFinalMessage agentResult
      sendMessage tgConfig (tmChatId msg) response
      putText
        <| "Responded to "
        <> userName
        <> " (cost: "
        <> tshow (Engine.resultTotalCost agentResult)
        <> " cents)"

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

  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
          provider = Provider.defaultOpenRouter (Text.pack key) "anthropic/claude-sonnet-4"
      runTelegramBot tgConfig provider