content: {-# LANGUAGE ExistentialQuantification #-}
  {-# LANGUAGE FlexibleContexts          #-}
  {-# LANGUAGE OverloadedStrings         #-}
  -- Copyright 2024 United States Government as represented by the Administrator
  -- of the National Aeronautics and Space Administration. All Rights Reserved.
  --
  -- Disclaimers
  --
  -- Licensed under the Apache License, Version 2.0 (the "License"); you may
  -- not use this file except in compliance with the License. You may obtain a
  -- copy of the License at
  --
  --      https://www.apache.org/licenses/LICENSE-2.0
  --
  -- Unless required by applicable law or agreed to in writing, software
  -- distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
  -- WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  -- License for the specific language governing permissions and limitations
  -- under the License.
  --
  -- | Transform a state diagram into a Copilot specification.
  module Command.Diagram
      ( diagram
      , DiagramOptions(..)
      , DiagramFormat(..)
      , DiagramMode(..)
      , DiagramPropFormat(..)
      , ErrorCode
      )
    where
  
  -- External imports
  import           Control.Exception                 as E
  import           Control.Monad                     (when)
  import           Data.Aeson                        (object, (.=))
  import           Data.ByteString.Lazy              (toStrict)
  import qualified Data.ByteString.Lazy              as B
  import           Data.Either                       (isLeft)
  import           Data.Foldable                     (for_)
  import           Data.Functor.Identity             (Identity)
  import           Data.GraphViz                     (graphEdges)
  import qualified Data.GraphViz                     as G
  import qualified Data.GraphViz.Attributes.Complete as Attributes
  import           Data.GraphViz.Commands.IO         (toUTF8)
  import qualified Data.GraphViz.Parsing             as G
  import           Data.GraphViz.PreProcessing       (preProcess)
  import qualified Data.GraphViz.Types.Generalised   as Gs
  import           Data.List                         (intercalate, nub, sort)
  import qualified Data.Set                          as Set
  import           Data.Text                         (Text)
  import qualified Data.Text                         as T
  import qualified Data.Text.Encoding                as T
  import           Data.Text.Lazy                    (pack)
  import qualified Data.Text.Lazy                    as LT
  import           Data.Void                         (Void)
  import           System.FilePath                   ((</>))
  import           Text.Megaparsec                   (ErrorFancy (ErrorFail),
                                                      ParsecT, empty,
                                                      errorBundlePretty,
                                                      fancyFailure, many,
                                                      manyTill, noneOf, parse)
  import           Text.Megaparsec.Char              (alphaNumChar, char,
                                                      digitChar, newline, space1,
                                                      string)
  import qualified Text.Megaparsec.Char.Lexer        as L
  
  
  -- External imports: auxiliary
  import Data.ByteString.Extra  as B ( safeReadFile )
  import System.Directory.Extra ( copyTemplate )
  
  -- External imports: parsing expressions.
  import qualified Language.Lustre.ParLustre as Lustre (myLexer, pBoolSpec)
  import qualified Language.SMV.ParSMV       as SMV (myLexer, pBoolSpec)
  
  -- Internal imports: auxiliary
  import Command.Result  (Result (..))
  import Data.Location   (Location (..))
  import Paths_ogma_core (getDataDir)
  
  -- Internal imports: language ASTs, transformers
  import           Language.SMV.Substitution     (substituteBoolExpr)
  import qualified Language.Trans.Lustre2Copilot as Lustre (boolSpec2Copilot,
                                                            boolSpecNames)
  import           Language.Trans.SMV2Copilot    as SMV (boolSpec2Copilot,
                                                         boolSpecNames)
  
  -- | Generate a new Copilot monitor that implements a state machine described
  -- in a diagram given as an input file.
  --
  -- PRE: The file given is readable, contains a valid file with recognizable
  -- format, the formulas in the file do not use any identifiers that exist in
  -- Copilot, or any of @stateMachine@, @externalState@, @noneOf@,
  -- @checkValidTransitions@, @main@, @spec@, @stateMachine1@, @clock@, @ftp@,
  -- @notPreviousNot@. All identifiers used are valid C99 identifiers. The
  -- template, if provided, exists and uses the variables needed by the diagram
  -- application generator. The target directory is writable and there's enough
  -- disk space to copy the files over.
  diagram :: FilePath       -- ^ Path to a file containing a diagram
          -> DiagramOptions -- ^ Customization options
          -> IO (Result ErrorCode)
  diagram fp options = do
    E.handle (return . diagramTemplateError fp) $ do
      -- Sub-parser for edge expressions.
      let functions = exprPair (diagramPropFormat options)
  
      -- Convert the diagram into elements in a Copilot spec.
      copilotSpecElems <- diagram' fp options functions
  
      -- Convert the elements into a success or error result.
      let (mOutput, result) = diagramResult fp copilotSpecElems
  
      -- If the result is success, expand the template.
      for_ mOutput $ \(streamDefs, handlerInputs) -> do
        let subst = object
                      [ "streamDefs"    .= pack streamDefs
                      , "specName"      .= pack (diagramFilename options)
                      , "input"         .= pack (diagramInputVar options)
                      , "state"         .= pack (diagramStateVar options)
                      , "handlerInputs" .= pack handlerInputs
                      ]
  
        templateDir <- case diagramTemplateDir options of
                         Just x  -> return x
                         Nothing -> do
                           dataDir <- getDataDir
                           return $ dataDir </> "templates" </> "diagram"
  
        let targetDir = diagramTargetDir options
  
        copyTemplate templateDir subst targetDir
  
      return result
  
  -- | Generate a new Copilot monitor that implements a state machine described
  -- in a diagram given as an input file, using a subexpression handler.
  --
  -- PRE: The file given is readable, contains a valid file with recognizable
  -- format, the formulas in the file do not use any identifiers that exist in
  -- Copilot, or any of @stateMachine@, @externalState@, @noneOf@,
  -- @checkValidTransitions@, @main@, @spec@, @stateMachine1@, @clock@, @ftp@,
  -- @notPreviousNot@. All identifiers used are valid C99 identifiers. The
  -- template, if provided, exists and uses the variables needed by the diagram
  -- application generator. The target directory is writable and there's enough
  -- disk space to copy the files over.
  diagram' :: FilePath
           -> DiagramOptions
           -> ExprPair
           -> IO (Either String (String, String))
  diagram' fp options exprP = do
    contentEither <- B.safeReadFile fp
    return $ do
      -- All of the following operations use Either to return error messages. The
      -- use of the monadic bind to pass arguments from one function to the next
      -- will cause the program to stop at the earliest error.
      diagFileContent <- contentEither
  
      -- Abtract representation of a state machine diagram.
      diagramR <- parseDiagram (diagramFormat options) diagFileContent exprP
  
      return $ diagramToCopilot diagramR (diagramMode options)
  
  -- | Options used to customize the conversion of diagrams to Copilot code.
  data DiagramOptions = DiagramOptions
    { diagramTargetDir   :: FilePath
    , diagramTemplateDir :: Maybe FilePath
    , diagramFormat      :: DiagramFormat
    , diagramPropFormat  :: DiagramPropFormat
    , diagramFilename    :: String
    , diagramMode        :: DiagramMode
    , diagramStateVar    :: String
    , diagramInputVar    :: String
    }
  
  -- | Modes of operation.
  data DiagramMode = CheckState   -- ^ Check if given state matches expectation
                   | ComputeState -- ^ Compute expected state
                   | CheckMoves   -- ^ Check if transitioning to a state would be
                                  --   possible.
    deriving (Eq, Show)
  
  -- | Diagram formats supported.
  data DiagramFormat = Mermaid
                     | Dot
    deriving (Eq, Show)
  
  -- | Property formats supported.
  data DiagramPropFormat = Lustre
                         | Inputs
                         | Literal
                         | SMV
    deriving (Eq, Show)
  
  -- * Error codes
  
  -- | Encoding of reasons why the command can fail.
  --
  -- The error code used is 1 for user error.
  type ErrorCode = Int
  
  -- | Error: the input file cannot be read due to it being unreadable or the
  -- format being incorrect.
  ecDiagramError :: ErrorCode
  ecDiagramError = 1
  
  -- | Error: diagram component generation failed during the copy/write
  -- process.
  ecDiagramTemplateError :: ErrorCode
  ecDiagramTemplateError = 2
  
  -- * Result
  
  -- | Process the result of the transformation function.
  diagramResult :: FilePath
                -> Either String a
                -> (Maybe a, Result ErrorCode)
  diagramResult fp result = case result of
    Left msg -> (Nothing, Error ecDiagramError msg (LocationFile fp))
    Right t  -> (Just t,  Success)
  
  -- | Report an error when trying to open or copy the template.
  diagramTemplateError :: FilePath
                       -> E.SomeException
                       -> Result ErrorCode
  diagramTemplateError fp exception =
      Error ecDiagramTemplateError msg (LocationFile fp)
    where
      msg =
        "Diagram monitor generation failed during copy/write operation. Check"
        ++ " that there's free space in the disk and that you have the necessary"
        ++ " permissions to write in the destination directory. "
        ++ show exception
  
  -- * Handler for boolean expressions in edges or transitions between states.
  
  -- | Handler for boolean expressions that knows how to parse them, replace
  -- variables in them, and convert them to Copilot.
  data ExprPair = forall a . ExprPair
    { _exprParse   :: String -> Either String a
    , _exprReplace :: [(String, String)] -> a -> a
    , _exprPrint   :: a -> String
    , _exprIdents  :: a -> [String]
    }
  
  -- | Return a handler depending on the format used for edge or transition
  -- properties.
  exprPair :: DiagramPropFormat -> ExprPair
  exprPair Lustre  = ExprPair (Lustre.pBoolSpec . Lustre.myLexer)
                              (\_ -> id)
                              Lustre.boolSpec2Copilot
                              Lustre.boolSpecNames
  exprPair Inputs  = ExprPair ((Right . read) :: String -> Either String Int)
                              (\_ -> id)
                              (\x -> "input == " ++ show x)
                              (const [])
  exprPair Literal = ExprPair Right
                              (\_ -> id)
                              id
                              (const [])
  exprPair SMV     = ExprPair (SMV.pBoolSpec . SMV.myLexer)
                              substituteBoolExpr
                              SMV.boolSpec2Copilot
                              SMV.boolSpecNames
  
  -- | Parse and print a value using an auxiliary Expression Pair.
  --
  -- Fails if the value has no valid parse.
  exprPairShow :: ExprPair -> String -> String
  exprPairShow (ExprPair parseProp _replace printProp _ids) =
    printProp . fromRight' . parseProp
  
  -- * Diagrams
  
  -- | Internal representation for diagrams.
  newtype Diagram = Diagram
      { diagramTransitions :: [(Int, String, Int)]
      }
    deriving (Show, Eq)
  
  -- * Diagram parsers
  
  -- | Generic function to parse a diagram.
  parseDiagram :: DiagramFormat          -- ^ Format of the input file
               -> B.ByteString           -- ^ Contents of the diagram
               -> ExprPair               -- ^ Subparser for conditions or edge
                                         -- expressions
               -> Either String Diagram
  parseDiagram Dot     = parseDiagramDot
  parseDiagram Mermaid = parseDiagramMermaid
  
  -- ** Dot parser
  
  -- | Parse a DOT / Graphviz diagram.
  parseDiagramDot :: B.ByteString -> ExprPair -> Either String Diagram
  parseDiagramDot contents exprP = do
      let contentsUTF8 = toUTF8 contents
      dg <- fst $ G.runParser G.parse $ preProcess contentsUTF8
      return $ makeDiagram dg
    where
      makeDiagram :: Gs.DotGraph LT.Text -> Diagram
      makeDiagram g = Diagram links
        where
          links = map edgeToLink (graphEdges g)
  
          edgeToLink edge =
              ( read (LT.unpack o)
              , exprPairShow exprP (LT.unpack e)
              , read (LT.unpack d)
              )
            where
              o = G.fromNode edge
              d = G.toNode edge
              e = getLabel (G.edgeAttributes edge)
  
              -- Extract the label from a list of attributes. If no label is
              -- found, it's assumed that the condition is the literal true.
              getLabel [] = "true"
              getLabel ((Attributes.Label (Attributes.StrLabel l)) : _) = l
              getLabel (_ : as) = getLabel as
  
  -- ** Mermaid parser
  
  -- | Parse a mermaid diagram.
  parseDiagramMermaid :: B.ByteString -> ExprPair -> Either String Diagram
  parseDiagramMermaid txtDia exprP =
      case parsingResult of
        Left e  -> Left (errorBundlePretty e)
        Right x -> Right x
    where
      txt           = T.decodeUtf8 (toStrict txtDia)
      parsingResult = parse (spaces *> pDiagram exprP) "<input>" txt
  
  -- | Type for parser for memaid diagrams.
  type MermaidParser = ParsecT Void Text Identity
  
  -- | Parser for a mermaid diagram.
  --
  -- This parser depends on an auxiliary parser for the expressions associated to
  -- the edges or connections between states.
  pDiagram :: ExprPair -> MermaidParser Diagram
  pDiagram exprP = do
    _ <- string "graph" <* spaces
    _name <- T.pack <$> manyTill alphaNumChar (char ';')
    _ <- newline
  
    transitions <- many (pTransition exprP)
  
    pure $ Diagram transitions
  
  -- | Parser for an edge in a state diagram.
  --
  -- This parser depends on an auxiliary parser for the expressions associated to
  -- the edges or connections between states.
  pTransition :: ExprPair -> MermaidParser (Int, String, Int)
  pTransition ep@(ExprPair { _exprParse = parseProp }) = do
    _ <- spaces
    stateFrom <- many digitChar
    _ <- string "-->|"
    edge <- many (noneOf ("|" :: [Char]))
  
    let x = parseProp edge
    when (isLeft x) $ fancyFailure $ Set.singleton $
      ErrorFail $ "Edge property has incorrect format: " ++ show edge
  
    _ <- char '|'
    stateTo <- many digitChar
    _ <- char ';'
    _ <- newline
    return (read stateFrom, exprPairShow ep edge, read stateTo)
  
  -- | Consume spaces
  spaces :: MermaidParser ()
  spaces = L.space space1 empty empty
  
  -- * Backend
  
  -- | Convert the diagram into a set of Copilot definitions, and a list of
  -- arguments for the top-level handler.
  diagramToCopilot :: Diagram -> DiagramMode -> (String, String)
  diagramToCopilot diag mode = (machine, arguments)
    where
      machine = unlines
        [ "stateMachineProp :: Stream Bool"
        , "stateMachineProp = " ++ propExpr
        , ""
        , "stateMachine1 :: Stream Word8"
        , "stateMachine1 = stateMachineGF (initialState, finalState, noInput, "
          ++ "transitions, badState)"
        , ""
        , "-- Check"
        , "initialState :: Word8"
        , "initialState = " ++ show initialState
        , ""
        , "-- Check"
        , "finalState :: Word8"
        , "finalState = " ++ show finalState
        , ""
        , "noInput :: Stream Bool"
        , "noInput = false"
        , ""
        , "badState :: Word8"
        , "badState = " ++ show badState
        , ""
        , "transitions = " ++ showTransitions
        ]
  
      -- Elements of the spec.
      propExpr     = case mode of
                       CheckState   -> "stateMachine1 == externalState"
                       ComputeState -> "true"
                       CheckMoves   -> "true"
      initialState = minimum states
      finalState   = maximum states
      badState     = maximum states + 1
  
      -- Arguments for the handler.
      arguments = "[ " ++ intercalate ", " (map ("arg " ++) argExprs) ++ " ]"
  
      argExprs = case mode of
        CheckState   -> [ "stateMachine1", "externalState", "input" ]
        ComputeState -> [ "stateMachine1", "externalState", "input" ]
        CheckMoves   -> map stateCheckExpr states
  
      stateCheckExpr stateId =
        "(checkValidTransition transitions externalState " ++ show stateId ++ ")"
  
      -- States and transitions from the diagram.
      transitions = diagramTransitions diag
      states      = nub $ sort $ concat [ [x, y] | (x, _, y) <- transitions ]
  
      showTransitions :: String
      showTransitions = "[" ++ showTransitions' transitions
  
      showTransitions' :: [(Int, String, Int)] -> String
      showTransitions' []         = "]"
      showTransitions' (x1:x2:xs) =
        showTransition x1 ++ ", " ++ showTransitions' (x2:xs)
      showTransitions' (x2:[])    = showTransition x2 ++ "]"
  
      showTransition :: (Int, String, Int) -> String
      showTransition (a, b, c) =
        "(" ++ show a ++ ", " ++ b ++ ", " ++ show c ++ ")"
  
  -- * Auxiliary functions
  
  -- | Unsafe fromRight. Fails if the value is a 'Left'.
  fromRight' :: Either a b -> b
  fromRight' (Right v) = v
  fromRight' _         = error "fromRight' applied to Left value."
  