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

module Language.Fiddle.Compiler where

import Control.Arrow
import Control.Category
import Control.Monad (when)
import Control.Monad.Identity (Identity)
import Control.Monad.RWS
import Control.Monad.Trans.Maybe
import Control.Monad.Writer
import Data.Aeson (FromJSON, ToJSON (toJSON), encode)
import qualified Data.ByteString.Lazy.Char8 as BL
import Data.Default
import Data.Kind (Type)
import Data.Typeable
import GHC.Generics (Generic)
import Language.Fiddle.Ast
import Language.Fiddle.Types
import Options.Applicative
import System.IO (hPutStrLn, stderr)
import Text.Parsec (sourceColumn, sourceLine, sourceName)
import Prelude hiding (id, (.))

-- | Represents the severity level of a diagnostic message. It indicates
-- how critical the diagnostic is.
--
-- - 'Error': Indicates a serious problem that causes the compilation to fail,
--   meaning the resulting AST will not be passed to the backend.
-- - 'Warning': Represents a less severe issue, suggesting that something may
--   need attention but does not halt the compilation.
-- - 'Info': Used for informational messages that do not indicate a problem.
data Level = Error | Warning | Info
  deriving (Eq, Ord, Show, Read, Enum, ToJSON, FromJSON, Generic)

-- | A general representation of a compilation artifact. This can be any
-- typeable value that is generated or collected during the compilation
-- process. The 'Artifact' type is used to carry various outputs or
-- intermediate results produced by the compiler.
data Artifact where
  Artifact :: forall t. (Typeable t) => t -> Artifact

-- | Represents a diagnostic message that may be emitted during compilation.
-- A 'Diagnostic' contains the severity level, a descriptive message, and
-- the location in the source code where the issue was found.
data Diagnostic = Diagnostic
  { -- | The severity of the diagnostic.
    diagnosticLevel :: Level,
    -- | The message describing the issue.
    diagnosticMessage :: String,
    -- | The location in the source code.
    diagnosticLocation :: SourceSpan
  }
  deriving (Generic, ToJSON, FromJSON, Typeable)

-- | Emits a diagnostic with 'Error' severity, indicating a critical issue
-- that will cause the compilation to fail. The diagnostic includes a message
-- and the location where the error occurred.
emitDiagnosticError :: String -> Commented SourceSpan -> Compile a ()
emitDiagnosticError str a = tell [Diagnostic Error str (unCommented a)]

-- | Emits a diagnostic with 'Warning' severity, indicating a potential issue
-- that does not prevent compilation but may require attention. Includes a
-- message and location information.
emitDiagnosticWarning :: String -> Commented SourceSpan -> Compile a ()
emitDiagnosticWarning str a = tell [Diagnostic Warning str (unCommented a)]

-- | Emits a diagnostic with 'Info' severity, used to provide informational
-- messages during compilation. Includes a message and the relevant source
-- location.
emitDiagnosticInfo :: String -> Commented SourceSpan -> Compile a ()
emitDiagnosticInfo str a = tell [Diagnostic Info str (unCommented a)]

-- | Emits a compilation artifact, which is a general output or data produced
-- during the compilation process. Artifacts are collected for later use, such
-- as generating output files or debugging information.
emitArtifact :: Artifact -> Compile s ()
emitArtifact art = Compile $ tell ([], [art])

-- | The primary monad for managing the compilation process. 'Compile' handles
-- emitting diagnostics and collecting artifacts while allowing the state to be
-- modified. The 'MaybeT' layer is used to represent potential compilation
-- failure.
newtype Compile s a = Compile (MaybeT (RWS () ([Diagnostic], [Artifact]) s) a)
  deriving newtype (Functor, Applicative, Monad)

-- | Signals a compilation failure, terminating the current compilation process.
-- No further diagnostics will be emitted. This should be used sparingly and
-- only when it is not possible to continue compilation.
compilationFailure :: Compile s a
compilationFailure = Compile $ MaybeT (return Nothing)

-- | Provides a 'MonadWriter' instance for the 'Compile' monad, allowing
-- diagnostics to be collected and manipulated during the compilation process.
-- The diagnostics are treated as the output to be written.
instance MonadWriter [Diagnostic] (Compile s) where
  -- \| Add diagnostics to the compilation output.
  tell t = Compile $ tell (t, [])

  -- \| Listen for diagnostics emitted during a computation, returning the result
  -- along with the collected diagnostics.
  listen (Compile fn) = Compile $ do
    (a, (ds, _)) <- listen fn
    return (a, ds)

  -- \| Modify the collected diagnostics using the given function.
  pass (Compile mfun) = Compile $ pass (second first <$> mfun)

-- | Provides a 'MonadState' instance for the 'Compile' monad, allowing
-- manipulation of the compiler's state during the compilation process.
instance MonadState s (Compile s) where
  -- \| Get the current state.
  get = Compile get

  -- \| Set a new state.
  put s = Compile $ put s

  -- \| Apply a state transformation function.
  state fn = Compile $ state fn

-- | Run a compilation routine under a temporary state. After the routine
-- completes, the original state is restored, and the result along with the
-- modified temporary state is returned.
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

-- | Save the current state and run a provided compilation routine with a
-- duplicated state. Once the routine finishes, the original state is restored.
pushState :: Compile s a -> Compile s a
pushState cp = do
  s <- get
  snd <$> subCompile s cp

-- | Runs a compilation routine with the given initial state and returns the
-- diagnostics, artifacts, and result. If there are any 'Error' diagnostics,
-- the result will be 'Nothing'.
compile :: Compile s a -> s -> ([Diagnostic], [Artifact], Maybe a)
compile (Compile fn) initState = do
  let (a, _, (w, as)) = runRWS (runMaybeT fn) () initState
   in if hasError w then (w, as, Nothing) else (w, as, a)
  where
    hasError = any (\(Diagnostic e _ _) -> e == Error)

-- | Runs a compilation routine with a default initial state, using the 'Default'
-- type class to provide the initial state.
compile_ :: (Default s) => Compile s a -> ([Diagnostic], [Artifact], Maybe a)
compile_ c = compile c def

-- | Converts a 'Maybe' value into the 'Compile' monad. If the value is 'Nothing',
-- the compilation fails.
hoistMaybe :: Maybe a -> Compile s a
hoistMaybe = Compile . MaybeT . return

-- | Represents a format for displaying diagnostics, providing a function
-- to convert a list of diagnostics into a string.
newtype DiagnosticFormat = DiagnosticFormat
  { diagnosticsToString :: [Diagnostic] -> String
  }

-- | Diagnostic format that outputs diagnostics as a JSON string.
jsonDiagnosticFormat :: DiagnosticFormat
jsonDiagnosticFormat = DiagnosticFormat $ BL.unpack . encode . toJSON

-- | Diagnostic format that outputs diagnostics in a colored format using ANSI
-- escape codes. The colors indicate the diagnostic level (Error, Warning, Info).
coloredFormat :: DiagnosticFormat
coloredFormat = DiagnosticFormat $ \diags ->
  unlines $
    map
      ( \(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)
      )
      diags
  where
    -- \| Helper function to format a 'SourcePos' as a line:column string.
    tellPos pos = do
      tell (show $ sourceLine pos)
      tell ":"
      tell (show $ sourceColumn pos)

-- | Prints a list of diagnostics to 'stderr' using the specified
-- diagnostic format.
printDiagnostics :: DiagnosticFormat -> [Diagnostic] -> IO ()
printDiagnostics _ [] = return ()
printDiagnostics fmt d = hPutStrLn stderr (diagnosticsToString fmt d)

-- | Prints a single diagnostic to 'stderr' using the 'coloredFormat'
-- to visually indicate the diagnostic level (Error, Warning, Info).
printDiagnostic :: Diagnostic -> IO ()
printDiagnostic d =
  hPutStrLn stderr (diagnosticsToString coloredFormat [d])

-- | Converts a 'Maybe' value into the 'Compile' monad, emitting an
-- 'Error' diagnostic and causing compilation to fail if the value is 'Nothing'.
-- If the value is 'Just a', it is returned.
fromMayberOrFail :: SourceSpan -> String -> Maybe a -> Compile s a
fromMayberOrFail sourceSpan err Nothing = do
  tell [Diagnostic Error err sourceSpan]
  compilationFailure
fromMayberOrFail _ _ (Just a) = return a

-- | Creates a pure compilation phase. The provided function is used to
-- transform the 'FiddleUnit' from one stage to another within the compilation
-- pipeline, without requiring any additional state or side effects.
pureCompilationPhase ::
  ( FiddleUnit
      stageFrom
      Identity
      (Commented SourceSpan) ->
    Compile
      ()
      ( FiddleUnit
          stageTo
          Identity
          (Commented SourceSpan)
      )
  ) ->
  CompilationPhase stageFrom stageTo
pureCompilationPhase fn =
  CompilationPhase (pure ()) (\_ _ -> return ([], Just ())) (\_ _ -> fn)

-- | Emits a list of compilation artifacts, adding them to the compilation
-- output. Artifacts are user-defined data that can be generated during
-- the compilation process.
emitArtifacts :: [Artifact] -> Compile s ()
emitArtifacts arts = Compile $ tell ([], arts)

-- | '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
          (Commented SourceSpan) ->
        Compile
          ()
          ( FiddleUnit
              stageTo
              Identity
              (Commented SourceSpan)
          )
    } ->
    CompilationPhase stageFrom stageTo

instance Category CompilationPhase where
  id = pureCompilationPhase return
  (.) = flip thenPhase

-- | '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
      )

-- | '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 ::
  (StageAnnotation Parsed ~ Commented SourceSpan) =>
  CompilationPhase Parsed s' ->
  Parser
    ( FiddleUnit Parsed Identity (Commented SourceSpan) ->
      IO
        ( [Diagnostic],
          [Artifact],
          Maybe
            ( FiddleUnit s' Identity (Commented SourceSpan)
            )
        )
    )
execCompilationPipelineWithCmdline
  (CompilationPhase flagsParser ioAct rest) = do
    fmap
      ( \opts ast -> do
          (diags, ms) <- ioAct opts ast
          case ms of
            Just s ->
              return $
                (\(a, b, c) -> (diags ++ a, b, c)) $
                  compile_ $
                    rest opts s ast
            Nothing -> return (diags, [], Nothing)
      )
      flagsParser

-- | Extracts all embedded diagnostics from the abstract syntax tree (AST)
-- and emits them as compilation warnings or errors. This function inspects
-- each subelement in the AST for diagnostic messages. If any subelement
-- contains errors (represented as a 'Left' value), the compilation fails.
-- Otherwise, the AST is returned with all deferred diagnostics resolved.
squeezeDiagnostics :: (Alter t) => t (Either [Diagnostic]) a -> Compile () (t Identity a)
squeezeDiagnostics ast = do
  -- Iterate over the AST, emitting any diagnostics found.
  _ <-
    alter
      ( \case
          (Left l) -> tell l >> return (Left l) -- Emit diagnostics if present.
          r -> return r -- Otherwise, leave the subelement unchanged.
      )
      return
      ast

  -- Attempt to transform the AST into a non-deferred form.
  -- If any subelement contained errors, the transformation will fail.
  case squeeze ast of
    (Left _) -> compilationFailure -- Fail compilation if there are unresolved diagnostics.
    (Right a) -> return a -- Return the resolved AST if no errors are found.

-- | Checks the result of a computation and either returns the result or
-- emits any associated diagnostics and fails the compilation. If the result
-- is a 'Right' value, the contained value is returned. If it is a 'Left'
-- value containing diagnostics, they are emitted and the compilation is
-- terminated.
resolveOrFail :: Either [Diagnostic] a -> Compile s a
resolveOrFail (Right a) = return a -- Return the value if no diagnostics.
resolveOrFail (Left l) = tell l >> compilationFailure -- Emit diagnostics and fail if errors are present.