-- SPDX-FileCopyrightText: 2023 David Mosbach <david.mosbach@campus.lmu.de>
--
-- SPDX-License-Identifier: AGPL-3.0-or-later

{-# LANGUAGE OverloadedRecordDot,
             OverloadedStrings,
             NoFieldSelectors,
             DuplicateRecordFields,
             TupleSections #-}

module Transpiler where
  import DSL
  import Data.YAML (ToYAML (..), mapping, (.=))
  import Data.Text (Text, pack)
  import YamlParser (AnchorData (..), YAMLNode (Sequence))
  import Control.Monad.State (State, evalState, runState, get, put, unless, when)
  import Data.Map (Map, empty)
  import qualified Data.Map as M
  import Control.Monad.Except (ExceptT, runExceptT, throwError)
  import Data.Either (fromLeft)
  import Data.Maybe (fromJust, isNothing)
  

  data ResolvedLiteral = Pred' { pred :: Predicate }
                       | Neg' { pred :: Predicate }  deriving Show
  
  data DNFLiteral = DNFLit {
    anchor  :: AnchorData,
    literal :: ResolvedLiteral
  } deriving (Show)

  type DNF = [[DNFLiteral]]

  data ResolvedSubStage = RSubStage {
    head :: Head,
    body :: DNF
  } deriving (Show)
  
  instance ToYAML ResolvedSubStage where
    toYAML(RSubStage head body) = mapping [
        "mode"      .= if head.required then "required" else "optional" :: Text,
        "show-when" .= case head.showWhen of
                          Always      -> "always"
                          Fulfilled   -> "fulfilled"
                          Unfulfilled -> "unfulfilled" :: Text,
        "display-label" .= (undefined :: Text),
        "predicate" .= mapping [ "dnf-terms" .= toYAML body ]
      ]

  instance ToYAML DNFLiteral where
    toYAML (DNFLit anchor pred) = mapping [
        "tag" .= tag,
        "var" .= mapping [
          "tag" .= predToText p,
          predToText p .= pack p.ref
        ]
      ]
        where
          (tag, p) = case pred of
            Pred' x -> ("variable" :: Text, x)
            Neg' x -> ("negated", x)
          predToText :: Predicate -> Text
          predToText (EdgeInHistory _) = "edge-in-history"
          predToText (NodeInHistory _) = "node-in-history"
          predToText (PayloadFilled _) = "payload-filled"
          predToText (PreviousNode  _) = "previous-node"
          predToText x = error $ show x ++ " is not fully resolved"
    

  newtype ResolveError = ResolveError String

  instance Show ResolveError where
    show (ResolveError s) = s

  data StateData = StateData {
    innerVariables :: Map String (Variable, Bool), -- True means "already used" => anchor ref. False means "not used before" => new anchor
    outerVariables :: Map String (Variable, Bool),
    disjunction    :: DNF
  }

  type Resolver = ExceptT ResolveError (State StateData)
  newtype Warning = Warning String deriving Show
  data ResolvedData = RData {
    subStage :: ResolvedSubStage,
    warnings :: [Warning]
  } deriving (Show)

  resolve :: SubStage -> Either ResolveError ResolvedData
  resolve (SubStage head body) = evaluation
    where
      (evaluation, state) = runState (runExceptT (RData <$> (RSubStage head <$> eval body) <*> warnings)) initState
      warnings = checkUnusedVariables
      initState = StateData empty (M.map (, False) body.variables) []

  checkUnusedVariables :: Resolver [Warning]
  checkUnusedVariables = do
    state <- get
    let unusedInner = M.foldl f [] state.innerVariables
    let unusedOuter = M.foldl f [] state.outerVariables
    return $ unusedInner ++ unusedOuter
      where
        f :: [Warning] -> (Variable, Bool) -> [Warning]
        f acc (_, True) = acc
        f acc (var, False) = Warning ("Unused variable: " ++ id) : acc
          where id = case var of
                        Single id' _ -> id'
                        Block id' _ -> id'


  
  eval :: Body -> Resolver DNF
  eval (Body variables []) = get >>= \s -> return s.disjunction
  eval (Body variables (c:dnf)) = do
    conjunction <- evalConjunction c []
    state <- get
    put $ state {innerVariables = empty, disjunction = conjunction : state.disjunction}
    eval $ Body variables dnf
      where
        evalConjunction :: Conjunction -> [DNFLiteral] -> Resolver [DNFLiteral]
        evalConjunction [] acc = return acc
        evalConjunction (l:ls) acc = do
          lit <- evalLiteral l
          case lit of
            Left literal -> evalConjunction ls (literal : acc)
            Right block -> evalConjunction ls (block ++ acc) -- Merge content of block conjunction variables
        evalLiteral :: Literal -> Resolver (Either DNFLiteral [DNFLiteral])
        evalLiteral n@(Neg _) = Left <$> evalNegation n
        evalLiteral p@(Pred _) = Left <$> evalPredicate p
        evalLiteral v@(Var _) = evalVariable False v
        evalNegation :: Literal -> Resolver DNFLiteral -- Resolves redundant negations, e.g. `not not x` and also `let x = not y; let z = not x`
        evalNegation (Neg n) = do
          let (lit, count) = countNot 1 n
          lit' <- case lit of {
            Pred _ -> evalPredicate lit;
            Var _ -> evalVariable True lit >>= \l -> return $ fromLeft (error "Preventing negated blocks failed") l;
            Neg _ -> throwError . ResolveError $ "Could not resolve negation of: " ++ show n;
          }
          if even count then return lit' else do
            let sign = case lit'.literal of {
              Neg' _ -> Pred';
              Pred' _ -> Neg';
            }
            return lit' { literal = sign lit'.literal.pred }
        evalNegation x = throwError . ResolveError $ "Wrongfully labelt as negation:  " ++ show x
        countNot :: Word -> Literal -> (Literal, Word)
        countNot x (Neg n) = countNot (x+1) n
        countNot x lit = (lit, x)
        evalPredicate :: Literal -> Resolver DNFLiteral
        evalPredicate (Pred (EdgesInHistory _)) = undefined -- Problem: how to handle negations without de morgan? forbid like negating block vars?
        evalPredicate (Pred (NodesInHistory _)) = undefined
        evalPredicate (Pred (PayloadsFilled _)) = undefined
        evalPredicate (Pred (PreviousNodes _)) = undefined
        evalPredicate (Pred p) = return $ DNFLit { anchor = NoAnchor, literal = Pred' p }
        evalPredicate x = throwError . ResolveError $ "Wrongfully labelt as predicate: " ++ show x
        evalVariable :: Bool -> Literal -> Resolver (Either DNFLiteral [DNFLiteral])
        evalVariable negated (Var v) = do
          state <- get
          case M.lookup v state.innerVariables of
            Just (var, alreadyUsed) -> processVarRef var alreadyUsed True negated
            Nothing -> case M.lookup v state.outerVariables of
              Just (var, alreadyUsed) -> processVarRef var alreadyUsed False negated
              Nothing -> throwError . ResolveError $ "Reference of unassigned variable: " ++ v
        processVarRef :: Variable -> Bool -> Bool -> Bool -> Resolver (Either DNFLiteral [DNFLiteral])
        processVarRef var alreadyUsed isInner negated = do
          let updateVars = M.adjust (\(x,_) -> (x,True)) var.id
          state <- get
          unless alreadyUsed . put $ if isInner
                                     then state { innerVariables = updateVars state.innerVariables }
                                     else state { outerVariables = updateVars state.outerVariables }
          let anchor = if alreadyUsed then AnchorAlias (pack var.id) else AnchorDef (pack var.id)
          case var of
            Single _ (Pred p) -> return $ Left DNFLit { anchor = anchor, literal = Pred' p }
            Single _ v'@(Var _) -> evalVariable negated v'
            Single _ n@(Neg _) -> Left <$> (evalNegation n >>= \x -> return $ if x.anchor == NoAnchor then x {anchor = anchor} else x)
            Block id conj -> preventBlockNegation negated id >> Right <$> evalConjunction conj []
        preventBlockNegation :: Bool -> String -> Resolver ()
        preventBlockNegation True s = throwError . ResolveError $ "Negating conjunction blocks is not permitted: " ++ s
        preventBlockNegation False _ = return ()