Connected.Framework/Connected.Data/Storage/DatabaseConnection.cs

using Connected.Entities.Storage;
using Connected.Middleware;
using Connected.ServiceModel;
using Connected.Threading;
using System.Collections.Concurrent;
using System.Collections.Immutable;
using System.Data;
using System.Data.Common;

namespace Connected.Data.Storage;

public abstract class DatabaseConnection : MiddlewareComponent, IStorageConnection
{
	private readonly AsyncLocker _lock = new();
	private IDbConnection _connection;
	private ConcurrentDictionary<IStorageCommand, IDbCommand> _commands = null;

	protected DatabaseConnection(ICancellationContext context)
	{
		Context = context;

		Commands = new();
	}

	protected ICancellationContext Context { get; }
	private IDbTransaction Transaction { get; set; }
	public StorageConnectionMode Behavior { get; set; }
	public string ConnectionString { get; private set; }

	public async Task Initialize(StorageConnectionArgs args)
	{
		ConnectionString = args.ConnectionString;
		Behavior = args.Behavior;

		await OnInitialize();
	}

	protected virtual async Task OnInitialize()
	{
		await Task.CompletedTask;
	}

	protected abstract Task<IDbConnection> OnCreateConnection();

	private async Task<IDbConnection> GetConnection()
	{
		if (_connection is null)
		{
			await _lock.LockAsync(1, async () =>
			{
				_connection ??= await OnCreateConnection();
			});
		}

		return _connection;
	}

	private ConcurrentDictionary<IStorageCommand, IDbCommand> Commands { get; }

	public async Task Commit()
	{
		if (Transaction is null || Transaction.Connection is null)
			return;

		await _lock.LockAsync(2, async () =>
		{
			if (Transaction is null || Transaction.Connection is null)
				return;

			if (Transaction is DbTransaction db)
				await db.CommitAsync(Context is null ? CancellationToken.None : Context.CancellationToken);
			else
				Transaction.Commit();

			Transaction.Dispose();
			Transaction = null;
		});

		await Task.CompletedTask;
	}

	public async Task Rollback()
	{
		if (Transaction is null || Transaction.Connection is null)
			return;

		await _lock.LockAsync(3, async () =>
		{
			if (Transaction is null || Transaction.Connection is null)
				return;

			try
			{
				if (Transaction is DbTransaction db)
					await db.RollbackAsync(Context is null ? CancellationToken.None : Context.CancellationToken);
				else
					Transaction.Rollback();
			}
			catch { }
		});

		await Task.CompletedTask;
	}

	private async Task Open()
	{
		var connection = await GetConnection();

		if (connection.State == ConnectionState.Open)
			return;

		await _lock.LockAsync(4, async () =>
			 {
				 if (connection.State != ConnectionState.Closed)
					 return;

				 if (connection is DbConnection db)
					 await db.OpenAsync(Context is null ? CancellationToken.None : Context.CancellationToken);
				 else
					 connection.Open();
			 });

		await Task.CompletedTask;
	}

	public async Task Close()
	{
		if (_connection is null)
			return;

		if (_connection is not null && _connection.State == ConnectionState.Open)
		{
			await _lock.LockAsync(5, async () =>
				 {
					 if (_connection is not null && _connection.State == ConnectionState.Open)
					 {
						 if (Transaction is not null && Transaction.Connection is not null)
						 {
							 try
							 {
								 if (Transaction is DbTransaction db)
									 await db.RollbackAsync(Context is null ? CancellationToken.None : Context.CancellationToken);
								 else
									 Transaction.Rollback();
							 }
							 catch { }

						 }

						 if (_connection is DbConnection dbc)
							 await dbc.CloseAsync();
						 else
							 _connection.Close();
					 }
				 });
		}

		await Task.CompletedTask;
	}

	public async Task<int> Execute(IStorageCommand command)
	{
		await EnsureOpen(true);

		var com = await ResolveCommand(command);

		SetupParameters(command, com);

		if (command.Operation.Parameters is not null)
		{
			foreach (var i in command.Operation.Parameters)
				SetParameterValue(com, i.Name, i.Value);
		}

		var recordsAffected = await OnExecute(command, com);

		if (command.Operation.Parameters is not null)
		{
			foreach (var i in command.Operation.Parameters)
			{
				if (i.Direction == ParameterDirection.Output)
					i.Value = GetParameterValue(com, i.Name);
			}
		}

		return recordsAffected;

	}
	protected virtual void SetParameterValue(IDbCommand command, string parameterName, object value)
	{

	}

	protected virtual object? GetParameterValue(IDbCommand command, string parameterName)
	{
		return default;
	}

	protected virtual void SetupParameters(IStorageCommand command, IDbCommand cmd)
	{
	}

	protected virtual async Task<int> OnExecute(IStorageCommand command, IDbCommand cmd)
	{
		if (cmd is DbCommand dbCommand)
			return await dbCommand.ExecuteNonQueryAsync(Context is null ? CancellationToken.None : Context.CancellationToken);
		else
			return cmd.ExecuteNonQuery();
	}

	public virtual async Task<ImmutableList<R>> Query<R>(IStorageCommand command)
	{
		await EnsureOpen(false);

		var com = await ResolveCommand(command);

		IDataReader rdr = null;

		try
		{
			SetupParameters(command, com);

			if (command.Operation.Parameters is not null)
			{
				foreach (var i in command.Operation.Parameters)
					SetParameterValue(com, i.Name, i.Value);
			}

			rdr = com is DbCommand db ? await db.ExecuteReaderAsync(Context is null ? CancellationToken.None : Context.CancellationToken) : com.ExecuteReader();
			var result = new List<R>();
			var mappings = new FieldMappings<R>(rdr);

			while (rdr.Read())
				result.Add(mappings.CreateInstance(rdr));

			return result.ToImmutableList();
		}
		finally
		{
			if (rdr != null && !rdr.IsClosed)
			{
				if (rdr is DbDataReader db)
					await db.CloseAsync();
				else
					rdr.Close();
			}
		}
	}

	public virtual async Task<R?> Select<R>(IStorageCommand command)
	{
		await EnsureOpen(false);

		var com = await ResolveCommand(command);

		IDataReader rdr = null;

		try
		{
			SetupParameters(command, com);

			if (command.Operation.Parameters is not null)
			{
				foreach (var i in command.Operation.Parameters)
					SetParameterValue(com, i.Name, i.Value);
			}

			rdr = com.ExecuteReader(CommandBehavior.SingleRow);
			var mappings = new FieldMappings<R>(rdr);

			if (rdr.Read())
				return mappings.CreateInstance(rdr);

			return default;
		}
		finally
		{
			if (rdr != null && !rdr.IsClosed)
			{
				if (rdr is DbDataReader db)
					await db.CloseAsync();
				else
					rdr.Close();
			}
		}
	}

	public virtual async Task<IDataReader?> OpenReader(IStorageCommand command)
	{
		await EnsureOpen(false);

		var com = await ResolveCommand(command);

		SetupParameters(command, com);

		if (command.Operation.Parameters is not null)
		{
			foreach (var i in command.Operation.Parameters)
				SetParameterValue(com, i.Name, i.Value);
		}

		return com.ExecuteReader();
	}

	protected virtual async Task<IDbCommand?> ResolveCommand(IStorageCommand command)
	{
		if (Commands.TryGetValue(command, out IDbCommand? existing))
			return existing;

		if (Commands.TryGetValue(command, out IDbCommand? existing2))
			return existing2;

		return await _lock.LockAsync(6, async () =>
		{
			var connection = await GetConnection();

			var r = connection.CreateCommand();

			r.CommandText = command.Operation.CommandText;
			r.CommandType = command.Operation.CommandType;
			r.CommandTimeout = command.Operation.CommandTimeout;

			if (Transaction is not null)
				r.Transaction = Transaction;

			Commands.TryAdd(command, r);

			return r;
		});
	}

	private async Task EnsureOpen(bool createTransaction)
	{
		var connection = await GetConnection();

		if (connection is null || connection.State == ConnectionState.Open)
			return;

		await _lock.LockAsync(7, async () =>
		{
			await Open();

			if (createTransaction && Transaction is null)
			{
				Transaction = connection is DbConnection dbc
					? await dbc.BeginTransactionAsync(Context is null ? CancellationToken.None : Context.CancellationToken)
					: connection.BeginTransaction(IsolationLevel.ReadCommitted);
			}
		});
	}
	public void Dispose()
	{
		OnDispose(true);
		GC.SuppressFinalize(this);
	}

	protected virtual void OnDispose(bool disposing)
	{
		AsyncUtils.RunSync(() => Close());

		if (_commands is not null)
		{
			foreach (var command in _commands)
				command.Value.Dispose();

			_commands = null;
		}

		if (Transaction is not null)
		{
			try
			{
				AsyncUtils.RunSync(Rollback);

				Transaction.Dispose();
			}
			catch { }

			Transaction = null;
		}

		if (_connection is not null)
		{
			_connection.Dispose();
			_connection = null;
		}
	}
	public async ValueTask DisposeAsync()
	{
		await OnDisposeAsyncCore().ConfigureAwait(false);

		OnDispose(false);
		GC.SuppressFinalize(this);
	}

	protected virtual async ValueTask OnDisposeAsyncCore()
	{
		await Close().ConfigureAwait(false);

		if (_commands is not null)
		{
			foreach (var command in _commands)
			{
				if (command.Value is DbCommand db)
					await db.DisposeAsync().ConfigureAwait(false);
				else
					command.Value.Dispose();
			}

			_commands = null;
		}

		if (Transaction is not null)
		{
			try
			{
				//No way to check if possible
				await Rollback().ConfigureAwait(false);

				if (Transaction is DbTransaction dbt)
					await dbt.DisposeAsync().ConfigureAwait(false);
				else
					Transaction.Dispose();
			}
			catch { }

			Transaction = null;
		}

		if (_connection is not null)
		{
			if (_connection is DbConnection dbc)
				await dbc.DisposeAsync().ConfigureAwait(false);
			else
				_connection.Dispose();

			_connection = null;
		}
	}
}