module Main

import Prelude
import Data.List1
import Data.String
import Data.IORef
import Data.SortedMap
import Lib.Common
import Lib.Compile
import Lib.Parser
import Lib.Elab
import Lib.Util
import Lib.Parser.Impl
import Lib.Prettier
import Lib.ProcessDecl
import Lib.Token
import Lib.Tokenizer
import Lib.TopContext
import Lib.Types
import Lib.Syntax
import Node
import Serialize

-- For editors, dump some information about the context (fc, name, type)
jsonTopContext : M Json
jsonTopContext = do
  top <- getTop
  let defs = join $ map (\mod => listValues mod.modDefs) $ listValues top.modules
  pure $ JsonObj (("context", JsonArray (map jsonDef $ defs)) :: Nil)
  where
    jsonDef : TopEntry -> Json
    -- There is no FC here...
    jsonDef (MkEntry fc (QN ns name) type def _) = JsonObj
      (  ("fc", toJson fc)
      :: ("name", toJson name)
      :: ("type", toJson (render 80 $ pprint Nil type) )
      :: Nil)

dumpContext : TopContext -> M Unit
dumpContext top = do
    putStrLn "Context:"
    go $ listValues top.defs
    putStrLn "---"
  where
    go : List TopEntry -> M Unit
    go Nil = pure MkUnit
    go (x :: xs) = putStrLn "    \{show x}" >> go xs

writeSource : String -> M Unit
writeSource fn = do
  docs <- compile
  let src = unlines $
        ( "\"use strict\";"
        :: "const bouncer = (f,ini) => { let obj = ini; while (obj.tag) obj = f(obj); return obj.h0 };"
        :: Nil)
        ++ map (render 90 ∘ noAlt) docs
  (Right _) <- liftIO {M} $ writeFile fn src
    | Left err => exitFailure (show err)
  -- (Right _) <- chmodRaw fn 493 | Left err => exitFailure (show err)
  pure MkUnit


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
      putStrLn $ showError "" err
      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

moduleHash : String → List (List String) → M String
moduleHash src imports = do
  srcHash <- liftIO $ checksum src
  top <- getTop
  let mods = mapMaybe (\x => lookupMap' x top.modules) imports
  let modHashes = map (\x => x.csum) mods
  liftIO $ checksum $ fastConcat $ srcHash :: modHashes

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

-- New style loader, one def at a time
processModule : FC -> String -> List String -> QName -> M String
processModule importFC base stk qn@(QN ns nm) = do
  top <- getTop
  -- TODO make top.loaded a List QName
  let modns = (snoc ns nm)
  let name = joinBy "." modns
  let (Nothing) = lookupMap modns top.modules | _ => pure ""
  -- dummy entry for processing
  modifyTop [modules := updateMap modns (emptyModCtx "") top.modules]
  let fn = (joinBy "/" (base :: ns)) ++ "/" ++ nm ++ ".newt"
  (Right src) <- liftIO {M} $ readFile fn
    | Left err => exitFailure "ERROR at \{show importFC}: error reading \{fn}: \{show err}"
  let (Right toks) = tokenise fn src
    | Left err => exitFailure (showError src err)

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

  log 1 $ \ _ => "scan imports for module \{modName}"
  let ns = split modName "."
  let (path, modName') = unsnoc $ split1 modName "."
  -- let bparts = split base "/"
  let (True) = qn == QN path modName'
    | _ => exitFailure "ERROR at \{show 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) => exitFailure (showError src err)

  imported <- for imports $ \case
    MkImport fc name' => do
      let (a,b) = unsnoc $ split1 name' "."
      let qname = QN a b
      -- we could use `fc` if it had a filename in it
      when (elem name' stk) $ \ _ => error emptyFC "import loop \{show name} -> \{show name'}"

      processModule fc base (name :: stk) qname
      pure $ split name' "."
  let imported = snoc imported primNS
  srcSum <- liftIO $ checksum src
  csum <- moduleHash srcSum imported

  putStrLn "module \{modName}"
  top <- getTop
  -- TODO we need a flag on this so `make newt3.js` properly tests self-compile
  (Nothing) <- loadModule qn csum
    | Just mod => do
      let modules = updateMap modns mod top.modules

      -- FIXME - we don't want stray operators in a module.
      -- inject module ops into top
      let ops = foldMap const top.ops $ toList mod.ctxOps
      modifyTop [modules := modules; ops := ops ]
      pure src  -- why am I returning this?

  log 1 $ \ _ => "MODNS " ++ show modns
  top <- getTop
  (decls, ops) <- parseDecls fn top.ops toks Lin

  top <- getTop
  let freshMC = MC emptyMap Nil 0 CheckAll
  -- set imported, mod, freshMC, ops  before processing
  modifyTop (\ top => MkTop top.modules imported emptyMap modns emptyMap freshMC top.verbose top.errors ops)
  for imported $ \ ns => do
    let (Just mod) = lookupMap' ns top.modules | _ => error emptyFC "namespace \{show ns} missing"
    importHints (listValues mod.modDefs)

  log 1 $ \ _ => "process Decls"
  traverse (tryProcessDecl src ns) (collectDecl decls)

  -- update modules with result, leave the rest of context in case this is top file
  top <- getTop

  let mod = MkModCtx csum top.defs top.metaCtx top.ops
  errors <- liftIO {M} $ readIORef top.errors
  if stk /= Nil && length' errors == 0
    then dumpModule qn src mod
    else pure MkUnit

  let modules = updateMap modns mod top.modules
  modifyTop [modules := modules]

  (Nil) <- liftIO {M} $ readIORef top.errors
    | errors => do
        -- we're now showing errors when they occur, so they're next to debug messages
        -- traverse (putStrLn ∘ showError src) errors
        exitFailure "Compile failed"
  logMetas $ reverse $ listValues top.metaCtx.metas
  pure src
  where
    tryProcessDecl : String -> List String → Decl -> M Unit
    tryProcessDecl src ns decl = do
      (Left err) <- tryError $ processDecl ns decl | _ => pure MkUnit
      putStrLn $ showError src err
      addError err

-- unwind the module part of the path name
baseDir : SnocList String -> SnocList String -> Either String String
baseDir dirs Lin = Right $ joinBy "/" (dirs <>> Nil)
baseDir (dirs :< d) (ns :< n) = if d == n
  then baseDir dirs ns
  else Left "module path doesn't match directory"
baseDir Lin _ = Left "module path doesn't match directory"

showErrors : String -> String -> M Unit
showErrors fn src = do
  top <- getTop
  -- TODO {M} needed to sort out scrutinee
  (Nil) <- liftIO {M} $ readIORef top.errors
    | errors => do
        traverse (putStrLn ∘ showError src) errors
        exitFailure "Compile failed"
  pure MkUnit

-- processFile called on the top level file
-- it sets up everything and then recurses into processModule
processFile : String -> M Unit
processFile fn = do
  putStrLn "*** Process \{fn}"
  let parts = split1 fn "/"
  let (dirs,file) = unsnoc parts
  let dir = if dirs == Nil then "." else joinBy "/" dirs
  let (name, ext) = splitFileName file
  log 1 $ \ _ => "\{show dir} \{show name} \{show ext}"

  (Right src) <- liftIO {M} $ readFile fn
    | Left err => error (MkFC fn (0,0)) "error reading \{fn}: \{show err}"
  let (Right toks) = tokenise fn src
    | Left err => throwError err
  let (Right ((nameFC, modName), _, _)) = partialParse fn parseModHeader emptyMap toks
    | Left (err,toks) => throwError err

  (base,qn) <- getBaseDir fn modName

  -- declare internal primitives

  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
  let modules = updateMap primNS (MkModCtx "" top.defs (MC emptyMap Nil 0 CheckAll) top.ops) top.modules
  modifyTop (\ top => MkTop modules (primNS :: Nil) emptyMap Nil emptyMap top.metaCtx top.verbose top.errors top.ops)

  src <- processModule emptyFC base Nil qn
  top <- getTop

  showErrors fn src
  pure MkUnit

cmdLine : List String -> M (Maybe String × List String)
cmdLine Nil = pure (Nothing, Nil)
cmdLine ("--top" :: args) = cmdLine args -- handled later
cmdLine ("-v" :: args) = do
  modifyTop [ verbose $= _+_ 1 ]
  cmdLine args
cmdLine ("-o" :: fn :: args) = do
  (out, files) <- cmdLine args
  pure ((out <|> Just fn), files)

cmdLine (fn :: args) = do
  let (True) = isSuffixOf ".newt" fn
    | _ => error emptyFC "Bad argument \{show fn}"
  (out, files) <- cmdLine args
  pure $ (out, fn :: files)

main' : M Unit
main' = do
  (arg0 :: args) <- liftIO {M} $ getArgs
    | _ => error emptyFC "error reading args"
  (out, files) <- cmdLine args
  traverse processFile files

  when (elem "--top" args) $ \ _ => do
    json <- jsonTopContext
    putStrLn "TOP:\{renderJson json}"

  case out of
    Nothing => pure MkUnit
    Just name => writeSource name


main : IO Unit
main = do
  -- we'll need to reset for each file, etc.
  ctx <- emptyTop
  (Right _)  <- .runM main' ctx
    | Left err => exitFailure "ERROR at \{show $ getFC err}: \{errorMsg err}"
  putStrLn "done"