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

-- | Concurrent logging with multi-line output support
module Omni.Log.Concurrent
  ( LineManager,
    BuildState (..),
    withLineManager,
    initializeLines,
    updateLine,
    updateLineState,
  )
where

import Alpha
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import qualified Data.Map as Map
import qualified Data.Text as Text
import Omni.Namespace (Namespace)
import qualified Omni.Namespace as Namespace
import Rainbow (chunk, fore, green, red)
import qualified Rainbow
import qualified System.Console.ANSI as ANSI
import qualified System.Environment as Env
import qualified System.IO as IO
import System.IO.Unsafe (unsafePerformIO)

data BuildState = Pending | Building | Success | Failed
  deriving (Eq, Show)

data LineManager = LineManager
  { lmNamespaces :: [Namespace],
    lmSupportsANSI :: Bool
  }

{-# NOINLINE currentLineManager #-}
currentLineManager :: IORef (Maybe LineManager)
currentLineManager = unsafePerformIO (newIORef Nothing)

{-# NOINLINE namespaceLines #-}
namespaceLines :: IORef (Map Namespace Int)
namespaceLines = unsafePerformIO (newIORef Map.empty)

withLineManager :: [Namespace] -> (LineManager -> IO a) -> IO a
withLineManager nss action = do
  supportsANSI <- checkANSISupport

  if not supportsANSI
    then do
      let mgr = LineManager {lmNamespaces = nss, lmSupportsANSI = False}
      writeIORef currentLineManager (Just mgr)
      result <- action mgr
      writeIORef currentLineManager Nothing
      writeIORef namespaceLines Map.empty
      pure result
    else do
      let numLines = length nss
      IO.hPutStrLn IO.stderr ""
      replicateM_ numLines (IO.hPutStrLn IO.stderr "")
      ANSI.hCursorUp IO.stderr numLines

      let mgr = LineManager {lmNamespaces = nss, lmSupportsANSI = True}
      writeIORef currentLineManager (Just mgr)

      -- Initialize the namespace -> line mapping
      writeIORef namespaceLines (Map.fromList <| zip nss [0 ..])

      result <- action mgr

      IO.hPutStrLn IO.stderr ""
      writeIORef currentLineManager Nothing
      writeIORef namespaceLines Map.empty
      pure result

checkANSISupport :: IO Bool
checkANSISupport = do
  term <- Env.lookupEnv "TERM"
  area <- Env.lookupEnv "AREA"
  pure <| case (term, area) of
    (Just "dumb", _) -> False
    (_, Just "Live") -> False
    (Nothing, _) -> False
    _ -> True

-- | Initialize all lines with pending status
initializeLines :: LineManager -> IO ()
initializeLines LineManager {..} =
  when lmSupportsANSI <| do
    nsMap <- readIORef namespaceLines
    forM_ (Map.toList nsMap) <| \(ns, _) -> do
      ANSI.hSetCursorColumn IO.stderr 0
      ANSI.hClearLine IO.stderr
      let nsText = Text.pack (Namespace.toPath ns)
      IO.hPutStrLn IO.stderr (Text.unpack <| "[…] " <> nsText)
      IO.hFlush IO.stderr

updateLine :: Namespace -> Text -> IO ()
updateLine ns output = do
  mMgr <- readIORef currentLineManager
  case mMgr of
    Nothing -> do
      IO.hPutStr IO.stderr (Text.unpack <| output <> "\r")
      IO.hFlush IO.stderr
    Just LineManager {..} ->
      if not lmSupportsANSI
        then do
          IO.hPutStr IO.stderr (Text.unpack <| output <> "\n")
          IO.hFlush IO.stderr
        else do
          nsMap <- readIORef namespaceLines
          case Map.lookup ns nsMap of
            Nothing -> pure ()
            Just lineNum -> do
              let numLines = length lmNamespaces
              -- Move to the target line from bottom
              ANSI.hCursorUp IO.stderr (numLines - lineNum)
              ANSI.hSetCursorColumn IO.stderr 0
              ANSI.hClearLine IO.stderr
              let nsText = Text.pack (Namespace.toPath ns)
              let formattedOutput = if Text.null output then "[~] " <> nsText else "[~] " <> nsText <> ": " <> output
              IO.hPutStr IO.stderr (Text.unpack formattedOutput)
              IO.hFlush IO.stderr
              -- Move back to bottom
              ANSI.hCursorDown IO.stderr (numLines - lineNum)

updateLineState :: Namespace -> BuildState -> IO ()
updateLineState ns buildState = do
  mMgr <- readIORef currentLineManager
  case mMgr of
    Nothing -> pure ()
    Just LineManager {..} ->
      when lmSupportsANSI <| do
        nsMap <- readIORef namespaceLines
        case Map.lookup ns nsMap of
          Nothing -> pure ()
          Just lineNum -> do
            let numLines = length lmNamespaces
            -- Move to the target line from bottom
            ANSI.hCursorUp IO.stderr (numLines - lineNum)
            ANSI.hSetCursorColumn IO.stderr 0
            ANSI.hClearLine IO.stderr
            let nsText = Text.pack (Namespace.toPath ns)
            case buildState of
              Success ->
                Rainbow.hPutChunks IO.stderr [fore green <| chunk <| "[✓] " <> nsText]
              Failed ->
                Rainbow.hPutChunks IO.stderr [fore red <| chunk <| "[x] " <> nsText]
              Pending ->
                IO.hPutStr IO.stderr (Text.unpack <| "[…] " <> nsText)
              Building ->
                IO.hPutStr IO.stderr (Text.unpack <| "[~] " <> nsText)
            IO.hFlush IO.stderr
            -- Move back to bottom
            ANSI.hCursorDown IO.stderr (numLines - lineNum)