{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE UndecidableInstances #-}

-----------------------------------------------------------------------------

-----------------------------------------------------------------------------

-- |
-- Module      :  Disco.Module
-- Copyright   :  (c) 2019 disco team (see LICENSE)
-- Maintainer  :  byorgey@gmail.com
--
-- SPDX-License-Identifier: BSD-3-Clause
--
-- The 'ModuleInfo' record representing a disco module, and functions
-- to resolve the location of a module on disk.
module Disco.Module where

import Data.Data (Data)
import GHC.Generics (Generic)

import Control.Lens (
  Getting,
  foldOf,
  makeLenses,
  view,
 )
import Control.Monad (filterM)
import Control.Monad.IO.Class (MonadIO (..))
import Data.Bifunctor (first)
import Data.Map (Map)
import qualified Data.Map as M
import Data.Maybe (listToMaybe)
import qualified Data.Set as S
import System.Directory (doesFileExist)
import System.FilePath (
  replaceExtension,
  (</>),
 )

import Unbound.Generics.LocallyNameless (
  Alpha,
  Bind,
  Name,
  Subst,
  bind,
 )
import Unbound.Generics.LocallyNameless.Unsafe (unsafeUnbind)

import Polysemy

import Disco.AST.Surface
import Disco.AST.Typed
import Disco.Context
import Disco.Extensions
import Disco.Names
import Disco.Pretty hiding ((<>))
import Disco.Typecheck.Erase (erase, erasePattern)
import Disco.Typecheck.Util (TyCtx)
import Disco.Types

import Paths_disco

------------------------------------------------------------
-- ModuleInfo and related types
------------------------------------------------------------

-- | When loading a module, we could be loading it from code entered
-- at the REPL, or from a standalone file.  The two modes have
-- slightly different behavior.
data LoadingMode = REPL | Standalone

-- | A definition consists of a name being defined, the types of any
--   pattern arguments (each clause must have the same number of
--   patterns), the type of the body of each clause, and a list of
--   clauses.  For example,
--
--   @
--   f x (0,z) = 3*x + z > 5
--   f x (y,z) = z == 9
--   @
--
--   might look like @Defn f [Z, Z*Z] B [clause 1 ..., clause 2 ...]@
data Defn = Defn (Name ATerm) [Type] Type [Clause]
  deriving (Show, Generic, Alpha, Data, Subst Type)

instance Pretty Defn where
  pretty (Defn x patTys ty clauses) =
    vcat $
      prettyTyDecl x (foldr (:->:) ty patTys)
        : map (pretty . (x,) . eraseClause) clauses

-- | A clause in a definition consists of a list of patterns (the LHS
--   of the =) and a term (the RHS).  For example, given the concrete
--   syntax @f n (x,y) = n*x + y@, the corresponding 'Clause' would be
--   something like @[n, (x,y)] (n*x + y)@.
type Clause = Bind [APattern] ATerm

eraseClause :: Clause -> Bind [Pattern] Term
eraseClause b = bind (map erasePattern ps) (erase t)
 where
  (ps, t) = unsafeUnbind b

-- | Type checking a module yields a value of type ModuleInfo which contains
--   mapping from terms to their relavent documenation, a mapping from terms to
--   properties, and a mapping from terms to their types.
data ModuleInfo = ModuleInfo
  { _miName :: ModuleName
  , _miImports :: Map ModuleName ModuleInfo
  , -- List of names declared by the module, in the order they occur
    _miNames :: [QName Term]
  , _miDocs :: Ctx Term Docs
  , _miProps :: Ctx ATerm [AProperty]
  , _miTys :: TyCtx
  , _miTydefs :: TyDefCtx
  , _miTermdefs :: Ctx ATerm Defn
  , _miTerms :: [(ATerm, PolyType)]
  , _miExts :: ExtSet
  }
  deriving (Show)

makeLenses ''ModuleInfo

instance Semigroup ModuleInfo where
  -- \| Two ModuleInfos
  --   are merged by joining their doc, type, type definition, and term
  --   contexts. The property context of the new module is the one
  --   obtained from the second module. The name of the new module is
  --   taken from the first. Definitions from later modules override
  --   earlier ones.  Note that this function should really only be used
  --   for the special top-level REPL module.
  ModuleInfo n1 is1 ns1 d1 _ ty1 tyd1 tm1 tms1 es1
    <> ModuleInfo _ is2 ns2 d2 p2 ty2 tyd2 tm2 tms2 es2 =
      ModuleInfo
        n1
        (is1 <> is2)
        (ns1 <> ns2)
        (d2 <> d1)
        p2
        (ty2 <> ty1)
        (tyd2 <> tyd1)
        (tm2 <> tm1)
        (tms1 <> tms2)
        (es1 <> es2)

instance Monoid ModuleInfo where
  mempty = emptyModuleInfo
  mappend = (<>)

-- | Get something from a module and its direct imports.
withImports :: Monoid a => Getting a ModuleInfo a -> ModuleInfo -> a
withImports l = view l <> foldOf (miImports . traverse . l)

-- | Get the types of all names bound in a module and its direct imports.
allTys :: ModuleInfo -> TyCtx
allTys = withImports miTys

-- | Get all type definitions from a module and its direct imports.
allTydefs :: ModuleInfo -> TyDefCtx
allTydefs = withImports miTydefs

-- | The empty module info record.
emptyModuleInfo :: ModuleInfo
emptyModuleInfo = ModuleInfo REPLModule M.empty [] emptyCtx emptyCtx emptyCtx M.empty emptyCtx [] S.empty

------------------------------------------------------------
-- Module resolution
------------------------------------------------------------

-- | A data type indicating where we should look for Disco modules to
--   be loaded.
data Resolver
  = -- | Load only from the stdlib               (standard lib modules)
    FromStdlib
  | -- | Load only from a specific directory     (:load)
    FromDir FilePath
  | -- | Load from current working dir or stdlib (import at REPL)
    FromCwdOrStdlib
  | -- | Load from specific dir or stdlib        (import in file)
    FromDirOrStdlib FilePath

-- | Add the possibility of loading imports from the stdlib.  For
--   example, this is what we want to do after a user loads a specific
--   file using `:load` (for which we will NOT look in the stdlib),
--   but then we need to recursively load modules which it imports
--   (which may either be in the stdlib, or the same directory as the
--   `:load`ed module).
withStdlib :: Resolver -> Resolver
withStdlib (FromDir fp) = FromDirOrStdlib fp
withStdlib r = r

-- | Given a module resolution mode and a raw module name, relavent
--   directories are searched for the file containing the provided
--   module name.  Returns Nothing if no module with the given name
--   could be found.
resolveModule :: Member (Embed IO) r => Resolver -> String -> Sem r (Maybe (FilePath, ModuleProvenance))
resolveModule resolver modname = do
  datadir <- liftIO getDataDir
  let searchPath =
        case resolver of
          FromStdlib -> [(datadir, Stdlib)]
          FromDir dir -> [(dir, Dir dir)]
          FromCwdOrStdlib -> [(datadir, Stdlib), (".", Dir ".")]
          FromDirOrStdlib dir -> [(datadir, Stdlib), (dir, Dir dir)]
  let fps = map (first (</> replaceExtension modname "disco")) searchPath
  fexists <- liftIO $ filterM (doesFileExist . fst) fps
  return $ listToMaybe fexists