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 %inline
maxLogLevel : Nat
maxLogLevel = 100

public export %inline
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

-- Q: why are you using `So` instead of `LT` and `Elem`
-- A: ① proof search gives up before finding a proof of e.g. ``99 `LT` 100``
--      (i.e. `LTESucc⁹⁹ LTEZero`)
--    ② the proofs aren't looked at in any way, i just wanted to make sure the
--      list of categories was consistent everywhere


||| a verbosity level from 0–100. higher is noisier. each log entry has a
||| verbosity level above which it will be printed, chosen, uh, based on vibes.
public export
LogLevel : Type
LogLevel = Subset Nat IsLogLevel

||| a logging category, like "check" (type checking), "whnf", or whatever.
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


||| verbosity levels for each category, if they differ from the default
public export
LevelMap : Type
LevelMap = List (LogCategory, LogLevel)

-- Q: why `List` instead of `SortedMap`
-- A: oof ouch my constant factors (maybe this one was more obvious)


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

public export
LevelStack : Type
LevelStack = List LogLevels

public export %inline
defaultLevel : LogLevel
defaultLevel = Element 0 Oh

export %inline
defaultLogLevels : LogLevels
defaultLogLevels = MkLogLevels defaultLevel []

export %inline
initStack : LevelStack
initStack = []

export %inline
getLevel1 : LogCategory -> LogLevels -> LogLevel
getLevel1 cat (MkLogLevels def lvls) = fromMaybe def $ lookup cat lvls

export %inline
getLevel : LogCategory -> LevelStack -> LogLevel
getLevel cat (lvls :: _) = getLevel1 cat lvls
getLevel cat []          = defaultLevel

export %inline
getCurLevels : LevelStack -> LogLevels
getCurLevels (lvls :: _) = lvls
getCurLevels []          = defaultLogLevels


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


private %inline
replace : Eq a => a -> b -> List (a, b) -> List (a, b)
replace k v kvs = (k, v) :: filter (\y => fst y /= k) kvs

private %inline
mergeLeft : Eq a => List (a, b) -> List (a, b) -> List (a, b)
mergeLeft l r = foldl (\lst, (k, v) => replace k v lst) r l


public export
data PushArg =
    SetDefault LogLevel
  | SetCat     LogCategory LogLevel
  | SetAll     LogLevel
%name PushArg push

export %inline
applyPush : LogLevels -> PushArg -> LogLevels
applyPush lvls (SetDefault def)     = {defLevel := def} lvls
applyPush lvls (SetCat     cat lvl) = {levels   $= replace cat lvl} lvls
applyPush lvls (SetAll     lvl)     = MkLogLevels lvl []

export %inline
fromPush : PushArg -> LogLevels
fromPush = applyPush defaultLogLevels


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
  ||| print some log messages
  SayMany : (cat : LogCategory) -> (loc : Loc) ->
            (msgs : List LogMsg) -> LogL lbl ()
  ||| set some verbosity levels
  Push : (push : List PushArg) -> LogL lbl ()
  ||| restore the previous verbosity levels.
  ||| returns False if the stack was already empty
  Pop : LogL lbl Bool
  ||| returns the current verbosity levels
  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 : List PushArg -> Eff fs ()
  pushAt lvls = send $ Push {lbl} lvls

  public export %inline
  push1At : PushArg -> Eff fs ()
  push1At lvl = pushAt [lvl]

  public export %inline
  popAt : Eff fs Bool
  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 : List PushArg -> Eff fs ()
  push = pushAt ()

  public export %inline
  push1 : PushArg -> Eff fs ()
  push1 = push1At ()

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

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


||| handles a `Log` effect with an existing `State` and `Writer`
export %inline
handleLogSW : (0 s : ts) -> (0 w : tw) ->
              Has (StateL s LevelStack) fs => Has (WriterL w LogDoc) fs =>
              LogL tag a -> Eff fs a
handleLogSW s w = \case
  Push push => modifyAt s $ \lst =>
    foldl applyPush (fromMaybe defaultLogLevels (head' lst)) push :: lst
  Pop => stateAt s $ maybe (False, []) (True,) . tail'
  SayMany cat loc msgs => do
    catLvl <- getsAt s $ fst . getLevel cat
    let loc = runPretty $ prettyLoc loc
    for_ msgs $ \(lvl :> msg) => when (lvl <= catLvl) $ tellAt w $
      hcat [loc, text cat.fst, "@", pshow lvl, ":"] <++> msg
  CurLevels =>
    getsAt s getCurLevels

export %inline
handleLogSW_ : LogL tag a -> Eff [State LevelStack, Writer LogDoc] a
handleLogSW_ = handleLogSW () ()

export %inline
handleLogIO : HasIO m => MonadRec m =>
              (FileError -> m ()) -> IORef LevelStack -> File ->
              LogL tag a -> m a
handleLogIO th lvls h act =
  runEff (handleLogSW_ act) [handleStateIORef lvls, handleWriter {m} printMsg]
where printMsg : LogDoc -> m ()
      printMsg msg = fPutStr h (render _ msg) >>= either th pure

export %inline
handleLogST : HasST m => MonadRec (m s) =>
              STRef s (SnocList LogDoc) -> STRef s LevelStack ->
              LogL tag a -> m s a
handleLogST docs lvls act =
  runEff (handleLogSW_ act) [handleStateSTRef lvls, handleWriterSTRef docs]

export %inline
handleLogDiscard : (0 s : ts) -> Has (StateL s Nat) fs =>
                   LogL tag a -> Eff fs a
handleLogDiscard s = \case
  Push _     => modifyAt s S
  Pop        => stateAt s $ \k => (k > 0, pred k)
  SayMany {} => pure ()
  CurLevels  => pure defaultLogLevels

export %inline
handleLogDiscard_ : LogL tag a -> Eff [State Nat] a
handleLogDiscard_ = handleLogDiscard ()

export %inline
handleLogDiscardST : HasST m => MonadRec (m s) => STRef s Nat ->
                     LogL tag a -> m s a
handleLogDiscardST ref act =
  runEff (handleLogDiscard_ act) [handleStateSTRef ref]

export %inline
handleLogDiscardIO : HasIO m => MonadRec m => IORef Nat ->
                     LogL tag a -> m a
handleLogDiscardIO ref act =
  runEff (handleLogDiscard_ act) [handleStateIORef ref]


||| approximate the push/pop effects in a discarded log by trimming a stack or
||| repeating its most recent element
export %inline
fixupDiscardedLog : Nat -> LevelStack -> LevelStack
fixupDiscardedLog want lvls =
  let len = length lvls in
  case compare len want of
    EQ => lvls
    GT => drop (len `minus` want) lvls
    LT => let new = fromMaybe defaultLogLevels $ head' lvls in
          replicate (want `minus` len) new ++ lvls