content: -- | Building blocks for "recompiling" (actually just type-checking)
  -- the Futhark program managed by the language server.  The challenge
  -- here is that if the program becomes type-invalid, we want to keep
  -- the old state around.
  module Futhark.LSP.Compile (tryTakeStateFromIORef, tryReCompile) where
  
  import Colog.Core (logStringStderr, (<&))
  import Control.Lens.Getter (view)
  import Control.Monad.IO.Class (MonadIO (liftIO))
  import Data.IORef (IORef, readIORef, writeIORef)
  import Data.Map qualified as M
  import Data.Maybe (fromMaybe)
  import Data.Text qualified as T
  import Futhark.Compiler.Program (LoadedProg, lpFilePaths, lpWarnings, noLoadedProg, reloadProg)
  import Futhark.LSP.Diagnostic (diagnosticSource, maxDiagnostic, publishErrorDiagnostics, publishWarningDiagnostics)
  import Futhark.LSP.State (State (..), emptyState, updateStaleContent, updateStaleMapping)
  import Futhark.LSP.Tool (computeMapping)
  import Language.Futhark.Warnings (listWarnings)
  import Language.LSP.Protocol.Types
    ( filePathToUri,
      fromNormalizedFilePath,
      toNormalizedUri,
      uriToNormalizedFilePath,
    )
  import Language.LSP.Server (LspT, flushDiagnosticsBySource, getVirtualFile, getVirtualFiles)
  import Language.LSP.VFS (VFS, vfsMap, virtualFileText)
  
  -- | Try to take state from IORef, if it's empty, try to compile.
  tryTakeStateFromIORef :: IORef State -> Maybe FilePath -> LspT () IO State
  tryTakeStateFromIORef state_mvar file_path = do
    old_state <- liftIO $ readIORef state_mvar
    case stateProgram old_state of
      Nothing -> do
        new_state <- tryCompile old_state file_path noLoadedProg
        liftIO $ writeIORef state_mvar new_state
        pure new_state
      Just prog -> do
        -- If this is in the context of some file that is not part of
        -- the program, try to reload the program from that file.
        let files = lpFilePaths prog
        state <- case file_path of
          Just file_path'
            | file_path' `notElem` files -> do
                logStringStderr <& ("File not part of program: " <> show file_path')
                logStringStderr <& ("Program contains: " <> show files)
                tryCompile old_state file_path noLoadedProg
          _ -> pure old_state
        liftIO $ writeIORef state_mvar state
        pure state
  
  -- | Try to (re)-compile, replace old state if successful.
  tryReCompile :: IORef State -> Maybe FilePath -> LspT () IO ()
  tryReCompile state_mvar file_path = do
    logStringStderr <& "(Re)-compiling ..."
    old_state <- liftIO $ readIORef state_mvar
    let loaded_prog = getLoadedProg old_state
    new_state <- tryCompile old_state file_path loaded_prog
    case stateProgram new_state of
      Nothing -> do
        logStringStderr <& "Failed to (re)-compile, using old state or Nothing"
        logStringStderr <& "Computing PositionMapping for: " <> show file_path
        mapping <- computeMapping old_state file_path
        liftIO $ writeIORef state_mvar $ updateStaleMapping file_path mapping old_state
      Just _ -> do
        logStringStderr <& "(Re)-compile successful"
        liftIO $ writeIORef state_mvar new_state
  
  -- | Try to compile, publish diagnostics on warnings and errors, return newly compiled state.
  --  Single point where the compilation is done, and shouldn't be exported.
  tryCompile :: State -> Maybe FilePath -> LoadedProg -> LspT () IO State
  tryCompile _ Nothing _ = pure emptyState
  tryCompile state (Just path) old_loaded_prog = do
    logStringStderr <& "Reloading program from " <> show path
    vfs <- getVirtualFiles
    res <- liftIO $ reloadProg old_loaded_prog [path] (transformVFS vfs) -- NOTE: vfs only keeps track of current opened files
    flushDiagnosticsBySource maxDiagnostic diagnosticSource
    case res of
      Right new_loaded_prog -> do
        publishWarningDiagnostics $ listWarnings $ lpWarnings new_loaded_prog
        maybe_virtual_file <- getVirtualFile $ toNormalizedUri $ filePathToUri path
        case maybe_virtual_file of
          Nothing -> pure $ State (Just new_loaded_prog) (staleData state) -- should never happen
          Just virtual_file ->
            pure $ updateStaleContent path virtual_file new_loaded_prog state
      -- Preserve files that have been opened should be enoguth.
      -- But still might need an update on re-compile logic, don't discard all state afterwards,
      -- try to compile from root file, if there is a depencency relatetion, improve performance and provide more dignostic.
      Left prog_error -> do
        logStringStderr <& "Compilation failed, publishing diagnostics"
        publishErrorDiagnostics prog_error
        pure emptyState
  
  -- | Transform VFS to a map of file paths to file contents.
  -- This is used to pass the file contents to the compiler.
  transformVFS :: VFS -> M.Map FilePath T.Text
  transformVFS vfs =
    M.foldrWithKey
      ( \uri virtual_file acc ->
          case uriToNormalizedFilePath uri of
            Nothing -> acc
            Just file_path ->
              M.insert (fromNormalizedFilePath file_path) (virtualFileText virtual_file) acc
      )
      M.empty
      (view vfsMap vfs)
  
  getLoadedProg :: State -> LoadedProg
  getLoadedProg state = fromMaybe noLoadedProg (stateProgram state)
  