module Quox.Log

import Quox.Loc
import Quox.Pretty

import Data.So
import Data.DPair
import Data.Maybe
import Data.List1
import Control.Eff
import Control.Monad.ST.Extra
import Data.IORef
import System.File
import Derive.Prelude

%default total
%language ElabReflection


public export
maxLogLevel : Nat
maxLogLevel = 100

public export
logCategories : List String
logCategories = ["whnf", "equal", "check"]

public export %inline
isLogLevel : Nat -> Bool
isLogLevel l = l <= maxLogLevel

public export
IsLogLevel : Nat -> Type
IsLogLevel l = So $ isLogLevel l

public export %inline
isLogCategory : String -> Bool
isLogCategory cat = cat `elem` logCategories

public export
IsLogCategory : String -> Type
IsLogCategory cat = So $ isLogCategory cat


public export
LogLevel : Type
LogLevel = Subset Nat IsLogLevel

public export
LogCategory : Type
LogCategory = Subset String IsLogCategory


public export %inline
toLogLevel : Nat -> Maybe LogLevel
toLogLevel l =
  case choose $ isLogLevel l of
    Left  y => Just $ Element l y
    Right _ => Nothing

public export %inline
toLogCategory : String -> Maybe LogCategory
toLogCategory c =
  case choose $ isLogCategory c of
    Left  y => Just $ Element c y
    Right _ => Nothing


public export
LevelMap : Type
LevelMap = List (LogCategory, LogLevel)
-- i tried SortedMap first, but it is too much overhead for LevelMaps

public export
record LogLevels where
  constructor MkLogLevels
  defLevel : LogLevel
  levels   : LevelMap
%name LogLevels lvls
%runElab derive "LogLevels" [Eq, Show]

public export
LevelStack : Type
LevelStack = List1 LogLevels

export %inline
defaultLogLevels : LogLevels
defaultLogLevels = MkLogLevels (Element 0 Oh) []

export %inline
initStack : LevelStack
initStack = singleton defaultLogLevels

||| right biased for the default and for overlapping elements
public export %inline
mergeLevels : LogLevels -> LogLevels -> LogLevels
mergeLevels (MkLogLevels _ map1) (MkLogLevels def map2) =
  MkLogLevels def $ map1 ++ map2

export %inline
getLevel : LogCategory -> LogLevels -> LogLevel
getLevel cat lvls = fromMaybe lvls.defLevel $ lookup cat lvls.levels


public export
LogDoc : Type
LogDoc = Doc (Opts {lineLength = 80})


public export
data PushArg = SetDefault LogLevel | SetCats LevelMap
%name PushArg push

export %inline
mergePush : PushArg -> LogLevels -> LogLevels
mergePush (SetDefault def) = {defLevel := def}
mergePush (SetCats    map) = {levels   $= (map ++)}


public export
record LogMsg where
  constructor (:>)
  level : Nat
  {auto 0 levelOk : IsLogLevel level}
  message : Lazy LogDoc
infix 0 :>
%name Log.LogMsg msg

public export
data LogL : (lbl : tag) -> Type -> Type where
  SayMany   : (cat : LogCategory) -> (loc : Loc) ->
              (msgs : List LogMsg) -> LogL lbl ()
  Push      : (push : PushArg) -> LogL lbl ()
  Pop       : LogL lbl ()
  CurLevels : LogL lbl LogLevels

public export
Log : Type -> Type
Log = LogL ()

parameters (0 lbl : tag) {auto _ : Has (LogL lbl) fs}
  public export %inline
  sayManyAt : (cat : String) -> (0 catOk : IsLogCategory cat) =>
              Loc -> List LogMsg -> Eff fs ()
  sayManyAt cat loc msgs {catOk} =
    send $ SayMany {lbl} (Element cat catOk) loc msgs

  public export %inline
  sayAt : (cat : String) -> (0 catOk : IsLogCategory cat) =>
          (lvl : Nat)    -> (0 lvlOk : IsLogLevel lvl) =>
          Loc -> Lazy LogDoc -> Eff fs ()
  sayAt cat lvl loc msg = sayManyAt cat loc [lvl :> msg]

  public export %inline
  pushAt : PushArg -> Eff fs ()
  pushAt lvls = send $ Push {lbl} lvls

  public export %inline
  popAt : Eff fs ()
  popAt = send $ Pop {lbl}

  public export %inline
  curLevelsAt : Eff fs LogLevels
  curLevelsAt = send $ CurLevels {lbl}

parameters {auto _ : Has Log fs}
  public export %inline
  sayMany : (cat : String) -> (0 catOk : IsLogCategory cat) =>
            Loc -> List LogMsg -> Eff fs ()
  sayMany = sayManyAt ()

  public export %inline
  say : (cat : String) -> (0 _ : IsLogCategory cat) =>
        (lvl : Nat)    -> (0 _ : IsLogLevel lvl) =>
        Loc -> Lazy LogDoc -> Eff fs ()
  say = sayAt ()

  public export %inline
  push : PushArg -> Eff fs ()
  push = pushAt ()

  public export %inline
  pop : Eff fs ()
  pop = popAt ()

  public export %inline
  curLevels : Eff fs LogLevels
  curLevels = curLevelsAt ()


export %inline
doPush : PushArg -> LevelStack -> LevelStack
doPush push list = mergePush push (head list) `cons` list

export %inline
doPop : List1 a -> List1 a
doPop (_ ::: x :: xs) = x ::: xs
doPop (x ::: [])      = x ::: []

export %inline
doSayMany : Applicative m =>
            LevelStack -> (LogDoc -> m ()) ->
            LogCategory -> Loc -> List LogMsg -> m ()
doSayMany (lvls ::: _) act cat loc msgs = do
  let Element catLvl _ = getLevel cat lvls
      loc              = runPretty $ prettyLoc loc
  for_ msgs $ \msg => when (msg.level <= catLvl) $
    act $ hcat [loc, text cat.fst, "@", pshow msg.level, ":"] <++>
          msg.message

export %inline
handleLogIO : HasIO m => (FileError -> m ()) ->
              IORef LevelStack -> File -> LogL tag a -> m a
handleLogIO th lvls h = \case
  Push push            => modifyIORef lvls $ doPush push
  Pop                  => modifyIORef lvls doPop
  SayMany cat loc msgs => doSayMany !(readIORef lvls) printMsg cat loc msgs
  CurLevels            => head <$> readIORef lvls
where printMsg : LogDoc -> m ()
      printMsg msg = fPutStr h (render _ msg) >>= either th pure

export %inline
handleLogST : (HasST m, Monad (m s)) =>
              STRef s (SnocList LogDoc) -> STRef s LevelStack ->
              LogL tag a -> m s a
handleLogST docs lvls = \case
  Push push            => modifySTRef' lvls $ doPush push
  Pop                  => modifySTRef' lvls doPop
  SayMany cat loc msgs => doSayMany !(readSTRef' lvls) printMsg cat loc msgs
  CurLevels            => head <$> readSTRef' lvls
where printMsg : LogDoc -> m s ()
      printMsg msg = modifySTRef' docs (:< msg)

export %inline
handleLogDiscard : Applicative m => LogL tag a -> m a
handleLogDiscard = \case
  SayMany {} => pure ()
  Push    {} => pure ()
  Pop        => pure ()
  CurLevels  => pure defaultLogLevels