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

module Omni.Agent.Log
  ( LogEntry (..),
    Status (..),
    initialStatus,
    updateStatus,
    renderStatus,
    parseLine,
    format,
  )
where

import Alpha
import Data.Aeson (FromJSON (..), (.:), (.:?))
import qualified Data.Aeson as Aeson
import qualified Data.ByteString.Lazy as BSL
import qualified Data.Set as Set

data LogEntry = LogEntry
  { leMessage :: Text,
    leLevel :: Maybe Text,
    leToolName :: Maybe Text,
    leBatches :: Maybe [[Text]],
    leMethod :: Maybe Text,
    lePath :: Maybe Text,
    leTimestamp :: Maybe Text
  }
  deriving (Show, Eq, Generic)

instance FromJSON LogEntry where
  parseJSON =
    Aeson.withObject "LogEntry" <| \v ->
      ( LogEntry
          </ (v .: "message")
      )
        <*> v
        .:? "level"
        <*> v
        .:? "toolName"
        <*> v
        .:? "batches"
        <*> v
        .:? "method"
        <*> v
        .:? "path"
        <*> v
        .:? "timestamp"

data Status = Status
  { sWorkerName :: Text,
    sTaskId :: Maybe Text,
    sFiles :: Set.Set Text,
    sStartTime :: Maybe Text,
    sLastActivity :: Text
  }
  deriving (Show, Eq, Generic)

initialStatus :: Text -> Status
initialStatus name =
  Status
    { sWorkerName = name,
      sTaskId = Nothing,
      sFiles = Set.empty,
      sStartTime = Nothing,
      sLastActivity = "Idle"
    }

updateStatus :: LogEntry -> Status -> Status
updateStatus e s =
  let s' = case format e of
        Just msg -> s {sLastActivity = msg}
        Nothing -> s
      s'' = case leTimestamp e of
        Just t -> if isNothing (sStartTime s) then s' {sStartTime = Just t} else s'
        Nothing -> s'
   in case (leMessage e, lePath e) of
        ("ide-fs", Just p) -> s'' {sFiles = Set.insert p (sFiles s'')}
        _ -> s''

renderStatus :: Status -> Text
renderStatus s =
  let line1 =
        "[Worker: "
          <> sWorkerName s
          <> "] "
          <> "Task: "
          <> fromMaybe "None" (sTaskId s)
          <> " | Files: "
          <> show (Set.size (sFiles s))
      line2 = sLastActivity s
   in line1 <> "\n" <> line2

parseLine :: Text -> Maybe LogEntry
parseLine line = Aeson.decode <| BSL.fromStrict <| encodeUtf8 line

format :: LogEntry -> Maybe Text
format e =
  case leMessage e of
    "executing 1 tools in 1 batch(es)" ->
      let tool = case leBatches e of
            Just ((t : _) : _) -> t
            _ -> "unknown"
       in Just <| "🤖  THOUGHT: Planning tool execution (" <> tool <> ")"
    "Tool Bash permitted - action: allow" ->
      Just "🔧  TOOL: Bash command executed"
    msg
      | "Processing tool completion for ledger" == msg && isJust (leToolName e) ->
          Just <| "✅  TOOL: " <> fromMaybe "" (leToolName e) <> " completed"
    "ide-fs" ->
      case leMethod e of
        Just "readFile" -> Just <| "📂  READ: " <> fromMaybe "" (lePath e)
        _ -> Nothing
    "System prompt build complete (no changes)" ->
      Just "🧠  THINKING..."
    "System prompt build complete (first build)" ->
      Just "🚀  STARTING new task context"
    msg ->
      case leLevel e of
        Just "error" -> Just <| "❌  ERROR: " <> msg
        _ -> Nothing