{-# 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

      ANSI.hCursorDown IO.stderr numLines
      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, lineNum) -> do
      ANSI.hSaveCursor IO.stderr
      ANSI.hSetCursorColumn IO.stderr 0
      when (lineNum > 0) <| ANSI.hCursorDown IO.stderr lineNum
      ANSI.hClearLine IO.stderr
      let nsText = Text.pack (Namespace.toPath ns)
      IO.hPutStr IO.stderr (Text.unpack <| "[…] " <> nsText)
      IO.hFlush IO.stderr
      ANSI.hRestoreCursor 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
              ANSI.hSaveCursor IO.stderr
              ANSI.hSetCursorColumn IO.stderr 0
              ANSI.hCursorUp IO.stderr (length lmNamespaces - lineNum)
              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
              ANSI.hRestoreCursor IO.stderr

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
            ANSI.hSaveCursor IO.stderr
            ANSI.hSetCursorColumn IO.stderr 0
            ANSI.hCursorUp IO.stderr (length lmNamespaces - lineNum)
            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
            ANSI.hRestoreCursor IO.stderr