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

-- | Tool for Ava to read her own audit logs and subagent traces.
--
-- Enables self-diagnosis: "Let me check my logs for that subagent run..."
--
-- : out omni-agent-tools-avalogs
-- : dep aeson
-- : dep time
module Omni.Agent.Tools.AvaLogs
  ( readAvaLogsTool,
    searchChatHistoryTool,
    main,
  )
where

import Alpha
import Data.Aeson ((.=))
import qualified Data.Aeson as Aeson
import qualified Data.Aeson.Key as Key
import qualified Data.Aeson.KeyMap as KeyMap
import qualified Data.Text as Text
import qualified Data.Time as Time
import qualified Omni.Agent.AuditLog as AuditLog
import qualified Omni.Agent.Engine as Engine

main :: IO ()
main = putText "Omni.Agent.Tools.AvaLogs - no standalone execution"

readAvaLogsTool :: Engine.Tool
readAvaLogsTool =
  Engine.Tool
    { Engine.toolName = "read_ava_logs",
      Engine.toolDescription =
        "Read Ava's audit logs or subagent traces for self-diagnosis. "
          <> "Use to review past conversations, inspect subagent runs, or debug issues. "
          <> "Pass subagent_id to view a specific subagent's trace, or last_n for recent Ava logs.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "subagent_id"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("Subagent ID to view (e.g. 'abc123')" :: Text)
                      ],
                  "last_n"
                    .= Aeson.object
                      [ "type" .= ("integer" :: Text),
                        "description" .= ("Number of recent log entries to return (default: 20)" :: Text)
                      ]
                ]
          ],
      Engine.toolExecute = executeReadLogs
    }

executeReadLogs :: Aeson.Value -> IO Aeson.Value
executeReadLogs v = do
  let maybeSubagentId = case v of
        Aeson.Object obj -> case KeyMap.lookup "subagent_id" obj of
          Just (Aeson.String sid) -> Just sid
          _ -> Nothing
        _ -> Nothing

  let lastN = case v of
        Aeson.Object obj -> case KeyMap.lookup "last_n" obj of
          Just (Aeson.Number n) -> round n
          _ -> 20
        _ -> 20

  case maybeSubagentId of
    Just sid -> do
      let subagentId = AuditLog.SubagentId sid
      entries <- AuditLog.readSubagentLogs subagentId
      pure
        <| Aeson.object
          [ "subagent_id" .= sid,
            "entry_count" .= length entries,
            "entries" .= map formatEntry entries
          ]
    Nothing -> do
      entries <- AuditLog.getRecentAvaLogs lastN
      today <- Time.utctDay </ Time.getCurrentTime
      pure
        <| Aeson.object
          [ "date" .= Time.formatTime Time.defaultTimeLocale "%Y-%m-%d" today,
            "entry_count" .= length entries,
            "entries" .= map formatEntry entries
          ]

formatEntry :: AuditLog.AuditLogEntry -> Aeson.Value
formatEntry entry =
  Aeson.object
    [ "timestamp" .= Time.formatTime Time.defaultTimeLocale "%H:%M:%S" (AuditLog.logTimestamp entry),
      "event_type" .= tshow (AuditLog.logEventType entry),
      "agent_id" .= AuditLog.unAgentId (AuditLog.logAgentId entry),
      "content" .= summarizeContent (AuditLog.logContent entry)
    ]

summarizeContent :: Aeson.Value -> Text
summarizeContent (Aeson.String t) = Text.take 200 t
summarizeContent (Aeson.Object obj) =
  let keys = KeyMap.keys obj
   in "object with keys: " <> Text.intercalate ", " (map Key.toText keys)
summarizeContent (Aeson.Array arr) = "array with " <> tshow (length arr) <> " items"
summarizeContent Aeson.Null = "null"
summarizeContent (Aeson.Bool b) = if b then "true" else "false"
summarizeContent (Aeson.Number n) = tshow n

searchChatHistoryTool :: Engine.Tool
searchChatHistoryTool =
  Engine.Tool
    { Engine.toolName = "search_chat_history",
      Engine.toolDescription =
        "Search your conversation history for specific content. "
          <> "Use this to find what was said in past conversations, recall context, "
          <> "or find when something was discussed. Searches message content.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "query"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("Text to search for in chat history" :: Text)
                      ],
                  "days_back"
                    .= Aeson.object
                      [ "type" .= ("integer" :: Text),
                        "description" .= ("How many days back to search (default: 7)" :: Text)
                      ],
                  "max_results"
                    .= Aeson.object
                      [ "type" .= ("integer" :: Text),
                        "description" .= ("Maximum results to return (default: 20)" :: Text)
                      ]
                ],
            "required" .= (["query"] :: [Text])
          ],
      Engine.toolExecute = executeSearchHistory
    }

executeSearchHistory :: Aeson.Value -> IO Aeson.Value
executeSearchHistory v = do
  let query = case v of
        Aeson.Object obj -> case KeyMap.lookup "query" obj of
          Just (Aeson.String q) -> q
          _ -> ""
        _ -> ""

  let daysBack = case v of
        Aeson.Object obj -> case KeyMap.lookup "days_back" obj of
          Just (Aeson.Number n) -> round n
          _ -> 7
        _ -> 7

  let maxResults = case v of
        Aeson.Object obj -> case KeyMap.lookup "max_results" obj of
          Just (Aeson.Number n) -> round n
          _ -> 20
        _ -> 20

  if Text.null query
    then pure <| Aeson.object ["error" .= ("query is required" :: Text)]
    else do
      today <- Time.utctDay </ Time.getCurrentTime
      let days = [Time.addDays (negate i) today | i <- [0 .. daysBack - 1]]
      allEntries <- concat </ traverse AuditLog.readAvaLogs days
      let matches = filter (matchesQuery query) allEntries
          limited = take maxResults matches
      pure
        <| Aeson.object
          [ "query" .= query,
            "days_searched" .= daysBack,
            "total_matches" .= length matches,
            "results" .= map formatSearchResult limited
          ]

matchesQuery :: Text -> AuditLog.AuditLogEntry -> Bool
matchesQuery query entry =
  let eventType = AuditLog.logEventType entry
      isMessage = eventType == AuditLog.UserMessage || eventType == AuditLog.AssistantMessage
      content = extractContent (AuditLog.logContent entry)
      queryLower = Text.toLower query
   in isMessage && queryLower `Text.isInfixOf` Text.toLower content

extractContent :: Aeson.Value -> Text
extractContent (Aeson.String t) = t
extractContent (Aeson.Object obj) =
  case KeyMap.lookup "text" obj of
    Just (Aeson.String t) -> t
    _ -> case KeyMap.lookup "content" obj of
      Just (Aeson.String t) -> t
      _ -> case KeyMap.lookup "message" obj of
        Just (Aeson.String t) -> t
        _ -> ""
extractContent _ = ""

formatSearchResult :: AuditLog.AuditLogEntry -> Aeson.Value
formatSearchResult entry =
  Aeson.object
    [ "timestamp" .= Time.formatTime Time.defaultTimeLocale "%Y-%m-%d %H:%M:%S" (AuditLog.logTimestamp entry),
      "role" .= roleText (AuditLog.logEventType entry),
      "user" .= AuditLog.logUserId entry,
      "content" .= Text.take 500 (extractContent (AuditLog.logContent entry))
    ]

roleText :: AuditLog.AuditEventType -> Text
roleText AuditLog.UserMessage = "user"
roleText AuditLog.AssistantMessage = "assistant"
roleText _ = "other"