module Language.Fiddle.Compiler where

import Control.Monad.State
import Control.Monad.Writer
import Data.Default
import Language.Fiddle.Ast
import Language.Fiddle.Types
import Text.Parsec (SourcePos, sourceColumn, sourceLine, sourceName)

data Level = Error | Warning | Info

data Diagnostic = Diagnostic Level String SourceSpan

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

instance Functor (Compile s) where
  fmap fn (Compile cfn) = Compile $ \s ->
    let (s', d', ma) = cfn s in (s, d', fmap fn ma)

instance Applicative (Compile s) where
  (<*>) mfn ma = do
    fn <- mfn
    fn <$> ma

  pure = return

instance Monad (Compile s) where
  return a = Compile (,[],Just a)

  -- m a -> (a -> m b) -> m b
  (>>=) (Compile cfn) fn = Compile $ \s ->
    let (s', diags, ma) = cfn s
     in case ma of
          Nothing -> (s', diags, Nothing)
          Just a ->
            let (Compile cfn') = fn a
                (s'', diags', mb) = cfn' s'
             in (s'', diags ++ diags', mb)

instance MonadWriter [Diagnostic] (Compile s) where
  tell diag = Compile (,diag,Just ())

  listen (Compile fn) = Compile $ \s ->
    let (s', diags, ma) = fn s in (s', diags, (,diags) <$> ma)

  -- Not really "correctly" implemented, but I suspect this function will not be
  -- used very much.
  pass (Compile fn) = Compile $ \s ->
    let (s', diags, mafn) = fn s
     in case mafn of
          Just (a, fn) -> (s', fn diags, Just a)
          Nothing -> (s', diags, Nothing)

instance MonadState s (Compile s) where
  get = Compile $ \s -> (s, [], Just s)

  put s = Compile $ const (s, [], Just ())

hoistMaybe :: Maybe a -> Compile s a
hoistMaybe ma = Compile (,[],ma)

-- 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 a
subCompile s' (Compile fn) = Compile $ \s ->
  let (_, diags, ma) = fn s' in (s, diags, ma)

-- Saves the state, runs the routine, then restores the state.
pushState :: Compile s a -> Compile s a
pushState cp = do
  s <- get
  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 =
  let (_, d, ma) = fn initState in (d, ma)

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

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 =
  putStrLn (diagnosticToString coloredFormat d)