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

-- | Concurrent logging with multi-line output support
module Omni.Log.Concurrent
  ( LineManager,
    BuildState (..),
    withLineManager,
    reserveLine,
    updateLine,
    releaseLine,
    updateCurrentLine,
    releaseCurrentLine,
  )
where

import Alpha
import Data.IORef (IORef, atomicModifyIORef', modifyIORef', 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 = Building | Success | Failed
  deriving (Eq, Show)

newtype BuildStatus = BuildStatus
  { bsLastOutput :: Text
  }

data LineManager = LineManager
  { lmLines :: IORef (Map Int (Maybe BuildStatus)),
    lmMaxLines :: Int,
    lmCurrentLine :: IORef Int,
    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 :: Int -> (LineManager -> IO a) -> IO a
withLineManager maxLines action = do
  supportsANSI <- checkANSISupport

  if not supportsANSI
    then do
      linesRef <- newIORef Map.empty
      currentRef <- newIORef 0
      let mgr =
            LineManager
              { lmLines = linesRef,
                lmMaxLines = 1,
                lmCurrentLine = currentRef,
                lmSupportsANSI = False
              }
      writeIORef currentLineManager (Just mgr)
      result <- action mgr
      writeIORef currentLineManager Nothing
      writeIORef namespaceLines Map.empty
      pure result
    else do
      IO.hPutStrLn IO.stderr ""
      replicateM_ maxLines (IO.hPutStrLn IO.stderr "")
      ANSI.hCursorUp IO.stderr maxLines

      linesRef <- newIORef (Map.fromList [(i, Nothing) | i <- [0 .. maxLines - 1]])
      currentRef <- newIORef 0

      let mgr =
            LineManager
              { lmLines = linesRef,
                lmMaxLines = maxLines,
                lmCurrentLine = currentRef,
                lmSupportsANSI = True
              }
      writeIORef currentLineManager (Just mgr)

      result <- action mgr

      forM_ [0 .. maxLines - 1] <| \_ -> do
        ANSI.hCursorDown IO.stderr 1
        ANSI.hSetCursorColumn IO.stderr 0
        ANSI.hClearLine 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

reserveLine :: LineManager -> Namespace -> IO (Maybe Int)
reserveLine LineManager {..} ns =
  if not lmSupportsANSI
    then pure Nothing
    else do
      mLine <-
        atomicModifyIORef' lmLines <| \linesMap ->
          case findFirstFree linesMap of
            Nothing -> (linesMap, Nothing)
            Just lineNum ->
              let status = BuildStatus ""
                  linesMap' = Map.insert lineNum (Just status) linesMap
               in (linesMap', Just lineNum)
      case mLine of
        Just lineNum -> do
          modifyIORef' namespaceLines (Map.insert ns lineNum)
          currentLine <- readIORef lmCurrentLine
          ANSI.hSaveCursor IO.stderr
          ANSI.hSetCursorColumn IO.stderr 0
          let linesToMove = currentLine - lineNum
          when (linesToMove > 0) <| ANSI.hCursorUp IO.stderr linesToMove
          when (linesToMove < 0) <| ANSI.hCursorDown IO.stderr (abs linesToMove)
          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
        Nothing -> pure ()
      pure mLine
  where
    findFirstFree :: Map Int (Maybe BuildStatus) -> Maybe Int
    findFirstFree m =
      Map.toList m
        |> filter (\(_, mbs) -> isNothing mbs)
        |> map fst
        |> listToMaybe

updateLine :: LineManager -> Maybe Int -> Namespace -> Text -> IO ()
updateLine LineManager {..} mLineNum ns output =
  if not lmSupportsANSI
    then do
      IO.hPutStr IO.stderr (Text.unpack <| output <> "\n")
      IO.hFlush IO.stderr
    else case mLineNum of
      Nothing -> pure ()
      Just lineNum -> do
        currentLine <- readIORef lmCurrentLine

        ANSI.hSaveCursor IO.stderr
        ANSI.hSetCursorColumn IO.stderr 0

        let linesToMove = currentLine - lineNum
        when (linesToMove > 0) <| ANSI.hCursorUp IO.stderr linesToMove
        when (linesToMove < 0) <| ANSI.hCursorDown IO.stderr (abs linesToMove)

        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

        modifyIORef' lmLines <| \linesMap ->
          Map.adjust (fmap (\bs -> bs {bsLastOutput = output})) lineNum linesMap

releaseLine :: LineManager -> Maybe Int -> Namespace -> BuildState -> IO ()
releaseLine LineManager {..} mLineNum ns buildState =
  case mLineNum of
    Nothing -> pure ()
    Just lineNum -> do
      modifyIORef' lmLines <| \linesMap ->
        Map.insert lineNum Nothing linesMap

      when lmSupportsANSI <| do
        current <- readIORef lmCurrentLine
        ANSI.hSaveCursor IO.stderr
        ANSI.hSetCursorColumn IO.stderr 0
        ANSI.hCursorUp IO.stderr (current - 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]
          Building ->
            IO.hPutStr IO.stderr (Text.unpack <| "[~] " <> nsText)
        IO.hFlush IO.stderr

        ANSI.hRestoreCursor IO.stderr

updateCurrentLine :: Namespace -> Text -> IO ()
updateCurrentLine ns output = do
  mMgr <- readIORef currentLineManager
  case mMgr of
    Nothing -> do
      IO.hPutStr IO.stderr (Text.unpack <| output <> "\r")
      IO.hFlush IO.stderr
    Just mgr -> do
      nsMap <- readIORef namespaceLines
      case Map.lookup ns nsMap of
        Nothing -> do
          IO.hPutStr IO.stderr (Text.unpack <| output <> "\r")
          IO.hFlush IO.stderr
        Just lineNum -> updateLine mgr (Just lineNum) ns output

releaseCurrentLine :: Namespace -> BuildState -> IO ()
releaseCurrentLine ns buildState = do
  mMgr <- readIORef currentLineManager
  case mMgr of
    Nothing -> pure ()
    Just mgr -> do
      nsMap <- readIORef namespaceLines
      case Map.lookup ns nsMap of
        Nothing -> pure ()
        Just lineNum -> do
          releaseLine mgr (Just lineNum) ns buildState
          modifyIORef' namespaceLines (Map.delete ns)