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

-- | Concurrent logging with efficient single-line updates
module Omni.Log.Concurrent
  ( LineManager,
    BuildState (..),
    withLineManager,
    initializeLines,
    updateLine,
    updateLineState,
  )
where

import Alpha
import Data.IORef (IORef, atomicModifyIORef', newIORef, readIORef, writeIORef)
import qualified Data.Map.Strict as Map
import qualified Data.Text as Text
import Omni.Log.Terminal (TerminalInfo (..), detectTerminal, truncateToWidth)
import Omni.Namespace (Namespace)
import qualified Omni.Namespace as Namespace
import Rainbow (chunk, fore, green, red, yellow)
import qualified Rainbow
import qualified System.Console.ANSI as ANSI
import qualified System.IO as IO
import System.IO.Unsafe (unsafePerformIO)

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

data LineManager = LineManager
  { lmNamespaces :: [Namespace],
    lmTermInfo :: TerminalInfo
  }

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

-- | Current state of each namespace
{-# NOINLINE namespaceStates #-}
namespaceStates :: IORef (Map Namespace (BuildState, Text))
namespaceStates = unsafePerformIO (newIORef Map.empty)

-- | Namespace to line number mapping (0-indexed from top of our area)
{-# NOINLINE namespaceLines #-}
namespaceLines :: IORef (Map Namespace Int)
namespaceLines = unsafePerformIO (newIORef Map.empty)

-- | Tracks if lines have been initialized
{-# NOINLINE linesInitialized #-}
linesInitialized :: IORef Bool
linesInitialized = unsafePerformIO (newIORef False)

-- | Global lock for all terminal operations
{-# NOINLINE terminalLock #-}
terminalLock :: MVar ()
terminalLock = unsafePerformIO (newMVar ())

withLineManager :: [Namespace] -> (LineManager -> IO a) -> IO a
withLineManager nss action = do
  existingMgr <- readIORef currentLineManager
  maybe createNewManager action existingMgr
  where
    createNewManager = do
      termInfo <- detectTerminal

      let mgr = LineManager {lmNamespaces = nss, lmTermInfo = termInfo}
      writeIORef currentLineManager (Just mgr)
      writeIORef linesInitialized False
      -- Initialize all namespaces to Pending state
      writeIORef namespaceStates (Map.fromList [(ns, (Pending, "")) | ns <- nss])
      -- Create line number mapping
      writeIORef namespaceLines (Map.fromList (zip nss [0 ..]))

      -- Reserve lines for the display area
      let numLines = length nss
      when (numLines > 0 && tiSupportsANSI termInfo) <| do
        -- Start fresh: clear from cursor to end of screen
        IO.hPutStr IO.stderr "\r"
        ANSI.hClearFromCursorToScreenEnd IO.stderr
        -- Print N blank lines to reserve space
        replicateM_ numLines (IO.hPutStrLn IO.stderr "")
        -- Move back to the first line and save this position
        ANSI.hCursorUp IO.stderr numLines
        ANSI.hSaveCursor IO.stderr
        IO.hFlush IO.stderr

      result <- action mgr

      -- Move to end and print final newline
      when (numLines > 0 && tiSupportsANSI termInfo) <| do
        -- Restore to top, then move down past our area
        ANSI.hRestoreCursor IO.stderr
        ANSI.hCursorDown IO.stderr numLines
        IO.hPutStrLn IO.stderr ""
        IO.hFlush IO.stderr

      writeIORef currentLineManager Nothing
      writeIORef namespaceStates Map.empty
      writeIORef namespaceLines Map.empty
      writeIORef linesInitialized False
      pure result

-- | Initialize display with all namespaces
initializeLines :: LineManager -> IO ()
initializeLines mgr = do
  alreadyInit <- readIORef linesInitialized
  unless alreadyInit <| do
    writeIORef linesInitialized True
    -- Draw all lines once at initialization
    redrawAll mgr

-- | Update progress text for a namespace
updateLine :: Namespace -> Text -> IO ()
updateLine ns output = do
  mMgr <- readIORef currentLineManager
  case mMgr of
    Nothing -> pure ()
    Just mgr -> do
      -- Update state and redraw just this line
      atomicModifyIORef' namespaceStates <| \m ->
        case Map.lookup ns m of
          Just (bState, _) -> (Map.insert ns (bState, output) m, ())
          Nothing -> (m, ())
      redrawLine mgr ns

-- | Update build state for a namespace
updateLineState :: Namespace -> BuildState -> IO ()
updateLineState ns buildState = do
  mMgr <- readIORef currentLineManager
  case mMgr of
    Nothing -> pure ()
    Just mgr -> do
      -- Update state and redraw just this line
      atomicModifyIORef' namespaceStates <| \m ->
        (Map.insert ns (buildState, "") m, ())
      redrawLine mgr ns

-- | Redraw a single line efficiently
redrawLine :: LineManager -> Namespace -> IO ()
redrawLine LineManager {..} ns = do
  states <- readIORef namespaceStates
  lineMap <- readIORef namespaceLines

  case Map.lookup ns lineMap of
    Nothing -> pure ()
    Just lineNum -> do
      let (bState, progress) = Map.findWithDefault (Pending, "") ns states
      withMVar terminalLock <| \_ -> do
        -- Restore to top of our area, then move down to the right line
        ANSI.hRestoreCursor IO.stderr
        when (lineNum > 0) <| ANSI.hCursorDown IO.stderr lineNum
        -- Clear and redraw this line
        ANSI.hClearLine IO.stderr
        drawNamespaceLine lmTermInfo ns bState progress
        IO.hFlush IO.stderr

-- | Redraw all lines (used only at initialization)
redrawAll :: LineManager -> IO ()
redrawAll LineManager {..} = do
  states <- readIORef namespaceStates

  withMVar terminalLock <| \_ -> do
    let numLines = length lmNamespaces
    when (numLines > 0) <| do
      -- Restore to saved position (top of our area)
      ANSI.hRestoreCursor IO.stderr
      -- Clear from here to end of screen
      ANSI.hClearFromCursorToScreenEnd IO.stderr

      -- Redraw each line
      forM_ lmNamespaces <| \ns -> do
        let (bState, progress) = Map.findWithDefault (Pending, "") ns states
        drawNamespaceLine lmTermInfo ns bState progress
        IO.hPutStrLn IO.stderr ""

      IO.hFlush IO.stderr

-- | Draw a single namespace line (without newline)
drawNamespaceLine :: TerminalInfo -> Namespace -> BuildState -> Text -> IO ()
drawNamespaceLine termInfo ns bState progress = do
  let nsText = Text.pack (Namespace.toPath ns)
  let (symbol, mColor) = stateSymbol bState
  let msg = case bState of
        Success -> symbol <> " " <> nsText
        Failed -> symbol <> " " <> nsText
        Skipped -> symbol <> " " <> nsText
        _
          | Text.null progress -> symbol <> " " <> nsText
          | otherwise -> symbol <> " " <> nsText <> ": " <> progress
  let truncated = truncateToWidth (tiWidth termInfo - 1) msg

  case mColor of
    Just color -> Rainbow.hPutChunks IO.stderr [fore color <| chunk truncated]
    Nothing -> IO.hPutStr IO.stderr (Text.unpack truncated)

stateSymbol :: BuildState -> (Text, Maybe Rainbow.Radiant)
stateSymbol = \case
  Success -> ("*", Just green)
  Failed -> ("x", Just red)
  Skipped -> ("_", Just yellow)
  Analyzing -> ("+", Nothing)
  Pending -> (".", Nothing)
  Building -> ("~", Nothing)