module Language.Fiddle.Compiler where

import Control.Monad (when)
import Control.Monad.RWS (RWS, RWST, evalRWS, mapRWS, runRWS)
import Control.Monad.State
import Control.Monad.Trans.Maybe
import Control.Monad.Writer
import Data.Default
import Language.Fiddle.Ast
import Language.Fiddle.Types
import Text.Parsec (SourcePos, sourceColumn, sourceLine, sourceName)
import System.IO (hPutStrLn, stderr)

data Level = Error | Warning | Info

data Diagnostic = Diagnostic Level String SourceSpan

-- Compilation monad. Has diagnostics. Optionally produces a value.
-- newtype Compile s a = Compile (s -> (s, [Diagnostic], Maybe a))

newtype Compile s a = Compile (MaybeT (RWS () [Diagnostic] s) a)
  deriving (Functor, Applicative, Monad)

compilationFailure :: Compile s a
compilationFailure = Compile $ MaybeT (return Nothing)

instance MonadWriter [Diagnostic] (Compile s) where
  tell = Compile . tell
  listen (Compile fn) = Compile $ listen fn
  pass (Compile fn) = Compile $ pass fn

instance MonadState s (Compile s) where
  get = Compile get
  put s = Compile $ put s
  state fn = Compile $ state fn

-- Runs a sub-compilation routine with the given state, but discards the
-- resulting state in favor of the original state.
subCompile :: s' -> Compile s' a -> Compile s (s', a)
subCompile s' (Compile mtrws) = Compile $ do
  let (a, s, w) = runRWS (runMaybeT mtrws) () s'
  tell w
  MaybeT $ return $ fmap (s,) a

-- mapMaybeT (mapRWS (\(a, _, w) -> (a, s', w))) mtrws

-- Saves the state, runs the routine, then restores the state.
pushState :: Compile s a -> Compile s a
pushState cp = do
  s <- get
  snd <$> subCompile s cp

-- Runs a compilation routine. It produces diagnostics and maybe a result.
-- Generally if the diagnostics contain an error, the result will be Nothing,
-- but if only Warnings are generated, then Just something will be returned.
--
-- Note that there is no actual type-level mechanism restricting this function
-- from returning something even if the diagnostics contain errors, but it
-- generally wouldn't make much sense for this to be the case.
compile :: Compile s a -> s -> ([Diagnostic], Maybe a)
compile (Compile fn) initState = do
  let (a, _, w) = runRWS (runMaybeT fn) () initState in (w, a)

compile_ :: (Default s) => Compile s a -> ([Diagnostic], Maybe a)
compile_ c = compile c def

hoistMaybe :: Maybe a -> Compile s a
hoistMaybe = Compile . MaybeT . return

newtype DiagnosticFormat = DiagnosticFormat (Diagnostic -> String)

coloredFormat :: DiagnosticFormat
coloredFormat = DiagnosticFormat $ \(Diagnostic level message (SourceSpan pos1 pos2)) ->
  execWriter $ do
    case level of
      Error -> tell "\x1b[01;31mError "
      Warning -> tell "\x1b[01;33mWarn "
      Info -> tell "\x1b[01;37mInfo "

    tell "\x1b[0m"
    tell (sourceName pos1)
    tell "("
    tellPos pos1
    when (pos2 /= pos1) $ do
      tell "-"
      tellPos pos2
    tell "): "
    tell (unwords $ words message)
  where
    tellPos pos = do
      tell (show $ sourceLine pos)
      tell ":"
      tell (show $ sourceColumn pos)

diagnosticToString :: DiagnosticFormat -> Diagnostic -> String
diagnosticToString (DiagnosticFormat f) = f

printDiagnostic :: Diagnostic -> IO ()
printDiagnostic d =
  hPutStrLn stderr (diagnosticToString coloredFormat d)

fromMayberOrFail :: SourceSpan -> String -> Maybe a -> Compile s a
fromMayberOrFail sourceSpan err Nothing = do
  tell [Diagnostic Error err sourceSpan]
  compilationFailure
fromMayberOrFail _ _ (Just a) = return a