spark/src/DataStreamSuite/DataStream/DataStreamExtensions.cs

namespace Bts.DataStream;

using BCG;

using bts.api.core;

using Bts.DataStreamCore;
using Bts.PubSub;

using Microsoft.Extensions.Logging;

using Spark.Core.Disposables;
using Spark.Core.Threading;

using System;
using System.Linq;
using System.Reactive.Linq;

public static class DataStreamExtensions {
   public static ViewStream AsViewStream(this IObservable<IDataPacket> source) =>
      source as ViewStream ?? new ObservableViewStream(source);

   public static ViewStream AsViewStream<T>(this IObservable<IDataPacket<T>> source, string? name = null, ILogger? logger = null) where T : IBlueRow =>
      source as ViewStream ?? new ObservableViewStream(source.Select(dp => new DataPacket(dp.newRows, dp.addedRows, dp.deletedRows)), name, logger);

   public static IObservable<IDataPacket> AsObservable(this ViewStream source) {
      return source as IObservable<IDataPacket> ?? Observable.Create<IDataPacket>(Subscribe);

      IDisposable Subscribe(IObserver<IDataPacket> observer) {
         source.onChange += OnChange;
         return Disposable.Create(() => source.onChange -= OnChange);

         void OnChange(object? newValue) =>
            observer.OnNext(newValue as IDataPacket ?? throw new ArgumentException(
               $"Expected IDataPacket, got {newValue?.GetType()}", nameof(newValue)));
      }
   }

   // public static IObservable<IDataPacket<T>> SelectDataPacket<T>(this IObservable<IDataPacket> source, Func<IBlueRow, T> selector) where T : IBlueRow =>
   //    source.Select(dp => {
   //       dp.get_new_added_deleted(out var newRows, out var addedRows, out var deletedRows);
   //       return new DataPacket<T>(CastRows(newRows), CastRows(addedRows), CastRows(deletedRows));

   //       IBlueRowCollection<T> CastRows(ICollection rows) =>
   //          new BlueTableSelectorWrapper<T>(
   //             rows as IBlueRowCollection ?? new BlueRowCollection(rows, dp.Schema), selector);
   //    });

   public static IObservable<TableUpdate<IBlueRow, int>> AsTables(this IObservable<IDataPacket> source) =>
      source.AsTables<IBlueRow, int>(r => r, r => r.Id());

   public static IObservable<TableUpdate<T, int>> AsTables<T>(this IObservable<IDataPacket> source, Func<IBlueRow, T> valueSelector) where T : IBlueRow =>
      source.AsTables<T, int>(valueSelector, r => r.Id());

   public static IObservable<TableUpdate<T, int>> AsTables<T>(
      this IObservable<IDataPacket> source, Func<IBlueRow, T> valueSelector, Func<T, int> keySelector) =>
      source.AsTables<T, int>(valueSelector, keySelector);

   public static IObservable<TableUpdate<T, TKey>> AsTables<T, TKey>(
      this IObservable<IDataPacket> source, Func<IBlueRow, T> valueSelector, Func<T, TKey> keySelector, Func<IBlueRow, TKey>? rowKeySelector = null) where TKey : IEquatable<TKey> {
      rowKeySelector ??= r => keySelector(valueSelector(r));
      return source.Scan(TableUpdate<T, TKey>.Empty(keySelector), (res, dp) => {
         var bd = res.ToBuilder();
         dp.get_new_added_deleted(out var newRows, out var addedRows, out var deletedRows);
         foreach (IBlueRow r in addedRows) {
            bd.AddOrUpdate(valueSelector(r));
         }
         foreach (IBlueRow r in deletedRows) {
            bd.Remove(rowKeySelector(r));
         }
         return bd.ToImmutable();
      });
   }

   public static IObservable<IDataPacket> AsDataPackets(this IObservable<TableUpdate<IBlueRow, int>> source) =>
      source.Scan(
         (Table: (BlueTableImmutable?)null, Dp: (IDataPacket?)null),
         (res, table) => {
            var resTable = res.Table;
            if (resTable is null) {
               if (table.Index.Count == 0) {
                  return (null, null);
               }
               resTable = new(table.Index.Values.First().schema);
            }
            var bd = resTable.ToDataPacketBuilder();
            foreach (var r in table.RemovedItems) {
               bd.Remove(r.Id());
            }
            foreach (var r in table.AddedItems) {
               bd.AddOrUpdate(r);
            }
            var resDp = bd.ToDataPacket(out var newRows);
            return (newRows, resDp);
         })
         .Where(t => t.Dp is not null && (t.Dp.addedRows.Count > 0 || t.Dp.deletedRows.Count > 0))
         .Select(t => t.Dp!);

   public static IObservable<IDataPacket> AsDataPackets<T>(this IObservable<TableUpdate<T, int>> source, IBlueTableSchema schema) where T : IBlueRow =>
      source.Scan(
         (Table: new BlueTableImmutable(schema), Dp: (IDataPacket)new DataPacket(schema)),
         (res, table) => {
            var bd = res.Table.ToDataPacketBuilder();
            foreach (var r in table.RemovedItems) {
               bd.Remove(r.Id());
            }
            foreach (var r in table.AddedItems) {
               bd.AddOrUpdate(r);
            }
            var resDp = bd.ToDataPacket(out var newRows);
            return (newRows, resDp);
         })
         .Where(t => t.Dp.addedRows.Count > 0 || t.Dp.deletedRows.Count > 0)
         .Select(t => t.Dp);

   public static IObservable<IDataPacket> AsDataPackets<T, TKey>(this IObservable<TableUpdate<T, TKey>> source, BlueTableSchemaAdapter<T> schema) where TKey : IEquatable<TKey> {
      return Observable.Create<IDataPacket>(Subscribe);

      IDisposable Subscribe(IObserver<IDataPacket> o) {
         var adapter = new BlueTableImmutableAdapter<T>(schema);
         IDataPacket initDataPacket = new DataPacket(schema);
         return source.Scan((Adapter: adapter, DataPacket: initDataPacket), OnUpdate).Select(res => res.DataPacket).Subscribe(o);

         (BlueTableImmutableAdapter<T> Adapter, IDataPacket DataPacket) OnUpdate(
            (BlueTableImmutableAdapter<T> Adapter, IDataPacket DataPacket) state, TableUpdate<T, TKey> tableUpdate) {
            var bd = state.Adapter.ToDataPacketBuilder();
            bd.Update(tableUpdate);
            var newRows = bd.ToImmutable(out var dp);
            return (newRows, dp);
         }
      }
   }

   public static ViewStream AsViewStream<T>(this IObservable<TableUpdate<T, int>> source, IBlueTableSchema schema) where T : IBlueRow =>
      source.AsDataPackets(schema).AsViewStream();

   public static int ResolveIndices(this IBlueTableSchema schema, string name) {
      var indices = ResolveIndices(schema, new[] { name });
      return indices[0];
   }

   public static (int, int) ResolveIndices(this IBlueTableSchema schema, string name1, string name2) {
      var indices = ResolveIndices(schema, new[] { name1, name2 });
      return (indices[0], indices[1]);
   }

   public static (int, int, int) ResolveIndices(this IBlueTableSchema schema, string name1, string name2, string name3) {
      var indices = ResolveIndices(schema, new[] { name1, name2, name3 });
      return (indices[0], indices[1], indices[2]);
   }

   public static (int, int, int, int) ResolveIndices(this IBlueTableSchema schema, string name1, string name2, string name3, string name4) {
      var indices = ResolveIndices(schema, new[] { name1, name2, name3, name4 });
      return (indices[0], indices[1], indices[2], indices[3]);
   }

   public static int[] ResolveIndices(this IBlueTableSchema schema, string[] columnNames) {
      var columns = schema.columnInfo;
      var indices = new int[columnNames.Length];
      int nResolved = 0;
      for (int i = 0; i < columnNames.Length; i++) {
         int origNResolved = nResolved;
         for (int j = 0; j < columns.Length; j++) {
            if (columns[j].colName == columnNames[i]) {
               indices[i] = j;
               nResolved++;
               break;
            }
         }
         if (nResolved == origNResolved) {
            throw new ArgumentException($"No column named '{columnNames[i]}'", nameof(schema));
         }
      }
      return indices;
   }

   public static IObservable<IDataPacket> OnWorkScheduler(this IObservable<IDataPacket> source, IWorkScheduler scheduler) =>
      Observable.Create<IDataPacket>(o => {
         var inputGate = new AtomicGate();
         var outputGate = new AtomicGate();
         BlueTableImmutable? lastTable = null;
         BlueTableImmutable.DataPacketBuilder? bd = null;
         IDataPacket? pendingDpOut = null;
         return Disposable.Create(source.Subscribe(OnNext), scheduler.Subscribe(DoWork));

         void OnNext(IDataPacket value) {
            using (inputGate.BeginTransaction()) {
               lastTable ??= new(value.newRows.schema);
               bd ??= lastTable.ToDataPacketBuilder();
               if (pendingDpOut is null) {
                  // The last outgoing datapacket was just fired. Don't merge this update with the
                  // previous one (otherwise its adds/removes would be fired twice).
                  bd.Reset();
               }
               foreach (IBlueRow r in value.deletedRows) {
                  bd.Remove(r.Id());
               }
               foreach (IBlueRow r in value.addedRows) {
                  bd.AddOrUpdate(r);
               }
               pendingDpOut = bd.ToDataPacket(out lastTable);
            }
         }

         void DoWork() {
            IDataPacket? dpOut = null;
            if (outputGate.TryEnter()) {
               using var exit = Disposable.Create(outputGate.Exit);
               (dpOut, pendingDpOut) = (pendingDpOut, dpOut);
               if (dpOut is not null) {
                  o.OnNext(dpOut);
               }
            }
         }
      });


}