module Lib.ProcessModule

import Prelude
import Lib.Types
import Lib.Common
import Lib.Syntax
import Lib.Error
import Lib.ProcessDecl
import Lib.TopContext
import Lib.Tokenizer
import Data.SortedMap
import Lib.Parser.Impl
import Lib.Parser
import Data.List1
import Lib.Elab

-- declare internal primitives
addPrimitives : M ModContext
addPrimitives = do
  modifyTop { currentMod := emptyModCtx "Prim" ""; hints := emptyMap; ops := emptyMap }
  processDecl primNS (PType emptyFC "Int" Nothing)
  processDecl primNS (PType emptyFC "String" Nothing)
  processDecl primNS (PType emptyFC "Char" Nothing)
  processDecl primNS (PType emptyFC "Lazy" (Just (RPi emptyFC (BI emptyFC "_" Explicit Zero) (RU emptyFC) (RU emptyFC))))
  -- TODO we shouldn't need this anymore
  setDef (QN primNS "PiType") emptyFC Zero (Erased emptyFC) (PrimFn "(h0, h1) => ({ tag: \"PiType\", h0, h1 });" [Many,Many] Nil) Nil
  top <- getTop
  modifyTop { modules $= updateMap primNS top.currentMod }
  pure top.currentMod

record FileSource where
  baseDir : String
  getFile : FC → String → M (String × String)

parseDecls : String → Operators → TokenList → SnocList Decl → M (List Decl × Operators)
parseDecls fn ops Nil acc = pure (acc <>> Nil, ops)
parseDecls fn ops toks@(first :: _) acc =
  case partialParse fn (sameLevel parseDecl) ops toks of
    Left (err, toks) => do
      addError err
      parseDecls fn ops (recover toks) acc
    Right (decl,ops,toks) => parseDecls fn ops toks (acc :< decl)
  where
    recover : TokenList → TokenList
    recover Nil = Nil
    -- skip to top token, but make sure there is progress
    recover (tok :: toks) = if tok.bounds.startCol == 0 && tok.bounds /= first.bounds
      then (tok :: toks)
      else recover toks

importToName : Import → List String
importToName  (MkImport fc (_,name)) = splitBy name '.'

importHints : List TopEntry → M Unit
importHints Nil = pure MkUnit
importHints (entry :: entries) = do
  when (elem Hint entry.eflags) $ addHint entry.name
  importHints entries

mergeOps : Operators → Operators → Operators
mergeOps mod top = foldMap (flip const) top $ toList mod

-- processModule might reset the currentModule in the topContext
-- do not rely on topContext state afterwards - it may or may not contain the module
processModule : FC → FileSource → List String → String → M ModContext
processModule importFC repo stk modns = do
  top <- getTop

  let (Nothing) = lookupMap' modns top.modules
    | Just mod => pure mod

  let (False) = modns == primNS | _ => addPrimitives
  let freshMC = MC emptyMap Nil Nil 0 CheckAll
  let parts = splitBy modns '.'
  let fn = joinBy "/" parts ++ ".newt"

  -- Dummy for initial load/parse
  let mod = MkModCtx modns "" emptyMap freshMC emptyMap Nil Nil Nil
  modifyTop { currentMod := mod; hints := emptyMap; ops := emptyMap }

  -- TODO now we can pass in the module name...
  Right (fn,src) <- tryError $ repo.getFile importFC fn
    | Left err => reportError err  -- TODO maybe want a better FC.
  modifyTop { currentMod $= { modSource := src }}

  let (Right toks) = tokenise fn src
    | Left err => reportError err

  let (Right ((nameFC, modName), ops, toks)) = partialParse fn parseModHeader top.ops toks
    | Left (err, toks) => reportError err

  log 1 "scan imports for module \{modName}"
  let (True) = modns == modName
    | _ => reportError $ E nameFC "module name \{show modName} doesn't match file name \{show fn}"

  let (Right (imports, ops, toks)) = partialParse fn parseImports ops toks
    | Left (err, toks) => reportError err

  imported <- for imports $ \case
    MkImport fc (nameFC, name') => do
      when (elem name' stk) $ error nameFC "import loop \{modns} → \{name'}"
      mod <- processModule nameFC repo (modns :: stk) name'
      pure $ name'

  processModule nameFC repo (modns :: stk) primNS
  let imported = snoc imported primNS

  putStrLn "module \{modName}"

  -- currentMod has been wiped by imports..
  let freshMC = MC emptyMap Nil Nil 0 CheckAll
  let mod = MkModCtx modns src emptyMap freshMC emptyMap imported Nil Nil
  modifyTop { currentMod := mod
            ; hints := emptyMap
            ; ops := ops
            }

  -- top hints / ops include all directly imported modules
  top <- getTop
  for_ imports $ \case
    (MkImport fc (nameFC, ns)) => do
      let (Just mod) = lookupMap' ns top.modules | _ => error emptyFC "namespace \{show ns} missing"
      importHints (listValues mod.modDefs)
      modifyTop { ops $= mergeOps mod.modOps }

      -- add error if an import contains an error
      -- maybe move this to after reporting
      case mod.modErrors of
        Nil => pure MkUnit
        _ => addError $ E nameFC "Error in import \{ns}"

  log 1 "parse Decls"
  top <- getTop
  (decls, ops) <- parseDecls fn top.ops toks Lin

  -- TODO only include this module's ops
  -- aside from reworking parsing, we could filter
  -- other options are removing updates from parsing (so we must use incremental parsing)
  -- or removing pratt from parsing (so it happens in elaboration)
  modifyTop { currentMod $= { modOps := ops } }

  log 1 "process Decls"
  traverse (tryProcessDecl src modns) (collectDecl decls)
  top <- getTop

  -- This has addErrors as a side-effect
  logMetas $ reverse $ listValues top.currentMod.modMetaCtx.metas
  logConstraints

  top <- getTop
  -- print errors (for batch processing case)
  for_ (reverse top.currentMod.modErrors) $ showError src >=> putStrLn
  -- let (Nil) = top.currentMod.modErrors
  --   | _ => fatalError "Stop"

  -- update modules with result, leave the rest of context in case this is top file
  top <- getTop
  let modules = updateMap modns top.currentMod top.modules
  modifyTop {modules := modules}

  pure top.currentMod
  where
    reportError : Error → M ModContext
    reportError err = do
      addError err
      top <- getTop
      modifyTop {modules $= updateMap modns top.currentMod }
      pure top.currentMod

    tryProcessDecl : String → String → Decl → M Unit
    tryProcessDecl src ns decl = do
      (Left err) <- tryError $ processDecl ns decl | _ => pure MkUnit
      addError err

-- invalidate Module and anyone who depends on it
invalidateModule : String -> M Unit
invalidateModule modname = do
  top <- getTop
  let modules = join $ map getDeps $ toList top.modules
  let revMap = map swap modules
  let deps = foldl accumulate emptyMap revMap
  modifyTop { modules $= go deps (modname :: Nil) }
  where
    accumulate : SortedMap String (List String) → String × String → SortedMap String (List String)
    accumulate deps (k,v) = let prev = fromMaybe Nil $ lookupMap' k deps
      in updateMap k (v :: prev) deps

    getDeps : String × ModContext → List (String × String)
    getDeps (nm, mod) = map (_,_ nm) mod.modDeps

    go : SortedMap String (List String) → List String → SortedMap String ModContext → SortedMap String ModContext
    go deps Nil mods = mods
    go deps (name :: names) mods =
      -- Have we hit this name already?
      let (Just _) = lookupMap name mods | _ => go deps names mods in
      let ds = fromMaybe Nil $ lookupMap' name deps in
      go deps (ds ++ names) (deleteMap name mods)