module Main

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

import Data.IORef
import Data.SortedSet
import Text.Show.PrettyVal
import Text.Show.Pretty
import System
import System.File
import Control.Eff

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


private
die : HasIO io => (opts : LayoutOpts) -> Doc opts -> io a
die opts err = do
  ignore $ fPutStr stderr $ render opts err
  exitFailure

private
runPretty : Options -> Eff Pretty a -> a
runPretty opts act =
  let doColor = opts.color && opts.outFile == Console
      hl      = if doColor then highlightSGR else noHighlight
  in
  runPrettyWith Outer opts.flavor hl 2 act

private
putErrLn : HasIO io => String -> io ()
putErrLn = ignore . fPutStrLn stderr

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

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

private
data Error =
  ParseError String Parser.Error
| FromParserError FromParser.Error
| EraseError Erase.Error
| WriteError FilePath FileError
| NoMain
| MultipleMains (List Id)
%hide FromParser.Error
%hide Erase.Error
%hide Lexer.Error
%hide Parser.Error


private
loadError : Loc -> FilePath -> FileError -> Error
loadError loc file err = FromParserError $ LoadError loc file err

private
prettyError : {opts : LayoutOpts} -> Error -> Eff Pretty (Doc opts)
prettyError (ParseError file e) = prettyParseError file e
prettyError (FromParserError e) = FromParser.prettyError True e
prettyError (EraseError      e) = Erase.prettyError True e
prettyError NoMain              = pure "no #[main] function given"
prettyError (MultipleMains xs)  =
  pure $ sep ["multiple #[main] functions given:",
              separateLoose "," !(traverse prettyId xs)]
prettyError (WriteError file e) = pure $
  hangSingle 2 (text "couldn't write file \{file}:") (pshow e)

private
data CompileTag = OPTS | STATE

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

private
runCompile : Options -> State -> Eff Compile a -> IO (Either Error a)
runCompile opts state act =
  fromIOErr $ runEff act $ with Union.(::)
    [handleExcept (\e => ioLeft e),
     handleReaderConst state,
     handleReaderConst opts,
     handleLoadFileIOE loadError ParseError state.seen opts.include,
     liftIO]


private
data StopTag = STOP

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

private
withEarlyStop : Has (FailL STOP) fs => Eff fs () -> Eff (fs - FailL STOP) ()
withEarlyStop = ignore . runFailAt STOP

private
stopHere : Has (FailL STOP) fs => Eff fs ()
stopHere = failAt STOP


private
FlexDoc : Type
FlexDoc = {opts : LayoutOpts} -> Doc opts


private
outputStr : Lazy String -> Eff Compile ()
outputStr str =
  case !(asksAt OPTS outFile) of
    None    => pure ()
    Console => putStr str
    File f  => do
      res <- withFile f WriteTruncate pure $ \h => fPutStr h str
      rethrow $ mapFst (WriteError f) res

private
outputDocs : (opts : Options) ->
             ({opts : LayoutOpts} -> List (Doc opts)) -> Eff Compile ()
outputDocs opts doc =
  outputStr $ concat $ map (render (Opts opts.width)) doc

private
outputDocStopIf : Phase ->
                  ({opts : LayoutOpts} -> Eff Pretty (List (Doc opts))) ->
                  Eff CompileStop ()
outputDocStopIf p docs = do
  opts <- askAt OPTS
  when (opts.until == Just p) $ Prelude.do
    lift $ outputDocs !(askAt OPTS) (runPretty opts docs)
    stopHere

private
liftFromParser : Eff FromParserIO a -> Eff CompileStop 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]

private
liftErase : Q.Definitions -> Eff Erase a -> Eff CompileStop a
liftErase defs act =
  runEff act
    [\case Err e => throw $ EraseError e,
     \case Ask   => pure defs,
     handleStateIORef !(asksAt STATE suf)]

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


private
oneMain : Has (Except Error) fs => List Id -> Eff fs Id
oneMain []    = throw NoMain
oneMain [x]   = pure x
oneMain mains = throw $ MultipleMains mains

private
processFile : String -> Eff Compile ()
processFile file = withEarlyStop $ do
  Just ast <- loadFile noLoc file
    | Nothing => pure ()
  -- putErrLn "checking \{file}"
  when (!(asksAt OPTS until) == Just Parse) $ do
    lift $ outputStr $ show ast
    stopHere
  defList <- liftFromParser $ concat <$> traverse fromPTopLevel ast
  outputDocStopIf Check $
    traverse (uncurry Q.prettyDef) defList
  let defs = SortedMap.fromList defList
  erased <- liftErase defs $
    traverse (\(x, d) => (x,) <$> eraseDef x d) defList
  outputDocStopIf Erase $
    traverse (uncurry U.prettyDef) erased
  (scheme, mains) <- liftScheme $ map catMaybes $
    traverse (uncurry defToScheme) erased
  outputDocStopIf Scheme $
    intersperse empty <$> traverse prettySexp scheme
  opts <- askAt OPTS
  main <- oneMain mains
  lift $ outputDocs opts $ intersperse empty $ runPretty opts $ do
    res    <- traverse prettySexp scheme
    runner <- makeRunMain main
    pure $ text Scheme.prelude :: res ++ [runner]

export
main : IO ()
main = do
  (_, opts, files) <- options
  case !(runCompile opts !newState $ traverse_ processFile files) of
    Right () => pure ()
    Left  e  => die (Opts opts.width) $
                runPretty ({outFile := Console} opts) $
                prettyError e


-----------------------------------
{-

private
text : PrettyOpts -> List String
text _ =
  ["",
   #" ___ ___ _____ __ __"#,
   #"/ _ `/ // / _ \\ \ /"#,
   #"\_, /\_,_/\___/_\_\"#,
   #" /_/"#,
   ""]

private
qtuwu : PrettyOpts -> List String
qtuwu opts =
  if opts.unicode then
    [#"      ___,-´⎠  "#,
     #"(·`──´ ◡    -´⎠"#,
     #" \/\/──´⎞/`──´ "#,
     #"  ⎜⎟───,-₎   ⎞ "#,
     #"  ⎝⎠  (‾‾)   ⎟ "#,
     #"     (‾‾‾)   ⎟ "#]
  else
    [#"      ___,-´/  "#,
     #"(.`--´ u    -´/"#,
     #" \/\/--´|/`--´ "#,
     #"  ||---,-,   \ "#,
     #"  `´  (--)   | "#,
     #"     (---)   | "#]

private
join1 : PrettyOpts -> String -> String -> String
join1 opts l r =
  if opts.color then
    "  " <+> show (colored Green l) <+> "  " <+> show (colored Magenta r)
  else
    "  " <+> l <+> "  " <+> r

export
banner : PrettyOpts -> String
banner opts = unlines $ zipWith (join1 opts) (qtuwu opts) (text opts)
-}