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

-- : out omni-agent-log
module Omni.Agent.Log where

import Alpha
import Data.Aeson (Value (..), decode)
import qualified Data.Aeson.KeyMap as KM
import qualified Data.ByteString.Lazy as BL
import Data.IORef (IORef, modifyIORef', newIORef, readIORef, writeIORef)
import qualified Data.Text.Encoding as TextEnc
import qualified Data.Text.IO as TIO
import qualified Data.Vector as V
import qualified System.Console.ANSI as ANSI
import qualified System.IO as IO
import System.IO.Unsafe (unsafePerformIO)
import Text.Printf (printf)

-- | Parsed log entry
data LogEntry = LogEntry
  { leMessage :: Maybe Text,
    leLevel :: Maybe Text,
    leToolName :: Maybe Text,
    leBatches :: Maybe [[Text]],
    leMethod :: Maybe Text,
    lePath :: Maybe Text
  }
  deriving (Show, Eq)

-- | Status of the agent for the UI
data Status = Status
  { statusWorker :: Text,
    statusTask :: Maybe Text,
    statusThreadId :: Maybe Text,
    statusFiles :: Int,
    statusCredits :: Double,
    statusTime :: Text, -- formatted time string
    statusActivity :: Text
  }
  deriving (Show, Eq)

emptyStatus :: Text -> Status
emptyStatus workerName =
  Status
    { statusWorker = workerName,
      statusTask = Nothing,
      statusThreadId = Nothing,
      statusFiles = 0,
      statusCredits = 0.0,
      statusTime = "00:00",
      statusActivity = "Idle"
    }

-- | Global state for the status bar
{-# NOINLINE currentStatus #-}
currentStatus :: IORef Status
currentStatus = unsafePerformIO (newIORef (emptyStatus "Unknown"))

-- | Initialize the status bar system
init :: Text -> IO ()
init workerName = do
  IO.hSetBuffering IO.stderr IO.LineBuffering
  writeIORef currentStatus (emptyStatus workerName)
  -- Reserve 5 lines at bottom
  replicateM_ 5 (IO.hPutStrLn IO.stderr "")
  ANSI.hCursorUp IO.stderr 5

-- | Update the status
update :: (Status -> Status) -> IO ()
update f = do
  modifyIORef' currentStatus f
  render

-- | Set the activity message
updateActivity :: Text -> IO ()
updateActivity msg = update (\s -> s {statusActivity = msg})

-- | Process a log line from the agent and update status if relevant
processLogLine :: Text -> IO ()
processLogLine line = do
  let entry = parseLine line
  for_ (entry +> formatLogEntry) updateActivity

-- | Parse a JSON log line into a LogEntry
parseLine :: Text -> Maybe LogEntry
parseLine line = do
  let lbs = BL.fromStrict (TextEnc.encodeUtf8 line)
  obj <- decode lbs
  case obj of
    Object o ->
      Just
        LogEntry
          { leMessage = getString "message" o,
            leLevel = getString "level" o,
            leToolName = getString "toolName" o,
            leBatches = getBatches o,
            leMethod = getString "method" o,
            lePath = getString "path" o
          }
    _ -> Nothing
  where
    getString k o =
      case KM.lookup k o of
        Just (String s) -> Just s
        _ -> Nothing

    getBatches o =
      case KM.lookup "batches" o of
        Just (Array b) ->
          Just
            <| mapMaybe
              ( \case
                  Array b0 ->
                    Just
                      <| mapMaybe
                        ( \case
                            String s -> Just s
                            _ -> Nothing
                        )
                        (V.toList b0)
                  _ -> Nothing
              )
              (V.toList b)
        _ -> Nothing

-- | Format a log entry into a user-friendly status message (NO EMOJIS)
formatLogEntry :: LogEntry -> Maybe Text
formatLogEntry LogEntry {..} =
  case leMessage of
    Just "executing 1 tools in 1 batch(es)" -> do
      let tools = fromMaybe [] leBatches
      let firstTool = case tools of
            ((t : _) : _) -> t
            _ -> "unknown"
      Just ("THOUGHT: Planning tool execution (" <> firstTool <> ")")
    Just "Tool Bash permitted - action: allow" ->
      Just "TOOL: Bash command executed"
    Just "Processing tool completion for ledger"
      | isJust leToolName ->
          Just ("TOOL: " <> fromMaybe "unknown" leToolName <> " completed")
    Just "ide-fs" | leMethod == Just "readFile" ->
      case lePath of
        Just p -> Just ("READ: " <> p)
        _ -> Nothing
    Just "System prompt build complete (no changes)" ->
      Just "THINKING..."
    Just "System prompt build complete (first build)" ->
      Just "STARTING new task context"
    Just msg
      | leLevel == Just "error" ->
          Just ("ERROR: " <> msg)
    _ -> Nothing

-- | Log a scrolling message (appears above status bars)
log :: Text -> IO ()
log msg = do
  -- Clear status bars (5 lines)
  ANSI.hClearLine IO.stderr
  ANSI.hCursorDown IO.stderr 1
  ANSI.hClearLine IO.stderr
  ANSI.hCursorDown IO.stderr 1
  ANSI.hClearLine IO.stderr
  ANSI.hCursorDown IO.stderr 1
  ANSI.hClearLine IO.stderr
  ANSI.hCursorDown IO.stderr 1
  ANSI.hClearLine IO.stderr
  ANSI.hCursorUp IO.stderr 4

  -- Print message (scrolls screen)
  TIO.hPutStrLn IO.stderr msg

  -- Re-render status bars at bottom
  -- (Since we scrolled, we are now on the line above where the first status line should be)
  render

-- | Render the 5 status lines (Vertical Layout)
render :: IO ()
render = do
  Status {..} <- readIORef currentStatus

  let taskStr = maybe "None" identity statusTask
      threadStr = maybe "None" identity statusThreadId

  -- Line 1: Worker + Time
  ANSI.hSetCursorColumn IO.stderr 0
  ANSI.hClearLine IO.stderr
  TIO.hPutStr IO.stderr <| "Worker: " <> statusWorker <> " | Time: " <> statusTime

  -- Line 2: Task
  ANSI.hCursorDown IO.stderr 1
  ANSI.hSetCursorColumn IO.stderr 0
  ANSI.hClearLine IO.stderr
  TIO.hPutStr IO.stderr <| "Task: " <> taskStr

  -- Line 3: Thread
  ANSI.hCursorDown IO.stderr 1
  ANSI.hSetCursorColumn IO.stderr 0
  ANSI.hClearLine IO.stderr
  TIO.hPutStr IO.stderr <| "Thread: " <> threadStr

  -- Line 4: Credits
  ANSI.hCursorDown IO.stderr 1
  ANSI.hSetCursorColumn IO.stderr 0
  ANSI.hClearLine IO.stderr
  TIO.hPutStr IO.stderr <| "Credits: $" <> str (printf "%.2f" statusCredits :: String)

  -- Line 5: Activity
  ANSI.hCursorDown IO.stderr 1
  ANSI.hSetCursorColumn IO.stderr 0
  ANSI.hClearLine IO.stderr
  TIO.hPutStr IO.stderr ("> " <> statusActivity)

  -- Return cursor to Line 1
  ANSI.hCursorUp IO.stderr 4
  IO.hFlush IO.stderr