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

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

import Alpha
import Data.IORef (IORef, atomicModifyIORef', modifyIORef', newIORef, readIORef)
import qualified Data.Map as Map
import qualified Data.Text as Text
import Omni.Namespace (Namespace)
import qualified Omni.Namespace as Namespace
import qualified System.Console.ANSI as ANSI
import qualified System.Environment as Env
import qualified System.IO as IO

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

data BuildStatus = BuildStatus
  { bsTarget :: Namespace,
    bsLastOutput :: Text,
    bsState :: BuildState
  }

data LineManager = LineManager
  { lmLines :: IORef (Map Int (Maybe BuildStatus)),
    lmMaxLines :: Int,
    lmCurrentLine :: IORef Int,
    lmSupportsANSI :: Bool
  }

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
      action
        LineManager
          { lmLines = linesRef,
            lmMaxLines = 1,
            lmCurrentLine = currentRef,
            lmSupportsANSI = False
          }
    else do
      replicateM_ maxLines (IO.hPutStrLn IO.stderr "")
      ANSI.hCursorUp IO.stderr maxLines

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

      result <-
        action
          LineManager
            { lmLines = linesRef,
              lmMaxLines = maxLines,
              lmCurrentLine = currentRef,
              lmSupportsANSI = True
            }

      ANSI.hCursorDown IO.stderr maxLines
      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
      atomicModifyIORef' lmLines <| \linesMap ->
        case findFirstFree linesMap of
          Nothing -> (linesMap, Nothing)
          Just lineNum ->
            let status = BuildStatus ns "" Building
                linesMap' = Map.insert lineNum (Just status) linesMap
             in (linesMap', Just lineNum)
  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
        IO.hPutStr IO.stderr (Text.unpack output)
        IO.hFlush IO.stderr

        ANSI.hRestoreCursor IO.stderr

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

releaseLine :: LineManager -> Maybe Int -> BuildState -> IO ()
releaseLine LineManager {..} mLineNum 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 statusChar = case buildState of
              Success -> "✓"
              Failed -> "✗"
              Building -> "…"
        IO.hPutStr IO.stderr statusChar
        IO.hFlush IO.stderr

        ANSI.hRestoreCursor IO.stderr