-- | 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 Control.Monad.IO.Class (MonadIO (liftIO))
import Data.IORef (IORef, readIORef, writeIORef)
import qualified Data.Map as M
import Data.Maybe (fromMaybe)
import qualified Data.Text 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 Futhark.Util (debug)
import Language.Futhark.Warnings (listWarnings)
import Language.LSP.Server (LspT, flushDiagnosticsBySource, getVirtualFile, getVirtualFiles)
import Language.LSP.Types
  ( filePathToUri,
    fromNormalizedFilePath,
    toNormalizedUri,
    uriToNormalizedFilePath,
  )
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
              debug $ "File not part of program: " <> show file_path'
              debug $ "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
  debug "(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
      debug "Failed to (re)-compile, using old state or Nothing"
      debug $ "Computing PositionMapping for: " <> show file_path
      mapping <- computeMapping old_state file_path
      liftIO $ writeIORef state_mvar $ updateStaleMapping file_path mapping old_state
    Just _ -> do
      debug "(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
  debug $ "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
      debug "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
    (vfsMap vfs)

getLoadedProg :: State -> LoadedProg
getLoadedProg state = fromMaybe noLoadedProg (stateProgram state)