{-# LANGUAGE LambdaCase #-}

-- | Low-level compilation parts.  Look at "Futhark.Compiler" for a
-- more high-level API.
module Futhark.Compiler.Program
  ( readLibrary,
    readUntypedLibrary,
    Imports,
    FileModule (..),
    E.Warnings,
    prettyWarnings,
    ProgError (..),
    LoadedProg (lpNameSource),
    noLoadedProg,
    lpImports,
    lpWarnings,
    lpFilePaths,
    reloadProg,
    extendProg,
    VFS,
  )
where

import Control.Concurrent (forkIO)
import Control.Concurrent.MVar
  ( MVar,
    modifyMVar,
    newEmptyMVar,
    newMVar,
    putMVar,
    readMVar,
  )
import Control.Monad
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad.State (execStateT, gets, modify)
import Data.Bifunctor (first)
import Data.List (intercalate, sort)
import Data.List.NonEmpty qualified as NE
import Data.Loc (Loc (..), Located, locOf)
import Data.Map qualified as M
import Data.Maybe (mapMaybe)
import Data.Text qualified as T
import Data.Text.IO qualified as T
import Data.Time.Clock (UTCTime, getCurrentTime)
import Futhark.FreshNames
import Futhark.Util (interactWithFileSafely, nubOrd, startupTime)
import Futhark.Util.Pretty (Doc, align, pretty)
import Language.Futhark qualified as E
import Language.Futhark.Parser (SyntaxError (..), parseFuthark)
import Language.Futhark.Prelude
import Language.Futhark.Prop (isBuiltin)
import Language.Futhark.Semantic
import Language.Futhark.TypeChecker qualified as E
import Language.Futhark.Warnings
import System.Directory (getModificationTime)
import System.FilePath (normalise, takeExtension)
import System.FilePath.Posix qualified as Posix

data LoadedFile fm = LoadedFile
  { lfPath :: FilePath,
    lfImportName :: ImportName,
    lfMod :: fm,
    -- | Modification time of the underlying file.
    lfModTime :: UTCTime
  }
  deriving (Eq, Ord, Show)

-- | Note that the location may be 'NoLoc'.  This essentially only
-- happens when the problem is that a root file cannot be found.
data ProgError
  = ProgError Loc (Doc ())
  | -- | Not actually an error, but we want them reported
    -- with errors.
    ProgWarning Loc (Doc ())

type WithErrors = Either (NE.NonEmpty ProgError)

instance Located ProgError where
  locOf (ProgError l _) = l
  locOf (ProgWarning l _) = l

-- | A mapping from absolute pathnames to pretty representing a virtual
-- file system.  Before loading a file from the file system, this
-- mapping is consulted.  If the desired pathname has an entry here,
-- the corresponding pretty is used instead of loading the file from
-- disk.
type VFS = M.Map FilePath T.Text

newtype UncheckedImport = UncheckedImport
  { unChecked ::
      WithErrors (LoadedFile E.UncheckedProg, [((ImportName, Loc), MVar UncheckedImport)])
  }

-- | If mapped to Nothing, treat it as present.  This is used when
-- reloading programs.
type ReaderState = MVar (M.Map ImportName (Maybe (MVar UncheckedImport)))

newState :: [ImportName] -> IO ReaderState
newState known = newMVar $ M.fromList $ zip known $ repeat Nothing

orderedImports ::
  [((ImportName, Loc), MVar UncheckedImport)] ->
  IO [(ImportName, WithErrors (LoadedFile E.UncheckedProg))]
orderedImports = fmap reverse . flip execStateT [] . mapM_ (spelunk [])
  where
    spelunk steps ((include, loc), mvar)
      | include `elem` steps = do
          let problem =
                ProgError loc . pretty $
                  "Import cycle: "
                    <> intercalate
                      " -> "
                      (map includeToString $ reverse $ include : steps)
          modify ((include, Left (NE.singleton problem)) :)
      | otherwise = do
          prev <- gets $ lookup include
          case prev of
            Just _ -> pure ()
            Nothing -> do
              res <- unChecked <$> liftIO (readMVar mvar)
              case res of
                Left errors ->
                  modify ((include, Left errors) :)
                Right (file, more_imports) -> do
                  mapM_ (spelunk (include : steps)) more_imports
                  modify ((include, Right file) :)

errorsToTop ::
  [(ImportName, WithErrors (LoadedFile E.UncheckedProg))] ->
  WithErrors [(ImportName, LoadedFile E.UncheckedProg)]
errorsToTop [] = Right []
errorsToTop ((_, Left x) : rest) =
  either (Left . (x <>)) (const (Left x)) (errorsToTop rest)
errorsToTop ((name, Right x) : rest) =
  fmap ((name, x) :) (errorsToTop rest)

newImportMVar :: IO UncheckedImport -> IO (Maybe (MVar UncheckedImport))
newImportMVar m = do
  mvar <- newEmptyMVar
  void $ forkIO $ putMVar mvar =<< m
  pure $ Just mvar

-- | Read the content and modification time of a file.
-- Check if the file exits in VFS before interact with file system directly.
contentsAndModTime :: FilePath -> VFS -> IO (Maybe (Either String (T.Text, UTCTime)))
contentsAndModTime filepath vfs = do
  case M.lookup filepath vfs of
    Nothing -> interactWithFileSafely $ (,) <$> T.readFile filepath <*> getModificationTime filepath
    Just file_contents -> do
      now <- getCurrentTime
      pure $ Just $ Right (file_contents, now)

readImportFile :: ImportName -> Loc -> VFS -> IO (Either ProgError (LoadedFile T.Text))
readImportFile include loc vfs = do
  -- First we try to find a file of the given name in the search path,
  -- then we look at the builtin library if we have to.  For the
  -- builtins, we don't use the search path.
  let filepath = includeToFilePath include
  r <- contentsAndModTime filepath vfs
  case (r, lookup prelude_str prelude) of
    (Just (Right (s, mod_time)), _) ->
      pure $ Right $ loaded filepath s mod_time
    (Just (Left e), _) ->
      pure $ Left $ ProgError loc $ pretty e
    (Nothing, Just s) ->
      pure $ Right $ loaded prelude_str s startupTime
    (Nothing, Nothing) ->
      pure $ Left $ ProgError loc $ pretty not_found
  where
    prelude_str = "/" Posix.</> includeToString include Posix.<.> "fut"

    loaded path s mod_time =
      LoadedFile
        { lfImportName = include,
          lfPath = path,
          lfMod = s,
          lfModTime = mod_time
        }

    not_found =
      "Could not find import " <> E.quote (includeToText include) <> "."

handleFile :: ReaderState -> VFS -> LoadedFile T.Text -> IO UncheckedImport
handleFile state_mvar vfs (LoadedFile file_name import_name file_contents mod_time) = do
  case parseFuthark file_name file_contents of
    Left (SyntaxError loc err) ->
      pure . UncheckedImport . Left . NE.singleton $ ProgError loc $ pretty err
    Right prog -> do
      let imports = map (first $ mkImportFrom import_name) $ E.progImports prog
      mvars <-
        mapMaybe sequenceA . zip imports
          <$> mapM (uncurry $ readImport state_mvar vfs) imports
      let file =
            LoadedFile
              { lfPath = file_name,
                lfImportName = import_name,
                lfModTime = mod_time,
                lfMod = prog
              }
      pure $ UncheckedImport $ Right (file, mvars)

readImport :: ReaderState -> VFS -> ImportName -> Loc -> IO (Maybe (MVar UncheckedImport))
readImport state_mvar vfs include loc =
  modifyMVar state_mvar $ \state ->
    case M.lookup include state of
      Just x -> pure (state, x)
      Nothing -> do
        prog_mvar <- newImportMVar $ do
          readImportFile include loc vfs >>= \case
            Left e -> pure $ UncheckedImport $ Left $ NE.singleton e
            Right file -> handleFile state_mvar vfs file
        pure (M.insert include prog_mvar state, prog_mvar)

readUntypedLibraryExceptKnown ::
  [ImportName] ->
  VFS ->
  [FilePath] ->
  IO (Either (NE.NonEmpty ProgError) [LoadedFile E.UncheckedProg])
readUntypedLibraryExceptKnown known vfs fps = do
  state_mvar <- liftIO $ newState known
  let prelude_import = mkInitialImport "/prelude/prelude"
  prelude_mvar <- liftIO $ readImport state_mvar vfs prelude_import mempty
  fps_mvars <- liftIO (mapM (onFile state_mvar) fps)
  let unknown_mvars = onlyUnknown (((prelude_import, mempty), prelude_mvar) : fps_mvars)
  fmap (map snd) . errorsToTop <$> orderedImports unknown_mvars
  where
    onlyUnknown = mapMaybe sequenceA
    onFile state_mvar fp =
      modifyMVar state_mvar $ \state -> do
        case M.lookup include state of
          Just prog_mvar -> pure (state, ((include, mempty), prog_mvar))
          Nothing -> do
            prog_mvar <- newImportMVar $ do
              if takeExtension fp /= ".fut"
                then
                  pure . UncheckedImport . Left . NE.singleton $
                    ProgError NoLoc $
                      pretty fp <> ": source files must have a .fut extension."
                else do
                  r <- contentsAndModTime fp vfs
                  case r of
                    Just (Right (fs, mod_time)) -> do
                      handleFile state_mvar vfs $
                        LoadedFile
                          { lfImportName = include,
                            lfMod = fs,
                            lfModTime = mod_time,
                            lfPath = fp
                          }
                    Just (Left e) ->
                      pure . UncheckedImport . Left . NE.singleton $
                        ProgError NoLoc $
                          pretty $
                            show e
                    Nothing ->
                      pure . UncheckedImport . Left . NE.singleton $
                        ProgError NoLoc $
                          pretty fp <> ": file not found."
            pure (M.insert include prog_mvar state, ((include, mempty), prog_mvar))
      where
        include = mkInitialImport fp_name
        (fp_name, _) = Posix.splitExtension fp

-- | A type-checked file.
data CheckedFile = CheckedFile
  { -- | The name generation state after checking this file.
    cfNameSource :: VNameSource,
    -- | The warnings that were issued from checking this file.
    cfWarnings :: Warnings,
    -- | The type-checked file.
    cfMod :: FileModule
  }

asImports :: [LoadedFile CheckedFile] -> Imports
asImports = map f
  where
    f lf = (lfImportName lf, cfMod $ lfMod lf)

typeCheckProg ::
  [LoadedFile CheckedFile] ->
  VNameSource ->
  [LoadedFile E.UncheckedProg] ->
  WithErrors ([LoadedFile CheckedFile], VNameSource)
typeCheckProg orig_imports orig_src =
  foldM f (orig_imports, orig_src)
  where
    roots = ["/prelude/prelude"]

    f (imports, src) (LoadedFile path import_name prog mod_time) = do
      let prog'
            | isBuiltin (includeToFilePath import_name) = prog
            | otherwise = prependRoots roots prog
      case E.checkProg (asImports imports) src import_name prog' of
        (prog_ws, Left (E.TypeError loc notes msg)) -> do
          let err' = msg <> pretty notes
              warningToError (wloc, wmsg) = ProgWarning (locOf wloc) wmsg
          Left $
            ProgError (locOf loc) err'
              NE.:| map warningToError (listWarnings prog_ws)
        (prog_ws, Right (m, src')) ->
          let warnHole (loc, t) =
                singleWarning (E.srclocOf loc) $ "Hole of type: " <> align (pretty t)
              prog_ws' = prog_ws <> foldMap warnHole (E.progHoles (fileProg m))
           in Right
                ( imports ++ [LoadedFile path import_name (CheckedFile src prog_ws' m) mod_time],
                  src'
                )

setEntryPoints ::
  [E.Name] ->
  [FilePath] ->
  [LoadedFile E.UncheckedProg] ->
  [LoadedFile E.UncheckedProg]
setEntryPoints extra_eps fps = map onFile
  where
    fps' = map normalise fps
    onFile lf
      | includeToFilePath (lfImportName lf) `elem` fps' =
          lf {lfMod = prog {E.progDecs = map onDec (E.progDecs prog)}}
      | otherwise =
          lf
      where
        prog = lfMod lf

    onDec (E.ValDec vb)
      | E.valBindName vb `elem` extra_eps =
          E.ValDec vb {E.valBindEntryPoint = Just E.NoInfo}
    onDec dec = dec

prependRoots :: [FilePath] -> E.UncheckedProg -> E.UncheckedProg
prependRoots roots (E.Prog doc ds) =
  E.Prog doc $ map mkImport roots ++ ds
  where
    mkImport fp =
      -- We do not use ImportDec here, because we do not want the
      -- type checker to issue a warning about a redundant import.
      E.LocalDec (E.OpenDec (E.ModImport fp E.NoInfo mempty) mempty) mempty

-- | A loaded, type-checked program.  This can be used to extract
-- information about the program, but also to speed up subsequent
-- reloads.
data LoadedProg = LoadedProg
  { lpRoots :: [FilePath],
    -- | The 'VNameSource' is the name source just *before* the module
    -- was type checked.
    lpFiles :: [LoadedFile CheckedFile],
    -- | Final name source.
    lpNameSource :: VNameSource
  }

-- | The 'Imports' of a 'LoadedProg', as expected by e.g. type
-- checking functions.
lpImports :: LoadedProg -> Imports
lpImports = map f . lpFiles
  where
    f lf = (lfImportName lf, cfMod $ lfMod lf)

-- | All warnings of a 'LoadedProg'.
lpWarnings :: LoadedProg -> Warnings
lpWarnings = foldMap (cfWarnings . lfMod) . lpFiles

-- | The absolute paths of the files that are part of this program.
lpFilePaths :: LoadedProg -> [FilePath]
lpFilePaths = map lfPath . lpFiles

unchangedImports ::
  MonadIO m =>
  VNameSource ->
  VFS ->
  [LoadedFile CheckedFile] ->
  m ([LoadedFile CheckedFile], VNameSource)
unchangedImports src _ [] = pure ([], src)
unchangedImports src vfs (f : fs)
  | isBuiltin (includeToFilePath (lfImportName f)) =
      first (f :) <$> unchangedImports src vfs fs
  | otherwise = do
      let file_path = lfPath f
      if M.member file_path vfs
        then pure ([], cfNameSource $ lfMod f)
        else do
          changed <-
            maybe True (either (const True) (> lfModTime f))
              <$> liftIO (interactWithFileSafely (getModificationTime file_path))
          if changed
            then pure ([], cfNameSource $ lfMod f)
            else first (f :) <$> unchangedImports src vfs fs

-- | A "loaded program" containing no actual files.  Use this as a
-- starting point for 'reloadProg'
noLoadedProg :: LoadedProg
noLoadedProg =
  LoadedProg
    { lpRoots = [],
      lpFiles = mempty,
      lpNameSource = newNameSource $ E.maxIntrinsicTag + 1
    }

-- | Find out how many of the old imports can be used.  Here we are
-- forced to be overly conservative, because our type checker
-- enforces a linear ordering.
usableLoadedProg :: MonadIO m => LoadedProg -> VFS -> [FilePath] -> m LoadedProg
usableLoadedProg (LoadedProg roots imports src) vfs new_roots
  | sort roots == sort new_roots = do
      (imports', src') <- unchangedImports src vfs imports
      pure $ LoadedProg [] imports' src'
  | otherwise =
      pure noLoadedProg

-- | Extend a loaded program with (possibly new) files.
extendProg ::
  LoadedProg ->
  [FilePath] ->
  VFS ->
  IO (Either (NE.NonEmpty ProgError) LoadedProg)
extendProg lp new_roots vfs = do
  new_imports_untyped <-
    readUntypedLibraryExceptKnown (map lfImportName $ lpFiles lp) vfs new_roots
  pure $ do
    (imports, src') <-
      typeCheckProg (lpFiles lp) (lpNameSource lp) =<< new_imports_untyped
    Right (LoadedProg (nubOrd (lpRoots lp ++ new_roots)) imports src')

-- | Load some new files, reusing as much of the previously loaded
-- program as possible.  This does not *extend* the currently loaded
-- program the way 'extendProg' does it, so it is always correct (if
-- less efficient) to pass 'noLoadedProg'.
reloadProg ::
  LoadedProg ->
  [FilePath] ->
  VFS ->
  IO (Either (NE.NonEmpty ProgError) LoadedProg)
reloadProg lp new_roots vfs = do
  lp' <- usableLoadedProg lp vfs new_roots
  extendProg lp' new_roots vfs

-- | Read and type-check some Futhark files.
readLibrary ::
  -- | Extra functions that should be marked as entry points; only
  -- applies to the immediate files, not any imports imported.
  [E.Name] ->
  -- | The files to read.
  [FilePath] ->
  IO (Either (NE.NonEmpty ProgError) (E.Warnings, Imports, VNameSource))
readLibrary extra_eps fps =
  ( fmap frob
      . typeCheckProg mempty (lpNameSource noLoadedProg)
      <=< fmap (setEntryPoints (E.defaultEntryPoint : extra_eps) fps)
  )
    <$> readUntypedLibraryExceptKnown [] M.empty fps
  where
    frob (y, z) = (foldMap (cfWarnings . lfMod) y, asImports y, z)

-- | Read (and parse) all source files (including the builtin prelude)
-- corresponding to a set of root files.
readUntypedLibrary ::
  [FilePath] ->
  IO (Either (NE.NonEmpty ProgError) [(ImportName, E.UncheckedProg)])
readUntypedLibrary = fmap (fmap (map f)) . readUntypedLibraryExceptKnown [] M.empty
  where
    f lf = (lfImportName lf, lfMod lf)