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 Lib.Syntax
import Node

primNS : List String
primNS = ("Prim" :: Nil)

jsonTopContext : M Json
jsonTopContext = do
  top <- get
  pure $ JsonObj (("context", JsonArray (map jsonDef $ listValues top.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 PiType = (h0, h1) => ({ tag: \"PiType\", h0, h1 })"
        :: 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

-- 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 <- get
  -- TODO make top.loaded a List QName
  let modns = (snoc ns nm)
  let name = joinBy "." modns
  let (Nothing) = lookupMap modns top.modules | _ => pure ""
  modify (\ top => MkTop (updateMap modns emptyModCtx top.modules) top.imported top.ns top.defs top.metaCtx top.verbose top.errors top.ops)
  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)

  putStrLn "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

  putStrLn $ "MODNS " ++ show modns
  top <- get
  (decls, ops) <- parseDecls fn top.ops toks Lin

  top <- get
  freshMC <- newIORef (MC EmptyMap 0 CheckAll)
  -- set imported, mod, freshMC, ops  before processing
  modify (\ top => MkTop top.modules imported modns EmptyMap freshMC top.verbose top.errors ops)
  putStrLn "process Decls"
  traverse (tryProcessDecl ns) (collectDecl decls)

  -- update modules with result, leave the rest of context in case this is top file
  top <- get
  mc <- readIORef top.metaCtx
  let modules = updateMap modns (MkModCtx top.defs mc top.ops) top.modules
  freshMC <- newIORef (MC EmptyMap 0 CheckAll)
  modify (\ top => MkTop modules top.imported top.ns top.defs top.metaCtx top.verbose top.errors top.ops)

  (Nil) <- liftIO {M} $ readIORef top.errors
    | errors => do
        for_ errors $ \err =>
          putStrLn (showError src err)
        exitFailure "Compile failed"
  if stk == Nil then logMetas $ reverse $ listValues mc.metas else pure MkUnit
  pure src
  where
    tryProcessDecl : List String -> Decl -> M Unit
    tryProcessDecl ns decl = do
      Left err <- tryError $ processDecl ns decl | _ => pure MkUnit
      addError err


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 <- get
  (Nil) <- liftIO {M} $ readIORef top.errors
    | errors => do
        for_ errors $ \err =>
          putStrLn (showError src err)
        -- if err.file == fn
        -- then putStrLn (showError src err)
        -- else putStrLn (showError "" err)
        exitFailure "Compile failed"
  pure MkUnit


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
  putStrLn "\{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)

  top <- get
  let modules = updateMap primNS (MkModCtx top.defs (MC EmptyMap 0 CheckAll) top.ops) top.modules
  modify (\ top => MkTop modules (primNS :: Nil) Nil EmptyMap top.metaCtx top.verbose top.errors top.ops)

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

  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
  modify (\ top => MkTop top.modules top.imported top.ns top.defs top.metaCtx True top.errors top.ops)
  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
  let (arg0 :: args) = 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"