module Data.GraphViz.Types.State
       ( Path
       , recursiveCall
         
       , GraphState
       , ClusterLookup
       , getGraphInfo
       , addSubGraph
       , addGraphGlobals
         
       , NodeState
       , NodeLookup
       , getNodeLookup
       , toDotNodes
       , addNodeGlobals
       , addNode
       , addEdgeNodes
         
       , EdgeState
       , getDotEdges
       , addEdgeGlobals
       , addEdge
       ) where
import Data.GraphViz.Attributes.Complete   (Attributes, usedByClusters,
                                            usedByGraphs)
import Data.GraphViz.Attributes.Same
import Data.GraphViz.Types.Internal.Common
import           Control.Arrow             ((&&&), (***))
import           Control.Monad             (when)
import           Control.Monad.Trans.State
import           Data.DList                (DList)
import qualified Data.DList                as DList
import           Data.Function             (on)
import           Data.Map                  (Map)
import qualified Data.Map                  as Map
import           Data.Sequence             (Seq, ViewL (..), (|>))
import qualified Data.Sequence             as Seq
import qualified Data.Set                  as Set
type GVState s a = State (StateValue s) a
data StateValue a = SV { globalAttrs :: SAttrs
                       , useGlobals  :: Bool
                       , globalPath  :: Path
                       , value       :: a
                       }
                  deriving (Eq, Ord, Show, Read)
type Path = Seq (Maybe GraphID)
modifyGlobal   :: (SAttrs -> SAttrs) -> GVState s ()
modifyGlobal f = modify f'
  where
    f' sv@(SV{globalAttrs = gas}) = sv{globalAttrs = f gas}
modifyValue   :: (s -> s) -> GVState s ()
modifyValue f = modify f'
  where
    f' sv@(SV{value = s}) = sv{value = f s}
addGlobals    :: Attributes -> GVState s ()
addGlobals as = do addG <- gets useGlobals
                   when addG $ modifyGlobal (`unionWith` as)
getGlobals :: GVState s SAttrs
getGlobals = gets globalAttrs
getPath :: GVState s Path
getPath = gets globalPath
modifyPath   :: (Path -> Path) -> GVState s ()
modifyPath f = modify f'
  where
    f' sv@(SV{globalPath = p}) = sv{globalPath = f p}
recursiveCall      :: Maybe (Maybe GraphID) -> GVState s () -> GVState s ()
recursiveCall mc s = do gas <- getGlobals
                        p   <- getPath
                        maybe (return ()) (modifyPath . flip (|>)) mc
                        s
                        modifyGlobal (const gas)
                        modifyPath (const p)
unionWith        :: SAttrs -> Attributes -> SAttrs
unionWith sas as = toSAttr as `Set.union` sas
type GraphState a = GVState ClusterLookup' a
type ClusterLookup = Map (Maybe GraphID) ([Path], GlobalAttributes)
type ClusterLookup' = Map (Maybe GraphID) ClusterInfo
type ClusterInfo = (DList Path, SAttrs)
getGraphInfo :: GraphState a -> (GlobalAttributes, ClusterLookup)
getGraphInfo = ((graphGlobal . globalAttrs) &&& (convert . value))
               . (`execState` initState)
  where
    convert = Map.map ((uniq . DList.toList) *** toGlobal)
    toGlobal = GraphAttrs . filter usedByClusters . unSame
    graphGlobal = GraphAttrs . filter usedByGraphs . unSame
    initState = SV Set.empty True Seq.empty Map.empty
    uniq = Set.toList . Set.fromList
mergeCInfos          :: ClusterInfo -> ClusterInfo -> ClusterInfo
mergeCInfos (p1,as1) = DList.append p1 *** Set.union as1
addCluster                 :: Maybe (Maybe GraphID) -> Path -> SAttrs
                              -> GraphState ()
addCluster Nothing    _ _  = return ()
addCluster (Just gid) p as = modifyValue $ Map.insertWith mergeCInfos gid ci
  where
    ci = (DList.singleton p, as)
addSubGraph           :: Maybe (Maybe GraphID) -> GraphState a -> GraphState ()
addSubGraph mid cntns = do pth <- getPath 
                           recursiveCall mid $ do cntns
                                                  
                                                  
                                                  gas <- getGlobals
                                                  addCluster mid pth gas
addGraphGlobals                 :: GlobalAttributes -> GraphState ()
addGraphGlobals (GraphAttrs as) = addGlobals as
addGraphGlobals _               = return ()
type NodeLookup n = Map n (Path, Attributes)
type NodeLookup' n = Map n NodeInfo
data NodeInfo = NI { atts     :: SAttrs
                   , gAtts    :: SAttrs 
                   , location :: Path
                   }
              deriving (Eq, Ord, Show, Read)
type NodeState n a = GVState (NodeLookup' n) a
toDotNodes :: (Ord n) => NodeLookup n -> [DotNode n]
toDotNodes = map (\(n,(_,as)) -> DotNode n as) . Map.assocs
getNodeLookup       :: (Ord n) => Bool -> NodeState n a -> NodeLookup n
getNodeLookup addGs = Map.map combine . value . (`execState` initState)
  where
    initState = SV Set.empty addGs Seq.empty Map.empty
    combine ni = (location ni, unSame $ atts ni `Set.union` gAtts ni)
mergeNInfos :: NodeInfo -> NodeInfo -> NodeInfo
mergeNInfos (NI a1 ga1 p1) (NI a2 ga2 p2) = NI (a1 `Set.union` a2)
                                                
                                               (ga2 `Set.union` ga1)
                                                
                                               (mergePs p2 p1)
mergePs :: Path -> Path -> Path
mergePs p1 p2 = mrg' p1 p2
  where
    mrg' = mrg `on` Seq.viewl
    mrg EmptyL      _           = p2
    mrg _           EmptyL      = p1
    mrg (c1 :< p1') (c2 :< p2')
      | c1 == c2                = mrg' p1' p2'
      | otherwise               = p1
addNodeGlobals                :: GlobalAttributes -> NodeState n ()
addNodeGlobals (NodeAttrs as) = addGlobals as
addNodeGlobals _              = return ()
mergeNode            :: (Ord n) => n -> Attributes -> SAttrs -> Path
                        -> NodeState n ()
mergeNode n as gas p = modifyValue $ Map.insertWith mergeNInfos n ni
  where
    ni = NI (toSAttr as) gas p
addNode                :: (Ord n) => DotNode n -> NodeState n ()
addNode (DotNode n as) = do gas <- getGlobals
                            p <- getPath
                            
                            mergeNode n as gas p
addEdgeNodes :: (Ord n) => DotEdge n -> NodeState n ()
addEdgeNodes (DotEdge f t _) = do gas <- getGlobals
                                  p <- getPath
                                  addEN f gas p
                                  addEN t gas p
  where
    addEN n = mergeNode n []
type EdgeState n a = GVState (DList (DotEdge n)) a
getDotEdges       :: Bool -> EdgeState n a -> [DotEdge n]
getDotEdges addGs = DList.toList . value . (`execState` initState)
  where
    initState = SV Set.empty addGs Seq.empty DList.empty
addEdgeGlobals                :: GlobalAttributes -> EdgeState n ()
addEdgeGlobals (EdgeAttrs as) = addGlobals as
addEdgeGlobals _              = return ()
addEdge :: DotEdge n -> EdgeState n ()
addEdge de@DotEdge{edgeAttributes = as}
  = do gas <- getGlobals
       let de' = de { edgeAttributes = unSame $ unionWith gas as }
       modifyValue $ (`DList.snoc` de')