{-# 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
import qualified Omni.Agent.Memory as Memory

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 using semantic similarity. "
          <> "Use this to find what was said in past conversations, recall context, "
          <> "or find when something was discussed. Finds semantically related messages.",
      Engine.toolJsonSchema =
        Aeson.object
          [ "type" .= ("object" :: Text),
            "properties"
              .= Aeson.object
                [ "query"
                    .= Aeson.object
                      [ "type" .= ("string" :: Text),
                        "description" .= ("What to search for (semantic search)" :: 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 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
      (_total, indexed) <- Memory.getChatHistoryStats
      if indexed > 0
        then do
          results <- Memory.searchChatHistorySemantic query maxResults
          pure
            <| Aeson.object
              [ "query" .= query,
                "search_type" .= ("semantic" :: Text),
                "indexed_messages" .= indexed,
                "results" .= map formatSemanticResult results
              ]
        else do
          today <- Time.utctDay </ Time.getCurrentTime
          let days = [Time.addDays (negate i) today | i <- [0 .. 6]]
          allEntries <- concat </ traverse AuditLog.readAvaLogs days
          let matches = filter (matchesQuery query) allEntries
              limited = take maxResults matches
          pure
            <| Aeson.object
              [ "query" .= query,
                "search_type" .= ("keyword_fallback" :: Text),
                "note" .= ("no indexed history - run backfill" :: Text),
                "total_matches" .= length matches,
                "results" .= map formatSearchResult limited
              ]

formatSemanticResult :: (Memory.ChatHistoryEntry, Float) -> Aeson.Value
formatSemanticResult (entry, score) =
  Aeson.object
    [ "timestamp" .= Time.formatTime Time.defaultTimeLocale "%Y-%m-%d %H:%M:%S" (Memory.cheCreatedAt entry),
      "role" .= Memory.cheRole entry,
      "sender" .= Memory.cheSenderName entry,
      "similarity" .= score,
      "content" .= Text.take 500 (Memory.cheContent entry)
    ]

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"