import { CollectionViewer, DataSource } from '@angular/cdk/collections';
import { FlatTreeControl } from '@angular/cdk/tree';
import {
  BehaviorSubject,
  merge,
  Observable,
  Subject,
  Subscription
} from 'rxjs';
import { map, share } from 'rxjs/operators';
import { MatTreeFlattener } from '@angular/material';
import { CdkVirtualScrollViewport } from '@angular/cdk/scrolling';

export class VirtualScrollTreeFlatDataSource<T, F> extends DataSource<F> {
  private flattenedDataSubject = new BehaviorSubject<F[]>([]);

  expandedDataSubject = new BehaviorSubject<F[]>([]);
  expandedData$: Observable<F[]>;

  private connectSubject: Subject<F[]>;
  private dataChangeSubscription: Subscription;
  private rangeChangeSubscription: Subscription;

  private rangeSubject: BehaviorSubject<{
    start: number;
    end: number;
  }>;

  // tslint:disable-next-line: variable-name
  private _cdkVirtualScrollViewport: CdkVirtualScrollViewport;
  set cdkVirtualScrollViewport(
    cdkVirtualScrollViewport: CdkVirtualScrollViewport
  ) {
    if (!cdkVirtualScrollViewport) {
      return;
    }
    this._cdkVirtualScrollViewport = cdkVirtualScrollViewport;
    this.rangeSubject = new BehaviorSubject<{ start: number; end: number }>({
      start: 0,
      end: 1
    });

    this.rangeChangeSubscription = cdkVirtualScrollViewport.renderedRangeStream.subscribe(
      range => {
        this.rangeSubject.next({
          start: range.start,
          end: range.end
        });
        if (!!this.connectSubject) {
          this.connectSubject.next(
            this.expandedDataSubject.value.slice(
              this.rangeSubject.value.start,
              this.rangeSubject.value.end
            )
          );
        }
      }
    );
  }

  // tslint:disable-next-line: variable-name
  private _data: BehaviorSubject<T[]>;
  get data() {
    return this._data.value;
  }
  set data(value: T[]) {
    this._data.next(value);
    this.flattenedDataSubject.next(this.treeFlattener.flattenNodes(this.data));
    this.treeControl.dataNodes = this.flattenedDataSubject.value;
  }

  constructor(
    private treeControl: FlatTreeControl<F>,
    private treeFlattener: MatTreeFlattener<T, F>,
    initialData: T[] = []
  ) {
    super();
    this._data = new BehaviorSubject<T[]>(initialData);
    this.expandedData$ = this.expandedDataSubject.asObservable().pipe(share());
  }

  connect(collectionViewer: CollectionViewer): Observable<F[]> {
    this.connectSubject = new Subject<F[]>();

    this.dataChangeSubscription = merge(
      collectionViewer.viewChange,
      this.treeControl.expansionModel.changed,
      this.flattenedDataSubject
    )
      .pipe(
        map(() => {
          this.expandedDataSubject.next(
            this.treeFlattener.expandFlattenedNodes(
              this.flattenedDataSubject.value,
              this.treeControl
            )
          );

          return !this.rangeSubject
            ? this.expandedDataSubject.value
            : this.expandedDataSubject.value.slice(
                this.rangeSubject.value.start,
                this.rangeSubject.value.end
              );
        })
      )
      .subscribe(datas => {
        this.connectSubject.next(datas);
        if (!!this._cdkVirtualScrollViewport) {
          this._cdkVirtualScrollViewport.checkViewportSize();
        }
      });

    return this.connectSubject.asObservable();
  }

  disconnect() {
    console.log('VirtualScrollTreeFlatDataSource disconnect');
    if (!!this.connectSubject) {
      this.connectSubject.next();
      this.connectSubject.unsubscribe();
    }

    if (!!this.dataChangeSubscription) {
      this.dataChangeSubscription.unsubscribe();
    }

    if (!!this.rangeChangeSubscription) {
      this.rangeChangeSubscription.unsubscribe();
    }
  }
}
