From 5dbf264b7b6c9f51b434130cf0db4c6d95a9a900 Mon Sep 17 00:00:00 2001
From: Marcin Sulecki <marcin.sulecki@sulmar.pl>
Date: Wed, 20 Dec 2023 15:07:38 +0100
Subject: [PATCH] Refactor GetInitialTransition

---
 src/Stateless/Graph/GraphStyleBase.cs   |  6 ++++++
 src/Stateless/Graph/StateGraph.cs       |  6 +-----
 src/Stateless/Graph/UmlDotGraphStyle.cs | 19 ++++++++++++++++++-
 3 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/src/Stateless/Graph/GraphStyleBase.cs b/src/Stateless/Graph/GraphStyleBase.cs
index 32788e86..ffa4e644 100644
--- a/src/Stateless/Graph/GraphStyleBase.cs
+++ b/src/Stateless/Graph/GraphStyleBase.cs
@@ -18,6 +18,12 @@ public abstract class GraphStyleBase
         /// <returns>Prefix text</returns>
         public abstract string GetPrefix();
 
+        /// <summary>
+        /// Get initial transition if present
+        /// </summary>
+        /// <returns></returns>
+        public abstract string GetInitialTransition(Reflection.StateInfo initialState);
+
         /// <summary>
         /// Returns the formatted text for a single state.
         /// For example, for DOT files this would be the description of a single node:
diff --git a/src/Stateless/Graph/StateGraph.cs b/src/Stateless/Graph/StateGraph.cs
index 74a2d9ca..3477bfd4 100644
--- a/src/Stateless/Graph/StateGraph.cs
+++ b/src/Stateless/Graph/StateGraph.cs
@@ -87,11 +87,7 @@ public string ToGraph(GraphStyleBase style)
                 dirgraphText += System.Environment.NewLine + transit;
 
             // Add initial transition if present
-            var initialStateName = initialState.UnderlyingState.ToString();
-            dirgraphText += System.Environment.NewLine + $" init [label=\"\", shape=point];";
-            dirgraphText += System.Environment.NewLine + $" init -> \"{initialStateName}\"[style = \"solid\"]";
-
-            dirgraphText += System.Environment.NewLine + "}";
+            dirgraphText += style.GetInitialTransition(initialState);
 
             return dirgraphText;
         }
diff --git a/src/Stateless/Graph/UmlDotGraphStyle.cs b/src/Stateless/Graph/UmlDotGraphStyle.cs
index b21f0599..9a053747 100644
--- a/src/Stateless/Graph/UmlDotGraphStyle.cs
+++ b/src/Stateless/Graph/UmlDotGraphStyle.cs
@@ -1,4 +1,5 @@
-using System;
+using Stateless.Reflection;
+using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Text;
@@ -122,5 +123,21 @@ internal string FormatOneLine(string fromNodeName, string toNodeName, string lab
         {
             return $"\"{fromNodeName}\" -> \"{toNodeName}\" [style=\"solid\", label=\"{label}\"];";
         }
+
+        /// <summary>
+        /// 
+        /// </summary>
+        /// <param name="initialState"></param>
+        /// <returns></returns>
+        public override string GetInitialTransition(StateInfo initialState)
+        {
+            var initialStateName = initialState.UnderlyingState.ToString();
+            string dirgraphText = System.Environment.NewLine + $" init [label=\"\", shape=point];";
+            dirgraphText += System.Environment.NewLine + $" init -> \"{initialStateName}\"[style = \"solid\"]";
+
+            dirgraphText += System.Environment.NewLine + "}";
+
+            return dirgraphText;
+        }
     }
 }