{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE RankNTypes #-}

module Language.Fiddle.Compiler where

import Control.Arrow
import Control.Monad (when)
import Control.Monad.Identity (Identity)
import Control.Monad.RWS
import Control.Monad.Trans.Maybe
import Control.Monad.Writer
import Data.Default
import Language.Fiddle.Ast
import Language.Fiddle.Types
import Options.Applicative
import System.IO (hPutStrLn, stderr)
import Text.Parsec (sourceColumn, sourceLine, sourceName)

data Level = Error | Warning | Info
  deriving (Eq, Ord, Show, Read, Enum)

data Diagnostic = Diagnostic Level String SourceSpan

emitDiagnosticError :: String -> Commented SourceSpan -> Compile a ()
emitDiagnosticError str a = tell [Diagnostic Error str (unCommented a)]

emitDiagnosticWarning :: String -> Commented SourceSpan -> Compile a ()
emitDiagnosticWarning str a = tell [Diagnostic Warning str (unCommented a)]

emitDiagnosticInfo :: String -> Commented SourceSpan -> Compile a ()
emitDiagnosticInfo str a = tell [Diagnostic Info str (unCommented a)]

newtype Compile s a = Compile (MaybeT (RWS () [Diagnostic] s) a)
  deriving newtype (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

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

pushState :: Compile s a -> Compile s a
pushState cp = do
  s <- get
  snd <$> subCompile s cp

compile :: Compile s a -> s -> ([Diagnostic], Maybe a)
compile (Compile fn) initState = do
  let (a, _, w) = runRWS (runMaybeT fn) () initState
   in if hasError w then (w, Nothing) else (w, a)
  where
    hasError = any (\(Diagnostic e _ _) -> e == Error)

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

pureCompilationPhase ::
  (CompilationStage stageFrom) =>
  ( FiddleUnit
      stageFrom
      Identity
      (StageAnnotation stageFrom) ->
    Compile
      ()
      ( FiddleUnit
          stageTo
          Identity
          (StageAnnotation stageTo)
      )
  ) ->
  CompilationPhase stageFrom stageTo
pureCompilationPhase fn =
  CompilationPhase (pure ()) (\_ _ -> return ([], Just ())) (\_ _ -> fn)

-- data IOActionExtraData = IOActionExtraData
--   { parseFile :: FilePath -> IO (TreeType FiddleUnit Parsed),
--     stage3Compile :: TreeType FiddleUnit Parsed ->
--       TreeType FiddleUnit Checked
--   }

-- | 'CompilationPhase' represents a phase in the compilation process.
-- It consists of an IO action that performs necessary side effects or state
-- preparations before the next stage, and a function that transforms the
-- 'FiddleUnit' from the current stage to the next.
data CompilationPhase stageFrom stageTo where
  CompilationPhase ::
    forall privateFlags privateState stageFrom stageTo.
    (CompilationStage stageFrom) =>
    { optionsParser :: Parser privateFlags,
      -- | 'ioAction' is an IO operation that runs after the ast is parsed. It
      -- takes the parsed 'FiddleUnit' and performs some side effect
      -- returning a private state that is passed to 'nextStage'. This is the
      -- only time a side effect may be performed.
      ioAction ::
        privateFlags ->
        FiddleUnit Parsed Identity (StageAnnotation Parsed) ->
        IO ([Diagnostic], Maybe privateState),
      -- | 'nextStage' is the function that transforms a 'FiddleUnit' from
      -- the current stage ('stageFrom') to the next stage ('stageTo'). It
      -- uses the private state obtained from 'ioAction' and outputs a
      -- potentially updated 'FiddleUnit' in the compilation pipeline.
      nextStage ::
        privateFlags ->
        privateState ->
        FiddleUnit
          stageFrom
          Identity
          (StageAnnotation stageFrom) ->
        Compile
          ()
          ( FiddleUnit
              stageTo
              Identity
              (StageAnnotation stageTo)
          )
    } ->
    CompilationPhase stageFrom stageTo

-- | 'thenPhase' composes two 'CompilationPhase' stages into a single pipeline
-- phase. It combines their IO actions and applies each stage in sequence.
thenPhase ::
  CompilationPhase stage1 stage2 ->
  CompilationPhase stage2 stage3 ->
  CompilationPhase stage1 stage3
thenPhase
  (CompilationPhase optParse1 ioAction1 compile1)
  (CompilationPhase optParse2 ioAction2 compile2) =
    CompilationPhase
      ((,) <$> optParse1 <*> optParse2)
      ( \(f1, f2) unit -> do
          (diags1, mst1) <- ioAction1 f1 unit
          case mst1 of
            Nothing -> return (diags1, Nothing)
            Just st1 -> do
              (diags2, mst2) <- ioAction2 f2 unit
              return $ case mst2 of
                Nothing -> (diags1 ++ diags2, Nothing)
                Just st2 -> (diags1 ++ diags2, Just (st1, st2))
      )
      ( \(f1, f2) (s1, s2) firstStage -> do
          secondStage <- compile1 f1 s1 firstStage
          compile2 f2 s2 secondStage
      )

-- | Infix operator for 'thenPhase' to chain compilation phases.
(>>>) :: CompilationPhase stage1 stage2 -> CompilationPhase stage2 stage3 -> CompilationPhase stage1 stage3
(>>>) = thenPhase

-- | 'execCompilationPipeline' executes a full compilation pipeline starting
-- from the 'Parsed' phase. It performs the IO action of the first phase and
-- then invokes the compilation function for the remaining stages. It returns
-- a tuple containing diagnostics and an optional final 'FiddleUnit'.
execCompilationPipelineWithCmdline ::
  CompilationPhase Parsed s' ->
  Parser
    ( FiddleUnit Parsed Identity (StageAnnotation Parsed) ->
      IO
        ( [Diagnostic],
          Maybe
            ( FiddleUnit s' Identity (StageAnnotation s')
            )
        )
    )
execCompilationPipelineWithCmdline
  (CompilationPhase flagsParser ioAct rest) = do
    fmap
      ( \opts ast -> do
          (diags, ms) <- ioAct opts ast
          case ms of
            Just s -> return $ first (diags ++) $ compile_ $ rest opts s ast
            Nothing -> return (diags, Nothing)
      )
      flagsParser

squeezeDiagnostics :: (Alter t) => t (Either [Diagnostic]) a -> Compile () (t Identity a)
squeezeDiagnostics ast = do
  _ <-
    alter
      ( \case
          (Left l) -> tell l >> return (Left l)
          r -> return r
      )
      return
      ast

  case squeeze ast of
    (Left _) -> compilationFailure
    (Right a) -> return a

resolveOrFail :: Either [Diagnostic] a -> Compile s a
resolveOrFail (Right a) = return a
resolveOrFail (Left l) = tell l >> compilationFailure