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

-- | Email tools for IMAP access via Telegram bot.
--
-- Provides email management for agents:
-- - Check for urgent/time-sensitive emails
-- - Identify emails needing response vs FYI
-- - Auto-unsubscribe from marketing
--
-- Uses HaskellNet for proper IMAP client support.
-- Password retrieved via `pass ben@bensima.com`.
--
-- : out omni-agent-tools-email
-- : dep aeson
-- : dep process
-- : dep regex-applicative
-- : dep http-conduit
-- : dep HaskellNet
-- : dep HaskellNet-SSL
module Omni.Agent.Tools.Email
  ( -- * Tools
    emailCheckTool,
    emailReadTool,
    emailUnsubscribeTool,
    emailArchiveTool,

    -- * All tools
    allEmailTools,

    -- * Direct API
    checkNewEmails,
    readEmail,
    unsubscribeFromEmail,
    archiveEmail,
    getPassword,

    -- * Scheduled Check
    emailCheckLoop,
    performScheduledCheck,

    -- * Testing
    main,
    test,
  )
where

import Alpha
import Data.Aeson ((.=))
import qualified Data.Aeson as Aeson
import qualified Data.Aeson.KeyMap as KeyMap
import qualified Data.ByteString.Char8 as BS8
import qualified Data.List as List
import qualified Data.Text as Text
import Data.Time (NominalDiffTime, UTCTime, addUTCTime, getCurrentTime)
import Data.Time.Format (defaultTimeLocale, formatTime, parseTimeM)
import Data.Time.LocalTime (TimeZone (..), utcToZonedTime)
import qualified Network.HTTP.Simple as HTTP
import qualified Network.HaskellNet.IMAP as IMAP
import Network.HaskellNet.IMAP.Connection (IMAPConnection)
import qualified Network.HaskellNet.IMAP.SSL as IMAPSSL
import qualified Omni.Agent.Engine as Engine
import qualified Omni.Test as Test
import System.Process (readProcessWithExitCode)
import Text.Regex.Applicative (RE, anySym, few, (=~))
import qualified Text.Regex.Applicative as RE

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

test :: Test.Tree
test =
  Test.group
    "Omni.Agent.Tools.Email"
    [ Test.unit "emailCheckTool has correct name" <| do
        Engine.toolName emailCheckTool Test.@=? "email_check",
      Test.unit "emailReadTool has correct name" <| do
        Engine.toolName emailReadTool Test.@=? "email_read",
      Test.unit "emailUnsubscribeTool has correct name" <| do
        Engine.toolName emailUnsubscribeTool Test.@=? "email_unsubscribe",
      Test.unit "emailArchiveTool has correct name" <| do
        Engine.toolName emailArchiveTool Test.@=? "email_archive",
      Test.unit "allEmailTools has 4 tools" <| do
        length allEmailTools Test.@=? 4,
      Test.unit "parseEmailHeaders extracts fields" <| do
        let headers =
              "From: test@example.com\r\n\
              \Subject: Test Subject\r\n\
              \Date: Mon, 1 Jan 2024 12:00:00 +0000\r\n\
              \\r\n"
        case parseEmailHeaders headers of
          Nothing -> Test.assertFailure "Failed to parse headers"
          Just email -> do
            emailFrom email Test.@=? "test@example.com"
            emailSubject email Test.@=? "Test Subject",
      Test.unit "parseUnsubscribeHeader extracts URL" <| do
        let header = "<https://example.com/unsubscribe>, <mailto:unsub@example.com>"
        case parseUnsubscribeUrl header of
          Nothing -> Test.assertFailure "Failed to parse unsubscribe URL"
          Just url -> ("https://example.com" `Text.isPrefixOf` url) Test.@=? True
    ]

imapServer :: String
imapServer = "bensima.com"

imapUser :: String
imapUser = "ben@bensima.com"

getPassword :: IO (Either Text Text)
getPassword = do
  result <- try <| readProcessWithExitCode "pass" ["ben@bensima.com"] ""
  case result of
    Left (e :: SomeException) ->
      pure (Left ("Failed to get password: " <> tshow e))
    Right (exitCode, stdoutStr, stderrStr) ->
      case exitCode of
        ExitSuccess -> pure (Right (Text.strip (Text.pack stdoutStr)))
        ExitFailure code ->
          pure (Left ("pass failed (" <> tshow code <> "): " <> Text.pack stderrStr))

withImapConnection :: (IMAPConnection -> IO a) -> IO (Either Text a)
withImapConnection action = do
  pwResult <- getPassword
  case pwResult of
    Left err -> pure (Left err)
    Right pw -> do
      result <-
        try <| do
          conn <- IMAPSSL.connectIMAPSSL imapServer
          IMAP.login conn imapUser (Text.unpack pw)
          r <- action conn
          IMAP.logout conn
          pure r
      case result of
        Left (e :: SomeException) -> pure (Left ("IMAP error: " <> tshow e))
        Right r -> pure (Right r)

data EmailSummary = EmailSummary
  { emailUid :: Int,
    emailFrom :: Text,
    emailSubject :: Text,
    emailDate :: Text,
    emailUnsubscribe :: Maybe Text
  }
  deriving (Show, Generic)

instance Aeson.ToJSON EmailSummary where
  toJSON e =
    Aeson.object
      [ "uid" .= emailUid e,
        "from" .= emailFrom e,
        "subject" .= emailSubject e,
        "date" .= formatDateAsEst (emailDate e),
        "has_unsubscribe" .= isJust (emailUnsubscribe e)
      ]

estTimezone :: TimeZone
estTimezone = TimeZone (-300) False "EST"

formatDateAsEst :: Text -> Text
formatDateAsEst dateStr =
  case parseEmailDate dateStr of
    Nothing -> dateStr
    Just utcTime ->
      let zonedTime = utcToZonedTime estTimezone utcTime
       in Text.pack (formatTime defaultTimeLocale "%a %b %d %H:%M EST" zonedTime)

parseEmailHeaders :: Text -> Maybe EmailSummary
parseEmailHeaders raw = do
  let headerLines = Text.lines raw
      fromLine = findHeader "From:" headerLines
      subjectLine = findHeader "Subject:" headerLines
      dateLine = findHeader "Date:" headerLines
      unsubLine = findHeader "List-Unsubscribe:" headerLines
  fromVal <- fromLine
  subject <- subjectLine
  dateVal <- dateLine
  pure
    EmailSummary
      { emailUid = 0,
        emailFrom = Text.strip (Text.drop 5 fromVal),
        emailSubject = Text.strip (Text.drop 8 subject),
        emailDate = Text.strip (Text.drop 5 dateVal),
        emailUnsubscribe = (parseUnsubscribeUrl <. Text.drop 16) =<< unsubLine
      }
  where
    findHeader :: Text -> [Text] -> Maybe Text
    findHeader prefix = List.find (prefix `Text.isPrefixOf`)

parseUnsubscribeUrl :: Text -> Maybe Text
parseUnsubscribeUrl header =
  let text = Text.unpack header
   in case text =~ urlInBrackets of
        Just url | "http" `List.isPrefixOf` url -> Just (Text.pack url)
        _ -> Nothing
  where
    urlInBrackets :: RE Char String
    urlInBrackets = few anySym *> RE.sym '<' *> few anySym <* RE.sym '>'

checkNewEmails :: Maybe Int -> Maybe Int -> IO (Either Text [EmailSummary])
checkNewEmails maybeLimit maybeHours = do
  withImapConnection <| \conn -> do
    IMAP.select conn "INBOX"
    uids <- IMAP.search conn [IMAP.UNFLAG IMAP.Seen]
    let limit = fromMaybe 20 maybeLimit
        recentUids = take limit (reverse (map fromIntegral uids))
    if null recentUids
      then pure []
      else do
        emails <-
          forM recentUids <| \uid -> do
            headerBytes <- IMAP.fetchHeader conn (fromIntegral uid)
            let headerText = Text.pack (BS8.unpack headerBytes)
            pure (parseEmailHeaders headerText, uid)
        let parsed =
              [ e {emailUid = uid}
                | (Just e, uid) <- emails
              ]
        case maybeHours of
          Nothing -> pure parsed
          Just hours -> do
            now <- getCurrentTime
            let cutoff = addUTCTime (negate (fromIntegral hours * 3600 :: NominalDiffTime)) now
            pure (filter (isAfterCutoff cutoff) parsed)

isAfterCutoff :: UTCTime -> EmailSummary -> Bool
isAfterCutoff cutoff email =
  case parseEmailDate (emailDate email) of
    Nothing -> False
    Just emailTime -> emailTime >= cutoff

parseEmailDate :: Text -> Maybe UTCTime
parseEmailDate dateStr =
  let cleaned = stripParenTz (Text.strip dateStr)
      formats =
        [ "%a, %d %b %Y %H:%M:%S %z",
          "%a, %d %b %Y %H:%M:%S %Z",
          "%d %b %Y %H:%M:%S %z",
          "%a, %d %b %Y %H:%M %z",
          "%a, %d %b %Y %H:%M:%S %z (%Z)"
        ]
      tryParse [] = Nothing
      tryParse (fmt : rest) =
        case parseTimeM True defaultTimeLocale fmt (Text.unpack cleaned) of
          Just t -> Just t
          Nothing -> tryParse rest
   in tryParse formats

stripParenTz :: Text -> Text
stripParenTz t =
  case Text.breakOn " (" t of
    (before, after)
      | Text.null after -> t
      | ")" `Text.isSuffixOf` after -> before
      | otherwise -> t

readEmail :: Int -> IO (Either Text Text)
readEmail uid =
  withImapConnection <| \conn -> do
    IMAP.select conn "INBOX"
    bodyBytes <- IMAP.fetch conn (fromIntegral uid)
    let bodyText = Text.pack (BS8.unpack bodyBytes)
    pure (Text.take 10000 bodyText)

unsubscribeFromEmail :: Int -> IO (Either Text Text)
unsubscribeFromEmail uid = do
  headerResult <-
    withImapConnection <| \conn -> do
      IMAP.select conn "INBOX"
      headerBytes <- IMAP.fetchHeader conn (fromIntegral uid)
      pure (Text.pack (BS8.unpack headerBytes))
  case headerResult of
    Left err -> pure (Left err)
    Right headerText ->
      case extractUnsubscribeUrl headerText of
        Nothing -> pure (Left "No unsubscribe URL found in this email")
        Just url -> do
          clickResult <- clickUnsubscribeLink url
          case clickResult of
            Left err -> pure (Left ("Failed to unsubscribe: " <> err))
            Right () -> do
              _ <- archiveEmail uid
              pure (Right ("Unsubscribed and archived email " <> tshow uid))

extractUnsubscribeUrl :: Text -> Maybe Text
extractUnsubscribeUrl headerText =
  let unsubLine = List.find ("List-Unsubscribe:" `Text.isInfixOf`) (Text.lines headerText)
   in (parseUnsubscribeUrl <. Text.drop 16 <. Text.strip) =<< unsubLine

clickUnsubscribeLink :: Text -> IO (Either Text ())
clickUnsubscribeLink url = do
  result <-
    try <| do
      req <- HTTP.parseRequest (Text.unpack url)
      _ <- HTTP.httpLBS req
      pure ()
  case result of
    Left (e :: SomeException) -> pure (Left (tshow e))
    Right () -> pure (Right ())

archiveEmail :: Int -> IO (Either Text Text)
archiveEmail uid =
  withImapConnection <| \conn -> do
    IMAP.select conn "INBOX"
    IMAP.copy conn (fromIntegral uid) "Archives.2025"
    IMAP.store conn (fromIntegral uid) (IMAP.PlusFlags [IMAP.Deleted])
    _ <- IMAP.expunge conn
    pure ("Archived email " <> tshow uid)

allEmailTools :: [Engine.Tool]
allEmailTools =
  [ emailCheckTool,
    emailReadTool,
    emailUnsubscribeTool,
    emailArchiveTool
  ]

emailCheckTool :: Engine.Tool
emailCheckTool =
  Engine.Tool
    { Engine.toolName = "email_check",
      Engine.toolDescription =
        "Check for new/unread emails. Returns a summary of recent unread emails "
          <> "including sender, subject, date, and whether they have an unsubscribe link. "
          <> "Use this to identify urgent items or emails needing response. "
          <> "Use 'hours' to filter to emails received in the last N hours (e.g., hours=6 for last 6 hours).",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "limit"
                    .= Aeson.object
                      [ "type" .= ("integer" :: Text),
                        "description" .= ("Max emails to return (default: 20)" :: Text)
                      ],
                  "hours"
                    .= Aeson.object
                      [ "type" .= ("integer" :: Text),
                        "description" .= ("Only return emails from the last N hours (e.g., 6 for last 6 hours)" :: Text)
                      ]
                ],
            "required" .= ([] :: [Text])
          ],
      Engine.toolExecute = executeEmailCheck
    }

executeEmailCheck :: Aeson.Value -> IO Aeson.Value
executeEmailCheck v = do
  let (limit, hours) = case v of
        Aeson.Object obj ->
          let l = case KeyMap.lookup "limit" obj of
                Just (Aeson.Number n) -> Just (round n :: Int)
                _ -> Nothing
              h = case KeyMap.lookup "hours" obj of
                Just (Aeson.Number n) -> Just (round n :: Int)
                _ -> Nothing
           in (l, h)
        _ -> (Nothing, Nothing)
  result <- checkNewEmails limit hours
  case result of
    Left err -> pure (Aeson.object ["error" .= err])
    Right emails ->
      pure
        ( Aeson.object
            [ "success" .= True,
              "count" .= length emails,
              "emails" .= emails
            ]
        )

emailReadTool :: Engine.Tool
emailReadTool =
  Engine.Tool
    { Engine.toolName = "email_read",
      Engine.toolDescription =
        "Read the full content of an email by its UID. "
          <> "Use after email_check to read emails that seem important or need a response.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "uid"
                    .= Aeson.object
                      [ "type" .= ("integer" :: Text),
                        "description" .= ("Email UID from email_check" :: Text)
                      ]
                ],
            "required" .= (["uid"] :: [Text])
          ],
      Engine.toolExecute = executeEmailRead
    }

executeEmailRead :: Aeson.Value -> IO Aeson.Value
executeEmailRead v = do
  let uidM = case v of
        Aeson.Object obj -> case KeyMap.lookup "uid" obj of
          Just (Aeson.Number n) -> Just (round n :: Int)
          _ -> Nothing
        _ -> Nothing
  case uidM of
    Nothing -> pure (Aeson.object ["error" .= ("Missing uid parameter" :: Text)])
    Just uid -> do
      result <- readEmail uid
      case result of
        Left err -> pure (Aeson.object ["error" .= err])
        Right body ->
          pure
            ( Aeson.object
                [ "success" .= True,
                  "uid" .= uid,
                  "body" .= body
                ]
            )

emailUnsubscribeTool :: Engine.Tool
emailUnsubscribeTool =
  Engine.Tool
    { Engine.toolName = "email_unsubscribe",
      Engine.toolDescription =
        "Unsubscribe from a mailing list by clicking the List-Unsubscribe link. "
          <> "Use for marketing/newsletter emails. Automatically archives the email after unsubscribing.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "uid"
                    .= Aeson.object
                      [ "type" .= ("integer" :: Text),
                        "description" .= ("Email UID to unsubscribe from" :: Text)
                      ]
                ],
            "required" .= (["uid"] :: [Text])
          ],
      Engine.toolExecute = executeEmailUnsubscribe
    }

executeEmailUnsubscribe :: Aeson.Value -> IO Aeson.Value
executeEmailUnsubscribe v = do
  let uidM = case v of
        Aeson.Object obj -> case KeyMap.lookup "uid" obj of
          Just (Aeson.Number n) -> Just (round n :: Int)
          _ -> Nothing
        _ -> Nothing
  case uidM of
    Nothing -> pure (Aeson.object ["error" .= ("Missing uid parameter" :: Text)])
    Just uid -> do
      result <- unsubscribeFromEmail uid
      case result of
        Left err -> pure (Aeson.object ["error" .= err])
        Right msg ->
          pure
            ( Aeson.object
                [ "success" .= True,
                  "message" .= msg
                ]
            )

emailArchiveTool :: Engine.Tool
emailArchiveTool =
  Engine.Tool
    { Engine.toolName = "email_archive",
      Engine.toolDescription =
        "Archive an email (move to Archives.2025 folder). "
          <> "Use for emails that don't need a response and are just FYI.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "uid"
                    .= Aeson.object
                      [ "type" .= ("integer" :: Text),
                        "description" .= ("Email UID to archive" :: Text)
                      ]
                ],
            "required" .= (["uid"] :: [Text])
          ],
      Engine.toolExecute = executeEmailArchive
    }

executeEmailArchive :: Aeson.Value -> IO Aeson.Value
executeEmailArchive v = do
  let uidM = case v of
        Aeson.Object obj -> case KeyMap.lookup "uid" obj of
          Just (Aeson.Number n) -> Just (round n :: Int)
          _ -> Nothing
        _ -> Nothing
  case uidM of
    Nothing -> pure (Aeson.object ["error" .= ("Missing uid parameter" :: Text)])
    Just uid -> do
      result <- archiveEmail uid
      case result of
        Left err -> pure (Aeson.object ["error" .= err])
        Right msg ->
          pure
            ( Aeson.object
                [ "success" .= True,
                  "message" .= msg
                ]
            )

emailCheckLoop :: (Int -> Maybe Int -> Text -> IO (Maybe Int)) -> Int -> IO ()
emailCheckLoop sendFn chatId =
  forever <| do
    let sixHours = 6 * 60 * 60 * 1000000
    threadDelay sixHours
    performScheduledCheck sendFn chatId

performScheduledCheck :: (Int -> Maybe Int -> Text -> IO (Maybe Int)) -> Int -> IO ()
performScheduledCheck sendFn chatId = do
  putText "Running scheduled email check..."
  result <- checkNewEmails (Just 50) (Just 6)
  case result of
    Left err -> putText ("Email check failed: " <> err)
    Right emails -> do
      let urgent = filter isUrgent emails
          needsResponse = filter needsResponsePred emails
          marketing = filter hasUnsubscribe emails
      when (not (null urgent) || not (null needsResponse)) <| do
        let msg = formatEmailSummary urgent needsResponse (length marketing)
        _ <- sendFn chatId Nothing msg
        pure ()
  where
    isUrgent :: EmailSummary -> Bool
    isUrgent email =
      let subj = Text.toLower (emailSubject email)
       in "urgent"
            `Text.isInfixOf` subj
            || "asap"
            `Text.isInfixOf` subj
            || "important"
            `Text.isInfixOf` subj
            || "action required"
            `Text.isInfixOf` subj

    needsResponsePred :: EmailSummary -> Bool
    needsResponsePred email =
      let sender = Text.toLower (emailFrom email)
          subj = Text.toLower (emailSubject email)
       in not (hasUnsubscribe email)
            && not (isUrgent email)
            && not ("noreply" `Text.isInfixOf` sender)
            && not ("no-reply" `Text.isInfixOf` sender)
            && ("?" `Text.isInfixOf` subj || "reply" `Text.isInfixOf` subj || "response" `Text.isInfixOf` subj)

    hasUnsubscribe :: EmailSummary -> Bool
    hasUnsubscribe = isJust <. emailUnsubscribe

    formatEmailSummary :: [EmailSummary] -> [EmailSummary] -> Int -> Text
    formatEmailSummary urgent needs marketingCount =
      Text.unlines
        <| ["📧 *email check*", ""]
        <> (if null urgent then [] else ["*urgent:*"] <> map formatOne urgent <> [""])
        <> (if null needs then [] else ["*may need response:*"] <> map formatOne needs <> [""])
        <> [tshow marketingCount <> " marketing emails (use email_check to review)"]

    formatOne :: EmailSummary -> Text
    formatOne e =
      "• " <> emailSubject e <> " (from: " <> emailFrom e <> ", uid: " <> tshow (emailUid e) <> ")"