module Lib.ProcessModule

import Prelude
import Lib.Types
import Lib.Common
import Lib.Syntax
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)
  setDef (QN primNS "PiType") emptyFC (Erased emptyFC) (PrimFn "(h0, h1) => ({ tag: \"PiType\", h0, h1 });" (S (S Z)) 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)) = split 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 parts = split modns "."
  let fn = joinBy "/" parts ++ ".newt"
  -- TODO now we can pass in the module name...
  (fn,src) <- repo.getFile importFC fn
  let (Right toks) = tokenise fn src
    | Left err => throwError err

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

  log 1 $ \ _ => "scan imports for module \{modName}"
  let (True) = modns == modName
    | _ => throwError $ 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) => throwError 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}"
  top <- getTop
  let freshMC = MC emptyMap 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
  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

  -- print errors (for batch processing case)
  for_ top.currentMod.modErrors $ \ err => putStrLn $ showError src err

  -- 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
    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
  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 → M Unit
    go deps Nil = pure MkUnit
    go deps (name :: names) = do
      modifyTop [modules $= deleteMap name]
      let ds = fromMaybe Nil $ lookupMap' name deps
      go deps $ ds ++ names