module CompileMonad

import Quox.Syntax as Q
import Quox.Definition as Q
import Quox.Untyped.Syntax as U
import Quox.Parser
import Quox.Untyped.Erase
import Quox.Untyped.Scheme
import Quox.Pretty
import Quox.Log
import Options
import Output
import Error

import System.File
import Data.IORef
import Data.Maybe
import Control.Eff

%default total

%hide Doc.(>>=)
%hide Core.(>>=)

%hide FromParser.Error
%hide Erase.Error
%hide Lexer.Error
%hide Parser.Error



public export
record State where
  constructor MkState
  seen : IORef SeenSet
  defs : IORef Q.Definitions
  ns   : IORef Mods
  suf  : IORef NameSuf
%name CompileMonad.State state

export %inline
newState : HasIO io => io State
newState = pure $ MkState {
  seen = !(newIORef empty),
  defs = !(newIORef empty),
  ns   = !(newIORef [<]),
  suf  = !(newIORef 0)
}


public export
data CompileTag = OPTS | STATE

public export
Compile : List (Type -> Type)
Compile =
  [Except Error,
   ReaderL STATE State, ReaderL OPTS Options, Log,
   LoadFile, IO]


export %inline
handleLog : IORef LevelStack -> OpenFile -> LogL x a -> IOErr Error a
handleLog ref f l = case f of
  OConsole ch => handleLogIO (const $ pure ()) ref (consoleHandle ch) l
  OFile _ h   => handleLogIO (const $ pure ()) ref h l
  ONone       => do
    lvls   <- readIORef ref
    lenRef <- newIORef (length lvls)
    res    <- handleLogDiscardIO lenRef l
    writeIORef ref $ fixupDiscardedLog !(readIORef lenRef) lvls
    pure res

private %inline
withLogFile : Options ->
              (IORef LevelStack -> OpenFile -> IO (Either Error a)) ->
              IO (Either Error a)
withLogFile opts act = do
  lvlStack <- newIORef $ singleton opts.logLevels
  withOutFile CErr opts.logFile fromError $ act lvlStack
where
  fromError : String -> FileError -> IO (Either Error a)
  fromError file err = pure $ Left $ WriteError file err

export covering %inline
runCompile : Options -> State -> Eff Compile a -> IO (Either Error a)
runCompile opts state act = do
  withLogFile opts $ \lvls, logFile =>
    fromIOErr $ runEff act $ with Union.(::)
      [handleExcept (\e => ioLeft e),
       handleReaderConst state,
       handleReaderConst opts,
       handleLog lvls logFile,
       handleLoadFileIOE loadError ParseError state.seen opts.include,
       liftIO]

private %inline
rethrowFileC : String -> Either FileError a -> Eff Compile a
rethrowFileC f = rethrow . mapFst (WriteError f)


export %inline
outputStr : OpenFile -> Lazy String -> Eff Compile ()
outputStr ONone           _   = pure ()
outputStr (OConsole COut) str = putStr str
outputStr (OConsole CErr) str = fPutStr stderr str >>= rethrowFileC "<stderr>"
outputStr (OFile f h)     str = fPutStr h str >>= rethrowFileC f

export %inline
outputDocs : OpenFile ->
             ({opts : LayoutOpts} -> Eff Pretty (List (Doc opts))) ->
             Eff Compile ()
outputDocs file docs = do
  opts <- askAt OPTS
  for_ (runPretty opts (toOutFile file) docs) $ \x =>
    outputStr file $ render (Opts opts.width) x

export %inline
outputDoc : OpenFile ->
            ({opts : LayoutOpts} -> Eff Pretty (Doc opts)) -> Eff Compile ()
outputDoc file doc = outputDocs file $ singleton <$> doc


public export
data StopTag = STOP

public export
CompileStop : List (Type -> Type)
CompileStop = FailL STOP :: Compile

export %inline
withEarlyStop : Eff CompileStop () -> Eff Compile ()
withEarlyStop = ignore . runFailAt STOP

export %inline
stopHere : Has (FailL STOP) fs => Eff fs ()
stopHere = failAt STOP


export %inline
liftFromParser : Eff FromParserIO a -> Eff Compile a
liftFromParser act =
  runEff act $ with Union.(::)
    [handleExcept $ \err => throw $ FromParserError err,
     handleStateIORef !(asksAt STATE defs),
     handleStateIORef !(asksAt STATE ns),
     handleStateIORef !(asksAt STATE suf),
     \g => send g,
     \g => send g]

export %inline
liftErase : Q.Definitions -> Eff Erase a -> Eff Compile a
liftErase defs act =
  runEff act
    [handleExcept $ \err => throw $ EraseError err,
     handleStateIORef !(asksAt STATE suf),
     \g => send g]

export %inline
liftScheme : Eff Scheme a -> Eff Compile (a, List Id)
liftScheme act = do
  runEff [|MkPair act (getAt MAIN)|]
    [handleStateIORef !(newIORef empty),
     handleStateIORef !(newIORef [])]