module Language.PureScript.Sugar.Names.Imports
  ( ImportDef
  , resolveImports
  , resolveModuleImport
  , findImports
  ) where

import Prelude

import Control.Monad (foldM, when)
import Control.Monad.Error.Class (MonadError(..))

import Data.Foldable (for_, traverse_)
import Data.Maybe (fromMaybe)
import Data.Map qualified as M
import Data.Set qualified as S

import Language.PureScript.AST (Declaration(..), DeclarationRef(..), ErrorMessageHint(..), ExportSource(..), ImportDeclarationType(..), Module(..), SourceSpan, internalModuleSourceSpan)
import Language.PureScript.Crash (internalError)
import Language.PureScript.Errors (MultipleErrors, SimpleErrorMessage(..), addHint, errorMessage', rethrow)
import Language.PureScript.Names (pattern ByNullSourcePos, ModuleName, Name(..), ProperName, ProperNameType(..), Qualified(..), QualifiedBy(..), byMaybeModuleName)
import Language.PureScript.Sugar.Names.Env (Env, Exports(..), ImportProvenance(..), ImportRecord(..), Imports(..), envModuleExports, nullImports)

type ImportDef = (SourceSpan, ImportDeclarationType, Maybe ModuleName)

-- |
-- Finds the imports within a module, mapping the imported module name to an optional set of
-- explicitly imported declarations.
--
findImports
  :: [Declaration]
  -> M.Map ModuleName [ImportDef]
findImports = foldr go M.empty
  where
  go (ImportDeclaration (pos, _) mn typ qual) =
    M.alter (return . ((pos, typ, qual) :) . fromMaybe []) mn
  go _ = id

-- |
-- Constructs a set of imports for a module.
--
resolveImports
  :: forall m
   . MonadError MultipleErrors m
  => Env
  -> Module
  -> m (Module, Imports)
resolveImports env (Module ss coms currentModule decls exps) =
  rethrow (addHint (ErrorInModule currentModule)) $ do
    let imports = findImports decls
        imports' = M.map (map (\(ss', dt, mmn) -> (ss', Just dt, mmn))) imports
        scope = M.insert currentModule [(internalModuleSourceSpan "<module>", Nothing, Nothing)] imports'
    (Module ss coms currentModule decls exps,) <$>
      foldM (resolveModuleImport env) nullImports (M.toList scope)

-- | Constructs a set of imports for a single module import.
resolveModuleImport
  :: forall m
   . MonadError MultipleErrors m
  => Env
  -> Imports
  -> (ModuleName, [(SourceSpan, Maybe ImportDeclarationType, Maybe ModuleName)])
  -> m Imports
resolveModuleImport env ie (mn, imps) = foldM go ie imps
  where
  go :: Imports
     -> (SourceSpan, Maybe ImportDeclarationType, Maybe ModuleName)
     -> m Imports
  go ie' (ss, typ, impQual) = do
    modExports <-
      maybe
        (throwError . errorMessage' ss . UnknownName . Qualified ByNullSourcePos $ ModName mn)
        (return . envModuleExports)
        (mn `M.lookup` env)
    let impModules = importedModules ie'
        qualModules = importedQualModules ie'
        ie'' = ie' { importedModules = maybe (S.insert mn impModules) (const impModules) impQual
                   , importedQualModules = maybe qualModules (`S.insert` qualModules) impQual
                   }
    resolveImport mn modExports ie'' impQual ss typ

-- |
-- Extends the local environment for a module by resolving an import of another module.
--
resolveImport
  :: forall m
   . MonadError MultipleErrors m
  => ModuleName
  -> Exports
  -> Imports
  -> Maybe ModuleName
  -> SourceSpan
  -> Maybe ImportDeclarationType
  -> m Imports
resolveImport importModule exps imps impQual = resolveByType
  where

  resolveByType :: SourceSpan -> Maybe ImportDeclarationType -> m Imports
  resolveByType ss Nothing =
    importAll ss (importRef Local)
  resolveByType ss (Just Implicit) =
    importAll ss (importRef FromImplicit)
  resolveByType _ (Just (Explicit refs)) =
    checkRefs False refs >> foldM (importRef FromExplicit) imps refs
  resolveByType ss (Just (Hiding refs)) =
    checkRefs True refs >> importAll ss (importNonHidden refs)

  -- Check that a 'DeclarationRef' refers to an importable symbol
  checkRefs :: Bool -> [DeclarationRef] -> m ()
  checkRefs isHiding = traverse_ check
    where
    check (ValueRef ss name) =
      checkImportExists ss IdentName (exportedValues exps) name
    check (ValueOpRef ss op) =
      checkImportExists ss ValOpName (exportedValueOps exps) op
    check (TypeRef ss name dctors) = do
      checkImportExists ss TyName (exportedTypes exps) name
      let (allDctors, _) = allExportedDataConstructors name
      for_ dctors $ traverse_ (checkDctorExists ss name allDctors)
    check (TypeOpRef ss name) =
      checkImportExists ss TyOpName (exportedTypeOps exps) name
    check (TypeClassRef ss name) =
      checkImportExists ss TyClassName (exportedTypeClasses exps) name
    check (ModuleRef ss name) | isHiding =
      throwError . errorMessage' ss $ ImportHidingModule name
    check r = internalError $ "Invalid argument to checkRefs: " ++ show r

  -- Check that an explicitly imported item exists in the module it is being imported from
  checkImportExists
    :: Ord a
    => SourceSpan
    -> (a -> Name)
    -> M.Map a b
    -> a
    -> m ()
  checkImportExists ss toName exports item
    = when (item `M.notMember` exports)
    . throwError . errorMessage' ss
    $ UnknownImport importModule (toName item)

  -- Ensure that an explicitly imported data constructor exists for the type it is being imported
  -- from
  checkDctorExists
    :: SourceSpan
    -> ProperName 'TypeName
    -> [ProperName 'ConstructorName]
    -> ProperName 'ConstructorName
    -> m ()
  checkDctorExists ss tcon exports dctor
    = when (dctor `notElem` exports)
    . throwError . errorMessage' ss
    $ UnknownImportDataConstructor importModule tcon dctor

  importNonHidden :: [DeclarationRef] -> Imports -> DeclarationRef -> m Imports
  importNonHidden hidden m ref | isHidden ref = return m
                               | otherwise = importRef FromImplicit m ref
    where
    -- TODO: rework this to be not confusing
    isHidden :: DeclarationRef -> Bool
    isHidden ref'@TypeRef{} = foldl (checkTypeRef ref') False hidden
    isHidden ref' = ref' `elem` hidden
    checkTypeRef :: DeclarationRef -> Bool -> DeclarationRef -> Bool
    checkTypeRef _ True _ = True
    checkTypeRef (TypeRef _ _ Nothing) acc (TypeRef _ _ (Just _)) = acc
    checkTypeRef (TypeRef _ name (Just dctor)) _ (TypeRef _ name' (Just dctor')) = name == name' && dctor == dctor'
    checkTypeRef (TypeRef _ name _) _ (TypeRef _ name' Nothing) = name == name'
    checkTypeRef _ acc _ = acc

  -- Import all symbols
  importAll :: SourceSpan -> (Imports -> DeclarationRef -> m Imports) -> m Imports
  importAll ss importer =
    foldM (\m (name, (dctors, _)) -> importer m (TypeRef ss name (Just dctors))) imps (M.toList (exportedTypes exps))
      >>= flip (foldM (\m (name, _) -> importer m (TypeOpRef ss name))) (M.toList (exportedTypeOps exps))
      >>= flip (foldM (\m (name, _) -> importer m (ValueRef ss name))) (M.toList (exportedValues exps))
      >>= flip (foldM (\m (name, _) -> importer m (ValueOpRef ss name))) (M.toList (exportedValueOps exps))
      >>= flip (foldM (\m (name, _) -> importer m (TypeClassRef ss name))) (M.toList (exportedTypeClasses exps))

  importRef :: ImportProvenance -> Imports -> DeclarationRef -> m Imports
  importRef prov imp (ValueRef ss name) = do
    let values' = updateImports (importedValues imp) (exportedValues exps) id name ss prov
    return $ imp { importedValues = values' }
  importRef prov imp (ValueOpRef ss name) = do
    let valueOps' = updateImports (importedValueOps imp) (exportedValueOps exps) id name ss prov
    return $ imp { importedValueOps = valueOps' }
  importRef prov imp (TypeRef ss name dctors) = do
    let types' = updateImports (importedTypes imp) (exportedTypes exps) snd name ss prov
    let (dctorNames, src) = allExportedDataConstructors name
        dctorLookup :: M.Map (ProperName 'ConstructorName) ExportSource
        dctorLookup = M.fromList $ map (, src) dctorNames
    traverse_ (traverse_ $ checkDctorExists ss name dctorNames) dctors
    let dctors' = foldl (\m d -> updateImports m dctorLookup id d ss prov) (importedDataConstructors imp) (fromMaybe dctorNames dctors)
    return $ imp { importedTypes = types', importedDataConstructors = dctors' }
  importRef prov imp (TypeOpRef ss name) = do
    let ops' = updateImports (importedTypeOps imp) (exportedTypeOps exps) id name ss prov
    return $ imp { importedTypeOps = ops' }
  importRef prov imp (TypeClassRef ss name) = do
    let typeClasses' = updateImports (importedTypeClasses imp) (exportedTypeClasses exps) id name ss prov
    return $ imp { importedTypeClasses = typeClasses' }
  importRef _ _ TypeInstanceRef{} = internalError "TypeInstanceRef in importRef"
  importRef _ _ ModuleRef{} = internalError "ModuleRef in importRef"
  importRef _ _ ReExportRef{} = internalError "ReExportRef in importRef"

  -- Find all exported data constructors for a given type
  allExportedDataConstructors
    :: ProperName 'TypeName
    -> ([ProperName 'ConstructorName], ExportSource)
  allExportedDataConstructors name =
    fromMaybe (internalError "Invalid state in allExportedDataConstructors")
      $ name `M.lookup` exportedTypes exps

  -- Add something to an import resolution list
  updateImports
    :: Ord a
    => M.Map (Qualified a) [ImportRecord a]
    -> M.Map a b
    -> (b -> ExportSource)
    -> a
    -> SourceSpan
    -> ImportProvenance
    -> M.Map (Qualified a) [ImportRecord a]
  updateImports imps' exps' expName name ss prov =
    let
      src = maybe (internalError "Invalid state in updateImports") expName (name `M.lookup` exps')
      rec = ImportRecord (Qualified (ByModuleName importModule) name) (exportSourceDefinedIn src) ss prov
    in
      M.alter
        (\currNames -> Just $ rec : fromMaybe [] currNames)
        (Qualified (byMaybeModuleName impQual) name)
        imps'