From 21318520c9d676386dafaff2ec780f66586d2c48 Mon Sep 17 00:00:00 2001
From: Rutherther <fandabohacek@gmail.com>
Date: Thu, 2 Feb 2023 20:16:32 +0100
Subject: [PATCH] feat(core): add possibility to add items to stateful
 repository manually

---
 Core/NosSmooth.Core/NosSmooth.Core.csproj     |  5 +-
 .../Stateful/StatefulRepository.cs            | 64 ++++++++++++++-----
 2 files changed, 51 insertions(+), 18 deletions(-)

diff --git a/Core/NosSmooth.Core/NosSmooth.Core.csproj b/Core/NosSmooth.Core/NosSmooth.Core.csproj
index b9f04da469a089c904adc05ec7967f9b4d373e25..96a65317477a945e88d6229bd1f5cafd23a994f4 100644
--- a/Core/NosSmooth.Core/NosSmooth.Core.csproj
+++ b/Core/NosSmooth.Core/NosSmooth.Core.csproj
@@ -6,9 +6,8 @@
         <TargetFrameworks>net7.0;netstandard2.1</TargetFrameworks>
         <Authors>Rutherther</Authors>
         <Description>NosSmooth Core library allowing implementing nostale client, handling packets and commands.</Description>
-        <VersionPrefix>3.3.1</VersionPrefix>
-        <PackageReleaseNotes>Add contract reached state.
-Avoid deadlocks.</PackageReleaseNotes>
+        <VersionPrefix>3.4.0</VersionPrefix>
+        <PackageReleaseNotes>Add possibility to add items to the stateful repository manually.</PackageReleaseNotes>
     </PropertyGroup>
 
     <ItemGroup>
diff --git a/Core/NosSmooth.Core/Stateful/StatefulRepository.cs b/Core/NosSmooth.Core/Stateful/StatefulRepository.cs
index f5df603f6258dcb0665a13b1777d6a67236c0058..f2480fec1e5693134f8e8fa989bb9c8112804fcf 100644
--- a/Core/NosSmooth.Core/Stateful/StatefulRepository.cs
+++ b/Core/NosSmooth.Core/Stateful/StatefulRepository.cs
@@ -5,6 +5,7 @@
 //  Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
 using Microsoft.Extensions.DependencyInjection;
@@ -15,16 +16,43 @@ namespace NosSmooth.Core.Stateful;
 /// <summary>
 /// Repository holding all the stateful entities for various NosTale clients.
 /// </summary>
-internal class StatefulRepository
+public class StatefulRepository
 {
-    private readonly Dictionary<INostaleClient, Dictionary<Type, object>> _statefulEntities;
+    private readonly ConcurrentDictionary<INostaleClient, ConcurrentDictionary<Type, object>> _statefulEntities;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="StatefulRepository"/> class.
     /// </summary>
     public StatefulRepository()
     {
-        _statefulEntities = new Dictionary<INostaleClient, Dictionary<Type, object>>();
+        _statefulEntities = new ConcurrentDictionary<INostaleClient, ConcurrentDictionary<Type, object>>();
+    }
+
+    /// <summary>
+    /// Remove items of the given client.
+    /// </summary>
+    /// <param name="client">The client to remove.</param>
+    public void Remove(INostaleClient client)
+    {
+        _statefulEntities.Remove(client, out _);
+    }
+
+    /// <summary>
+    /// Set entity of the given type to the given client.
+    /// </summary>
+    /// <remarks>
+    /// If the entity is not set manually, there will be an attempt to create an instance.
+    /// </remarks>
+    /// <param name="client">The nostale client.</param>
+    /// <param name="entity">The entity.</param>
+    /// <typeparam name="TEntity">The type of the entity.</typeparam>
+    public void SetEntity<TEntity>(INostaleClient client, TEntity entity)
+        where TEntity : notnull
+    {
+        _statefulEntities.TryAdd(client, new ConcurrentDictionary<Type, object>());
+        var values = _statefulEntities[client];
+
+        values.AddOrUpdate(typeof(TEntity), (k) => entity, (k, v) => entity);
     }
 
     /// <summary>
@@ -36,17 +64,23 @@ internal class StatefulRepository
     /// <returns>The obtained entity.</returns>
     public object GetEntity(IServiceProvider services, INostaleClient client, Type statefulEntityType)
     {
-        if (!_statefulEntities.ContainsKey(client))
-        {
-            _statefulEntities.Add(client, new Dictionary<Type, object>());
-        }
-
-        var value = _statefulEntities[client];
-        if (!value.ContainsKey(statefulEntityType))
-        {
-            value.Add(statefulEntityType, ActivatorUtilities.CreateInstance(services, statefulEntityType));
-        }
-
-        return value[statefulEntityType];
+        var dict = _statefulEntities.AddOrUpdate
+        (
+            client,
+            _ =>
+            {
+                var objectDictionary = new ConcurrentDictionary<Type, object>();
+                objectDictionary.TryAdd
+                    (statefulEntityType, () => ActivatorUtilities.CreateInstance(services, statefulEntityType));
+                return objectDictionary;
+            },
+            (_, objectDictionary) =>
+            {
+                objectDictionary.TryAdd(statefulEntityType, ActivatorUtilities.CreateInstance(services, statefulEntityType));
+                return objectDictionary;
+            }
+        );
+
+        return dict[statefulEntityType];
     }
 }
\ No newline at end of file