{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{- | This module contains graphviz utilities. -}
module Network.Legion.Discovery.Graphviz (
  toDotGraph,
) where

import Data.GraphViz.Attributes.Complete (Attribute(Label, Style),
  Label(StrLabel), StyleName(Rounded, Dashed), StyleItem(SItem))
import Data.GraphViz.Types.Canonical (DotGraph(DotGraph), DotStatements(
  DotStmts), strictGraph, directedGraph, graphID, graphStatements,
  attrStmts, subGraphs, nodeStmts, edgeStmts, DotSubGraph(DotSG),
  DotNode(DotNode), DotEdge(DotEdge), isCluster, subGraphID,
  subGraphStmts, GlobalAttributes(GraphAttrs), GraphID(Str))
import Data.Map (Map)
import Data.Monoid ((<>))
import Data.Set (Set)
import Data.Version (showVersion, Version)
import Distribution.Version (withinRange)
import Network.Legion.Discovery.LegionApp (Service(Service), name,
  instances, InstanceInfo(InstanceInfo), version, unEntityName,
  Client(Client), requests, cName, cVersion, EntityName,
  RequestInfo(RequestInfo), riRange, ServiceAddr, unServiceAddr)
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.Text.Lazy as TL


toDotGraph :: [Service] -> DotGraph TL.Text
toDotGraph services = DotGraph {
      strictGraph = True,
      directedGraph = True,
      graphID = Nothing,
      graphStatements = DotStmts {
          attrStmts = [GraphAttrs [Style [SItem Rounded []]]],
          subGraphs,
          nodeStmts,
          edgeStmts
        }
    }
  where
    subGraphs :: [DotSubGraph TL.Text]
    subGraphs = [
        serviceCluster service
        | service <- services
      ]

    nodeStmts = [
        DotNode clientName [Style cStyle]
        | Service {requests} <- services
        , (client, _) <- Map.toList requests
        , (clientName, cStyle) <- clientNames client
      ]

    edgeStmts :: [DotEdge TL.Text]
    edgeStmts = [
        DotEdge clientName (instanceToNodeId name inst iaddy) [Style cStyle]
        | Service {name, instances, requests} <- services
        , (client, RequestInfo {riRange}) <- Map.toList requests
        , (clientName, cStyle) <- clientNames client
        , (iaddy, inst) <- Map.toList instances
        , withinRange (version inst) riRange
      ]

    {- | Organize the services by name and version. -}
    byNameAndVersion :: Map (EntityName, Version) (Set ServiceAddr)
    byNameAndVersion = Map.fromList [
        ((name, v), addys)
        | Service {name, instances} <- services
        , (v, addys) <- Map.toList (byVersion instances)
      ]

    clientNames :: Client -> [(TL.Text, [StyleItem])]
    clientNames client@Client {cName, cVersion} = case cVersion of
      Nothing -> [clientToNodeId client Nothing]
      Just v -> [
          clientToNodeId client addy
          | let
            addys = case Map.lookup (cName, v) byNameAndVersion of
              Just as | not (Set.null as) -> Just <$> Set.toList as
              _ -> [Nothing]
          , addy <- addys
        ]

{- | Show a 'EntityName' as 'TL.Text'. -}
showName :: EntityName -> TL.Text
showName = TL.fromStrict . unEntityName


{- | Show a version as 'TL.Text'. -}
showV :: Version -> TL.Text
showV = TL.pack . showVersion


{- | Convert a 'Client' to a 'DotNode'. -}
clientToNodeId :: Client -> Maybe ServiceAddr -> (TL.Text, [StyleItem])
clientToNodeId Client {cName, cVersion} addy =
  (
    showName cName
      <> maybe "" (("/" <>) . showV) cVersion
      <> maybe "" ((":" <>) . showAddy) addy,
    case addy of
      Nothing -> [SItem Dashed []]
      _ -> []
  )


{- | Convert a service instance to a 'DotNode'. -}
instanceToNodeId :: EntityName -> InstanceInfo -> ServiceAddr -> TL.Text
instanceToNodeId name info addy =
  showName name <> "/" <> showV (version info) <> ":" <> showAddy addy


{- | Convert a service into a cluster subgraph. -}
serviceCluster :: Service -> DotSubGraph TL.Text
serviceCluster Service {name, instances} = DotSG {
    isCluster = True,
    subGraphID = Just (Str (showName name)),
    subGraphStmts = DotStmts {
        attrStmts = [GraphAttrs [Label (StrLabel (showName name))]],
        subGraphs = [
            versionCluster name v addys
            | (v, addys) <- Map.toList (byVersion instances)
          ],
        nodeStmts = [],
        edgeStmts = []
      }
  }

{- | Create a version cluster subgraph. -}
versionCluster
  :: EntityName
  -> Version
  -> Set ServiceAddr
  -> DotSubGraph TL.Text
versionCluster n v addys = DotSG {
    isCluster = True,
    subGraphID = Just (Str (showName n <> "/" <> showV v)),
    subGraphStmts = DotStmts {
        attrStmts = [
            GraphAttrs [Label (StrLabel (showName n <> "/" <> showV v))]
          ],
        subGraphs = [],
        nodeStmts = [
            DotNode (showName n <> "/" <> showV v <> ":" <> showAddy a) []
            | a <- Set.toList addys
          ],
        edgeStmts = []
      }
  }


{- | Organize some instances by version. -}
byVersion
  :: Map ServiceAddr InstanceInfo
  -> Map Version (Set ServiceAddr)
byVersion instances =
    foldr push Map.empty (Map.toList instances)
  where
    push (addy, InstanceInfo {version}) acc =
      case version `Map.lookup` acc of
        Nothing -> Map.insert version (Set.singleton addy) acc
        Just addys -> Map.insert version (Set.insert addy addys) acc


{- | Show a service address as a 'TL.Text'. -}
showAddy :: ServiceAddr -> TL.Text
showAddy = TL.fromStrict . unServiceAddr