{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TemplateHaskell       #-}
{-# LANGUAGE TypeFamilies          #-}
-- | This module implements import resolution.
module KOSC.Compiler.ImportResolution
  ( ImportResolution (..)
  , ImportedModuleInfo (..)
  , resolveImports
    -- * Lenses
  , importResolutionMainModule
  , importResolutionModules
  , moduleInfoAST
  , moduleInfoExportedTypes
  , moduleInfoExportedVars
  , moduleInfoImports
  ) where

import           Control.Lens
import           Control.Monad.Except
import           Control.Monad.State
import           Data.List
import           Data.Map.Strict              (Map)
import qualified Data.Map.Strict              as Map
import           Data.Set                     (Set)
import qualified Data.Set                     as Set

import           KOSC.Compiler.Common
import qualified KOSC.Language.AST            as AST


-- | Results of running the import resolver.
data ImportResolution = ImportResolution
  { _importResolutionMainModule :: AST.ModuleName
  -- ^ name of the main module
  , _importResolutionModules    :: Map AST.ModuleName ImportedModuleInfo
  -- ^ map of all transitively imported modules to some additional information
  } deriving (Show)

-- | Information generated by the import resolver for each module
data ImportedModuleInfo = ImportedModuleInfo
  { _moduleInfoAST           :: AST.RawModule
  -- ^ raw AST of the imported module
  , _moduleInfoExportedTypes :: Set AST.Ident
  -- ^ names of the exported types
  , _moduleInfoExportedVars  :: Set AST.Ident
  -- ^ names of the exported variables and functions
  , _moduleInfoImports       :: [AST.ImportDecl]
  -- ^ import declarations of this module
  } deriving (Show)

makeLenses ''ImportResolution
makeLenses ''ImportedModuleInfo

-- | @getImportList m@ returns the list of import declarations in @m@.
getImportList :: AST.Module name -> [AST.ImportDecl]
getImportList mod = [ decl | AST.DeclImport decl <- view AST.declarations mod ]

-- | @getExportedTermNames m@ returns a set of exported term-level names, i.e. functions and global variables.
getExportedTermNames :: AST.Module name -> [AST.Ident]
getExportedTermNames mod = [ name | decl <- view AST.declarations mod, name <- getname decl ] where
  getname (AST.DeclImport _) = []
  getname (AST.DeclFun fd) = [fd ^. AST.funDeclSignature . AST.funSigName]
  getname (AST.DeclVar vd) = [vd ^. AST.varDeclSignature . AST.varSigName]
  getname (AST.DeclRec _) = []
  getname (AST.DeclBuiltin (AST.BuiltinStruct _)) = [] -- struct is only a type name
  getname (AST.DeclBuiltin (AST.BuiltinFun fsig)) = [ fsig ^. AST.funSigName ]
  getname (AST.DeclBuiltin (AST.BuiltinVar vsig)) = [ vsig ^. AST.varSigName ]

-- | @getExportedTermNames m@ returns a set of exported type-level names (for now nothing, but will be records)
getExportedTypeNames :: AST.Module name -> [AST.Ident]
getExportedTypeNames mod = [ name | decl <- view AST.declarations mod, name <- getname decl ] where
  getname (AST.DeclImport _) = []
  getname (AST.DeclFun _) = []
  getname (AST.DeclVar _) = []
  getname (AST.DeclRec r) = [r ^. AST.recDeclName]
  getname (AST.DeclBuiltin (AST.BuiltinStruct ssig)) = [ ssig ^. AST.structSigName ]
  getname (AST.DeclBuiltin (AST.BuiltinFun _)) = [ ]
  getname (AST.DeclBuiltin (AST.BuiltinVar _)) = [ ]

-- | Transitively resolves all imports, starting at the given main module.
resolveImports :: Monad m
               => (AST.ModuleName -> KOSCCompilerT m AST.RawModule) -- ^ function loading a module
               -> AST.RawModule                                     -- ^ main module where the search for imports is started
               -> KOSCCompilerT m ImportResolution
resolveImports loadModule mainModule = execStateT (go [mainModuleName]) initialState where
  mainModuleName = view AST.moduleName mainModule

  initialState = ImportResolution
    { _importResolutionMainModule = mainModule ^. AST.moduleName
    , _importResolutionModules = Map.empty
    }

  go [] = return ()
  go (name:rest) = use (importResolutionModules . at name) >>= \case
    Nothing -> do
      imports <- enterModule name $ do
        mod <- if name == mainModuleName then pure mainModule else lift $ loadModule name
        let expTerms = getExportedTermNames mod
            expTypes = getExportedTypeNames mod
            expTermSet = Set.fromList expTerms
            expTypeSet = Set.fromList expTypes
            duplTerms = expTerms \\ Set.toList expTermSet
            duplTypes = expTypes \\ Set.toList expTypeSet
        forM duplTerms $ \dupl -> lift $ messageWithContext MessageError (MessageDuplicateDeclaration dupl)
        forM duplTypes $ \dupl -> lift $ messageWithContext MessageError (MessageDuplicateDeclaration dupl)
        let imports = getImportList mod
        importResolutionModules . at name .= Just (ImportedModuleInfo mod expTypeSet expTermSet imports)
        return imports
      go $ map (view AST.importModuleName) imports ++ rest
    Just _ -> go rest -- already imported elsewhere
