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

-- | Audit logging for Ava and subagents.
--
-- Persists all agent events to JSONL files for debugging and diagnosis.
-- Logs are stored in @AVA_DATA_ROOT/logs/@.
--
-- Structure:
--   - @logs/ava/YYYY-MM-DD.jsonl@ - Daily Ava conversation logs
--   - @logs/subagents/S-{id}.jsonl@ - Per-subagent traces
--
-- : out omni-agent-auditlog
-- : dep aeson
-- : dep bytestring
-- : dep directory
-- : dep time
-- : dep uuid
module Omni.Agent.AuditLog
  ( -- * Types
    AuditLogEntry (..),
    AuditEventType (..),
    LogMetadata (..),
    SubagentId (..),
    SessionId (..),
    AgentId (..),

    -- * Writing logs
    writeAvaLog,
    writeSubagentLog,

    -- * Reading logs
    readAvaLogs,
    readSubagentLogs,
    getRecentAvaLogs,

    -- * SubagentId
    newSubagentId,
    subagentLogPath,

    -- * Paths
    avaLogsDir,
    subagentLogsDir,

    -- * Helpers
    mkLogEntry,
    emptyMetadata,

    -- * Testing
    main,
    test,
  )
where

import Alpha
import Data.Aeson ((.=))
import qualified Data.Aeson as Aeson
import qualified Data.ByteString.Lazy as LBS
import qualified Data.Text as Text
import qualified Data.Text.Encoding as Text
import qualified Data.Time as Time
import qualified Data.UUID as UUID
import qualified Data.UUID.V4 as UUID
import qualified Omni.Agent.Paths as Paths
import qualified Omni.Test as Test
import qualified System.Directory as Dir
import System.FilePath ((</>))

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

test :: Test.Tree
test =
  Test.group
    "Omni.Agent.AuditLog"
    [ Test.unit "SubagentId JSON roundtrip" <| do
        let sid = SubagentId "abc123"
        case Aeson.decode (Aeson.encode sid) of
          Nothing -> Test.assertFailure "Failed to decode SubagentId"
          Just decoded -> decoded Test.@=? sid,
      Test.unit "AuditEventType JSON roundtrip" <| do
        let types = [UserMessage, AssistantMessage, ToolCall, ToolResult, SubagentSpawn, SubagentComplete, ErrorOccurred]
        forM_ types <| \t ->
          case Aeson.decode (Aeson.encode t) of
            Nothing -> Test.assertFailure ("Failed to decode: " <> show t)
            Just decoded -> decoded Test.@=? t,
      Test.unit "AuditLogEntry JSON roundtrip" <| do
        now <- Time.getCurrentTime
        let entry =
              AuditLogEntry
                { logTimestamp = now,
                  logSessionId = SessionId "sess-123",
                  logAgentId = AgentId "ava",
                  logUserId = Just "ben",
                  logEventType = AssistantMessage,
                  logContent = Aeson.String "Hello",
                  logMetadata = emptyMetadata
                }
        case Aeson.decode (Aeson.encode entry) of
          Nothing -> Test.assertFailure "Failed to decode AuditLogEntry"
          Just decoded -> logEventType decoded Test.@=? AssistantMessage,
      Test.unit "subagentLogPath constructs correct path" <| do
        let sid = SubagentId "abc123"
        let path = subagentLogPath sid
        (Text.pack "abc123.jsonl" `Text.isInfixOf` Text.pack path) Test.@=? True
    ]

newtype SubagentId = SubagentId {unSubagentId :: Text}
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON SubagentId where
  toJSON (SubagentId sid) = Aeson.String sid

instance Aeson.FromJSON SubagentId where
  parseJSON = Aeson.withText "SubagentId" (pure <. SubagentId)

newtype SessionId = SessionId {unSessionId :: Text}
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON SessionId where
  toJSON (SessionId sid) = Aeson.String sid

instance Aeson.FromJSON SessionId where
  parseJSON = Aeson.withText "SessionId" (pure <. SessionId)

newtype AgentId = AgentId {unAgentId :: Text}
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON AgentId where
  toJSON (AgentId aid) = Aeson.String aid

instance Aeson.FromJSON AgentId where
  parseJSON = Aeson.withText "AgentId" (pure <. AgentId)

data AuditEventType
  = UserMessage
  | AssistantMessage
  | ToolCall
  | ToolResult
  | SubagentSpawn
  | SubagentComplete
  | ExtendedThinking
  | CostUpdate
  | ErrorOccurred
  | SessionStart
  | SessionEnd
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON AuditEventType where
  toJSON UserMessage = Aeson.String "user_message"
  toJSON AssistantMessage = Aeson.String "assistant_message"
  toJSON ToolCall = Aeson.String "tool_call"
  toJSON ToolResult = Aeson.String "tool_result"
  toJSON SubagentSpawn = Aeson.String "subagent_spawn"
  toJSON SubagentComplete = Aeson.String "subagent_complete"
  toJSON ExtendedThinking = Aeson.String "extended_thinking"
  toJSON CostUpdate = Aeson.String "cost_update"
  toJSON ErrorOccurred = Aeson.String "error"
  toJSON SessionStart = Aeson.String "session_start"
  toJSON SessionEnd = Aeson.String "session_end"

instance Aeson.FromJSON AuditEventType where
  parseJSON =
    Aeson.withText "AuditEventType" <| \case
      "user_message" -> pure UserMessage
      "assistant_message" -> pure AssistantMessage
      "tool_call" -> pure ToolCall
      "tool_result" -> pure ToolResult
      "subagent_spawn" -> pure SubagentSpawn
      "subagent_complete" -> pure SubagentComplete
      "extended_thinking" -> pure ExtendedThinking
      "cost_update" -> pure CostUpdate
      "error" -> pure ErrorOccurred
      "session_start" -> pure SessionStart
      "session_end" -> pure SessionEnd
      _ -> empty

data LogMetadata = LogMetadata
  { metaInputTokens :: Maybe Int,
    metaOutputTokens :: Maybe Int,
    metaCostCents :: Maybe Double,
    metaModelId :: Maybe Text,
    metaParentAgentId :: Maybe AgentId,
    metaDurationMs :: Maybe Int
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON LogMetadata where
  toJSON m =
    Aeson.object
      <| catMaybes
        [ ("input_tokens" .=) </ metaInputTokens m,
          ("output_tokens" .=) </ metaOutputTokens m,
          ("cost_cents" .=) </ metaCostCents m,
          ("model_id" .=) </ metaModelId m,
          ("parent_agent_id" .=) </ metaParentAgentId m,
          ("duration_ms" .=) </ metaDurationMs m
        ]

instance Aeson.FromJSON LogMetadata where
  parseJSON =
    Aeson.withObject "LogMetadata" <| \v ->
      LogMetadata
        </ (v Aeson..:? "input_tokens")
        <*> (v Aeson..:? "output_tokens")
        <*> (v Aeson..:? "cost_cents")
        <*> (v Aeson..:? "model_id")
        <*> (v Aeson..:? "parent_agent_id")
        <*> (v Aeson..:? "duration_ms")

emptyMetadata :: LogMetadata
emptyMetadata =
  LogMetadata
    { metaInputTokens = Nothing,
      metaOutputTokens = Nothing,
      metaCostCents = Nothing,
      metaModelId = Nothing,
      metaParentAgentId = Nothing,
      metaDurationMs = Nothing
    }

data AuditLogEntry = AuditLogEntry
  { logTimestamp :: Time.UTCTime,
    logSessionId :: SessionId,
    logAgentId :: AgentId,
    logUserId :: Maybe Text,
    logEventType :: AuditEventType,
    logContent :: Aeson.Value,
    logMetadata :: LogMetadata
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON AuditLogEntry where
  toJSON e =
    Aeson.object
      [ "timestamp" .= logTimestamp e,
        "session_id" .= logSessionId e,
        "agent_id" .= logAgentId e,
        "user_id" .= logUserId e,
        "event_type" .= logEventType e,
        "content" .= logContent e,
        "metadata" .= logMetadata e
      ]

instance Aeson.FromJSON AuditLogEntry where
  parseJSON =
    Aeson.withObject "AuditLogEntry" <| \v ->
      AuditLogEntry
        </ (v Aeson..: "timestamp")
        <*> (v Aeson..: "session_id")
        <*> (v Aeson..: "agent_id")
        <*> (v Aeson..:? "user_id")
        <*> (v Aeson..: "event_type")
        <*> (v Aeson..: "content")
        <*> (v Aeson..:? "metadata" Aeson..!= emptyMetadata)

avaLogsDir :: FilePath
avaLogsDir = Paths.avaLogsDir

subagentLogsDir :: FilePath
subagentLogsDir = Paths.subagentLogsDir

newSubagentId :: IO SubagentId
newSubagentId = do
  uuid <- UUID.nextRandom
  pure <| SubagentId <| Text.take 6 <| UUID.toText uuid

subagentLogPath :: SubagentId -> FilePath
subagentLogPath (SubagentId sid) =
  subagentLogsDir </> Text.unpack sid <> ".jsonl"

todayLogPath :: IO FilePath
todayLogPath = do
  today <- Time.utctDay </ Time.getCurrentTime
  let dateStr = Time.formatTime Time.defaultTimeLocale "%Y-%m-%d" today
  pure (avaLogsDir </> dateStr <> ".jsonl")

mkLogEntry ::
  SessionId ->
  AgentId ->
  Maybe Text ->
  AuditEventType ->
  Aeson.Value ->
  LogMetadata ->
  IO AuditLogEntry
mkLogEntry session agent user eventType content metadata = do
  now <- Time.getCurrentTime
  pure
    AuditLogEntry
      { logTimestamp = now,
        logSessionId = session,
        logAgentId = agent,
        logUserId = user,
        logEventType = eventType,
        logContent = content,
        logMetadata = metadata
      }

writeAvaLog :: AuditLogEntry -> IO ()
writeAvaLog entry = do
  Dir.createDirectoryIfMissing True avaLogsDir
  path <- todayLogPath
  let line = Aeson.encode entry <> "\n"
  LBS.appendFile path line

writeSubagentLog :: SubagentId -> AuditLogEntry -> IO ()
writeSubagentLog sid entry = do
  Dir.createDirectoryIfMissing True subagentLogsDir
  let path = subagentLogPath sid
  let line = Aeson.encode entry <> "\n"
  LBS.appendFile path line

readSubagentLogs :: SubagentId -> IO [AuditLogEntry]
readSubagentLogs sid = do
  let path = subagentLogPath sid
  exists <- Dir.doesFileExist path
  if exists
    then parseJsonlFile path
    else pure []

readAvaLogs :: Time.Day -> IO [AuditLogEntry]
readAvaLogs day = do
  let dateStr = Time.formatTime Time.defaultTimeLocale "%Y-%m-%d" day
  let path = avaLogsDir </> dateStr <> ".jsonl"
  exists <- Dir.doesFileExist path
  if exists
    then parseJsonlFile path
    else pure []

getRecentAvaLogs :: Int -> IO [AuditLogEntry]
getRecentAvaLogs n = do
  today <- Time.utctDay </ Time.getCurrentTime
  entries <- readAvaLogs today
  pure (take n (reverse entries))

parseJsonlFile :: FilePath -> IO [AuditLogEntry]
parseJsonlFile path = do
  contents <- LBS.readFile path
  let textLines = Text.lines <| Text.decodeUtf8 <| LBS.toStrict contents
  pure <| mapMaybe (Aeson.decodeStrict <. Text.encodeUtf8) textLines