module Lib.TopContext

import Data.IORef
import Data.SortedMap
import Data.String
import Prelude
import Lib.Common
import Lib.Types


-- TODO move the def in here (along with M) or merge this into types
-- The Monad can be its own file if we pull all of the monad update functions there.

lookup : QName -> TopContext -> Maybe TopEntry
lookup qn@(QN ns nm) top =
  if ns == top.currentMod.modName
    then lookupMap' qn top.currentMod.modDefs
    else case lookupMap' ns top.modules of
      Just mod => lookupMap' qn mod.modDefs
      Nothing => Nothing

lookupImported : String → TopContext -> List TopEntry
lookupImported raw top =
  mapMaybe (flip lookup top) $ (QN top.currentMod.modName raw) :: map (flip QN raw) top.currentMod.modDeps

-- For repl / out of scope errors
lookupAll : String → TopContext -> List TopEntry
lookupAll raw top =
  mapMaybe (flip lookup top) $ (QN top.currentMod.modName raw) :: map (flip QN raw) (map fst $ toList top.modules)

lookupRaw : String -> TopContext -> Maybe TopEntry
lookupRaw raw top =
  case lookupMap' (QN top.currentMod.modName raw) top.currentMod.modDefs of
    Just entry => Just entry
    Nothing => go top.currentMod.modDeps
  where
    go : List String → Maybe TopEntry
    go Nil = Nothing
    go (ns :: nss) = case lookupMap' ns top.modules of
      Nothing => go nss
      Just mod => case lookupMap' (QN ns raw) mod.modDefs of
        Just entry => Just entry
        Nothing => go nss

instance Show TopContext where
  show top = "\nContext:\n [\{ joinBy "\n" $ map (show ∘ snd) $ toList top.currentMod.modDefs}]"

emptyTop : TopContext
emptyTop = MkTop emptyMap emptyMap (emptyModCtx "" "") 0 emptyMap 0

setFlag : QName → FC → EFlag → M Unit
setFlag name fc flag = do
  top <- getTop
  let (Just (MkEntry fc nm ty def flags)) = lookupMap' name top.currentMod.modDefs
    | Nothing => error fc "\{show name} not declared"
  modifyTop { currentMod $= { modDefs $= (updateMap name (MkEntry fc name ty def (flag :: flags)))} }

setDef : QName -> FC -> Tm -> Def → List EFlag -> M Unit
setDef name fc ty def flags = do
  top <- getTop
  let (Nothing) = lookupMap' name top.currentMod.modDefs
    | Just (MkEntry fc' nm' ty' def' _) => error fc "\{show name} is already defined at \{show fc'}"

  modifyTop {currentMod $= { modDefs $= (updateMap name (MkEntry fc name ty def flags))} }

updateDef : QName -> FC -> Tm -> Def -> M Unit
updateDef name fc ty def = do
  top <- getTop
  let (Just (MkEntry fc' nm' ty' def' flags)) = lookupMap' name top.currentMod.modDefs
    | Nothing => error fc "\{show name} not declared"
  modifyTop { currentMod $= { modDefs := updateMap name (MkEntry fc' name ty def flags) top.currentMod.modDefs } }

typeName : Tm → Maybe QName
typeName (Pi fc nm Explicit rig t u) = Nothing
typeName (Pi fc name icit rig t u) = typeName u
typeName (App fc t u) = typeName t
typeName (Ref _ nm) = Just nm
typeName _ = Nothing

addHint : QName → M Unit
addHint qn = do
  top <- getTop
  case lookup qn top of
    Just entry => do
      let (Just tyname) = typeName entry.type
        | Nothing => error entry.fc "can't find tcon name for \{show qn}"
      let xs = fromMaybe Nil $ lookupMap' tyname top.hints
      let hints = updateMap tyname ((qn, entry.type) :: xs) top.hints
      modifyTop { hints := hints }
    Nothing => pure MkUnit

addError : Error -> M Unit
addError err = modifyTop { currentMod $= { modErrors $= (err ::) } }

addInfo : EditorInfo → M Unit
addInfo info = modifyTop { currentMod $= {modInfos $= (info ::) } }

-- temporary? used in derive for now
freshName : String → M String
freshName nm = do
  top <- getTop
  modifyTop { freshIx $= 1 + }
  pure $ "f$" ++ nm ++ show top.freshIx