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

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

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

-- | Status of the agent for the UI
data Status = Status
  { statusWorker :: Text,
    statusTask :: 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,
      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 2 lines at bottom
  IO.hPutStrLn IO.stderr ""
  IO.hPutStrLn IO.stderr ""
  ANSI.hCursorUp IO.stderr 2

-- | 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 lbs = BL.fromStrict (TextEnc.encodeUtf8 line)
  case decode lbs of
    Just (Object obj) -> do
      let message =
            case KM.lookup "message" obj of
              Just (String m) -> Just m
              _ -> Nothing

      let toolName =
            case KM.lookup "toolName" obj of
              Just (String t) -> Just t
              _ -> Nothing

      let level =
            case KM.lookup "level" obj of
              Just (String l) -> Just l
              _ -> Nothing

      case message of
        Just "executing 1 tools in 1 batch(es)" -> do
          let batchTool =
                case KM.lookup "batches" obj of
                  Just (Array b) ->
                    case V.toList b of
                      (Array b0 : _) ->
                        case V.toList b0 of
                          (String t : _) -> Just t
                          _ -> Nothing
                      _ -> Nothing
                  _ -> Nothing
          updateActivity ("THOUGHT: Planning tool execution (" <> fromMaybe "unknown" batchTool <> ")")

        Just "Tool Bash permitted - action: allow" ->
          updateActivity "TOOL: Bash command executed"

        Just msg | toolName /= Nothing && msg == "Processing tool completion for ledger" ->
          updateActivity ("TOOL: " <> fromMaybe "unknown" toolName <> " completed")

        Just "ide-fs" -> do
          let method =
                case KM.lookup "method" obj of
                  Just (String m) -> Just m
                  _ -> Nothing
          case method of
            Just "readFile" -> do
              let path =
                    case KM.lookup "path" obj of
                      Just (String p) -> Just p
                      _ -> Nothing
              case path of
                Just p -> updateActivity ("READ: " <> p)
                Nothing -> pure ()
            _ -> pure ()

        Just "System prompt build complete (no changes)" ->
          updateActivity "THINKING..."

        Just "System prompt build complete (first build)" ->
          updateActivity "STARTING new task context"

        Just msg | level == Just "error" ->
          updateActivity ("ERROR: " <> msg)

        _ -> pure ()

    _ -> pure ()

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

  -- 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 two status lines
render :: IO ()
render = do
  Status {..} <- readIORef currentStatus

  -- Line 1: Meta
  -- [Worker: name] Task: t-123 | Files: 3 | Credits: $0.45 | Time: 05:23
  let taskStr = maybe "None" identity statusTask
      meta =
        "[Worker: "
          <> statusWorker
          <> "] Task: "
          <> taskStr
          <> " | Files: "
          <> tshow statusFiles
          <> " | Credits: $"
          <> tshow statusCredits
          <> " | Time: "
          <> statusTime

  ANSI.hSetCursorColumn IO.stderr 0
  ANSI.hClearLine IO.stderr
  TIO.hPutStr IO.stderr meta

  -- Line 2: Activity
  -- [14:05:22] > Thinking...
  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 1
  IO.hFlush IO.stderr