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

-- | Agent Event types for observability and streaming.
--
-- Captures all events during agent execution for logging,
-- streaming to web UI, and future interactive chat.
module Omni.Agent.Event
  ( AgentEvent (..),
    EventType (..),
    eventToJSON,
    eventFromJSON,
    formatEventForTerminal,
  )
where

import Alpha
import Data.Aeson ((.=))
import qualified Data.Aeson as Aeson
import qualified Data.Text as Text
import Data.Time (UTCTime, defaultTimeLocale, formatTime)

-- | Types of agent events
data EventType
  = Assistant -- LLM text response
  | ToolCall -- Tool invocation with arguments
  | ToolResult -- Tool execution result
  | UserMessage -- For future interactive chat
  | Cost -- Token usage and cost info
  | Error -- Failures and errors
  | Complete -- Session ended successfully
  deriving (Show, Eq, Read)

-- | A single agent event with timestamp and content
data AgentEvent = AgentEvent
  { eventType :: EventType,
    eventTimestamp :: UTCTime,
    eventContent :: Aeson.Value
  }
  deriving (Show, Eq)

-- | Convert event to JSON for storage/streaming
eventToJSON :: AgentEvent -> Aeson.Value
eventToJSON e =
  Aeson.object
    [ "type" .= show (eventType e),
      "timestamp" .= eventTimestamp e,
      "content" .= eventContent e
    ]

-- | Parse event from JSON
eventFromJSON :: Aeson.Value -> Maybe AgentEvent
eventFromJSON v = do
  obj <- case v of
    Aeson.Object o -> Just o
    _ -> Nothing
  typeStr <- case Aeson.lookup "type" (Aeson.toList obj) of
    Just (Aeson.String t) -> Just (Text.unpack t)
    _ -> Nothing
  eventT <- readMaybe typeStr
  ts <- case Aeson.lookup "timestamp" (Aeson.toList obj) of
    Just t -> Aeson.parseMaybe Aeson.parseJSON t
    _ -> Nothing
  content <- Aeson.lookup "content" (Aeson.toList obj)
  pure
    AgentEvent
      { eventType = eventT,
        eventTimestamp = ts,
        eventContent = content
      }
  where
    Aeson.lookup k pairs = snd </ find (\(k', _) -> k' == k) pairs
    Aeson.toList (Aeson.Object o) = map (first Aeson.toText) (Aeson.toList o)
    Aeson.toList _ = []
    Aeson.toText = id
    first f (a, b) = (f a, b)

-- | Format event for terminal display
formatEventForTerminal :: AgentEvent -> Text
formatEventForTerminal e =
  let ts = Text.pack <| formatTime defaultTimeLocale "%H:%M:%S" (eventTimestamp e)
      content = case eventType e of
        Assistant -> case eventContent e of
          Aeson.String t -> "Assistant: " <> truncate' 100 t
          _ -> "Assistant: <message>"
        ToolCall -> case eventContent e of
          Aeson.Object _ ->
            let toolName = getField "tool" (eventContent e)
             in "Tool: " <> toolName
          _ -> "Tool: <call>"
        ToolResult -> case eventContent e of
          Aeson.Object _ ->
            let toolName = getField "tool" (eventContent e)
                success = getField "success" (eventContent e)
             in "Result: " <> toolName <> " (" <> success <> ")"
          _ -> "Result: <result>"
        UserMessage -> case eventContent e of
          Aeson.String t -> "User: " <> truncate' 100 t
          _ -> "User: <message>"
        Cost -> case eventContent e of
          Aeson.Object _ ->
            let tokens = getField "tokens" (eventContent e)
                cents = getField "cents" (eventContent e)
             in "Cost: " <> tokens <> " tokens, " <> cents <> " cents"
          _ -> "Cost: <info>"
        Error -> case eventContent e of
          Aeson.String t -> "Error: " <> t
          _ -> "Error: <error>"
        Complete -> "Complete"
   in "[" <> ts <> "] " <> content
  where
    truncate' n t = if Text.length t > n then Text.take n t <> "..." else t
    getField key val = case val of
      Aeson.Object o -> case Aeson.lookup key (Aeson.toList o) of
        Just (Aeson.String s) -> s
        Just (Aeson.Number n) -> Text.pack (show n)
        Just (Aeson.Bool b) -> if b then "ok" else "failed"
        _ -> "<" <> key <> ">"
      _ -> "<" <> key <> ">"
      where
        Aeson.lookup k pairs = snd </ find (\(k', _) -> k' == k) pairs
        Aeson.toList (Aeson.Object o') = map (first' Aeson.toText) (Aeson.toList o')
        Aeson.toList _ = []
        Aeson.toText = id
        first' f (a, b) = (f a, b)

-- Helper constructors for common events

mkAssistantEvent :: UTCTime -> Text -> AgentEvent
mkAssistantEvent ts content =
  AgentEvent
    { eventType = Assistant,
      eventTimestamp = ts,
      eventContent = Aeson.String content
    }

mkToolCallEvent :: UTCTime -> Text -> Aeson.Value -> AgentEvent
mkToolCallEvent ts toolName args =
  AgentEvent
    { eventType = ToolCall,
      eventTimestamp = ts,
      eventContent = Aeson.object ["tool" .= toolName, "args" .= args]
    }

mkToolResultEvent :: UTCTime -> Text -> Bool -> Text -> AgentEvent
mkToolResultEvent ts toolName success output =
  AgentEvent
    { eventType = ToolResult,
      eventTimestamp = ts,
      eventContent =
        Aeson.object
          [ "tool" .= toolName,
            "success" .= success,
            "output" .= output
          ]
    }

mkCostEvent :: UTCTime -> Int -> Int -> AgentEvent
mkCostEvent ts tokens cents =
  AgentEvent
    { eventType = Cost,
      eventTimestamp = ts,
      eventContent = Aeson.object ["tokens" .= tokens, "cents" .= cents]
    }

mkErrorEvent :: UTCTime -> Text -> AgentEvent
mkErrorEvent ts msg =
  AgentEvent
    { eventType = Error,
      eventTimestamp = ts,
      eventContent = Aeson.String msg
    }

mkCompleteEvent :: UTCTime -> AgentEvent
mkCompleteEvent ts =
  AgentEvent
    { eventType = Complete,
      eventTimestamp = ts,
      eventContent = Aeson.Null
    }