import type {
  MapExecutionFrame,
  MapExecutionScope,
  PlayCheckpoint,
  PlayExecutionEvent,
} from './ctx-types';

export function cloneMapFrame(frame: MapExecutionFrame): MapExecutionFrame {
  return {
    ...frame,
    completedRowKeys: [...frame.completedRowKeys],
    pendingRowKeys: [...frame.pendingRowKeys],
  };
}

export class MapExecutionFrameStore {
  constructor(
    private readonly checkpoint: PlayCheckpoint,
    private readonly emitExecutionEvent: (event: PlayExecutionEvent) => void,
  ) {}

  set(frame: MapExecutionFrame): void {
    this.checkpoint.mapFrames = {
      ...(this.checkpoint.mapFrames ?? {}),
      [frame.mapInvocationId]: cloneMapFrame(frame),
    };
  }

  start(input: {
    scope: MapExecutionScope;
    totalRows: number;
    completedRowKeys: readonly string[];
    pendingRowKeys: readonly string[];
  }): void {
    const now = Date.now();
    this.set({
      mapInvocationId: input.scope.mapInvocationId,
      mapNodeId: input.scope.mapNodeId ?? null,
      logicalNamespace: input.scope.logicalNamespace,
      artifactTableNamespace: input.scope.artifactTableNamespace,
      status: 'running',
      totalRows: input.totalRows,
      completedRowKeys: [...input.completedRowKeys],
      pendingRowKeys: [...input.pendingRowKeys],
      startedAt: now,
      updatedAt: now,
    });
    this.emitExecutionEvent({
      type: 'map.started',
      mapInvocationId: input.scope.mapInvocationId,
      mapNodeId: input.scope.mapNodeId ?? null,
      logicalNamespace: input.scope.logicalNamespace,
      artifactTableNamespace: input.scope.artifactTableNamespace,
      totalRows: input.totalRows,
      completedRows: input.completedRowKeys.length,
      pendingRows: input.pendingRowKeys.length,
      at: Date.now(),
    });
  }

  updateProgress(input: {
    scope: MapExecutionScope;
    totalRows: number;
    status?: MapExecutionFrame['status'];
    completedRowKey?: string | null;
    pendingRowKey?: string | null;
    activeBoundaryId?: string | null;
    emitEventType?: PlayExecutionEvent['type'];
  }): void {
    const existing =
      this.checkpoint.mapFrames?.[input.scope.mapInvocationId] ?? null;
    if (!existing) {
      return;
    }
    const completedRowKeys = new Set(existing.completedRowKeys);
    const pendingRowKeys = new Set(existing.pendingRowKeys);
    const completedBefore = completedRowKeys.size;
    if (input.completedRowKey?.trim()) {
      const completedRowKey = input.completedRowKey.trim();
      completedRowKeys.add(completedRowKey);
      pendingRowKeys.delete(completedRowKey);
    }
    if (input.pendingRowKey?.trim()) {
      pendingRowKeys.add(input.pendingRowKey.trim());
    }
    const nextFrame: MapExecutionFrame = {
      ...existing,
      status: input.status ?? existing.status,
      completedRowKeys: [...completedRowKeys],
      pendingRowKeys: [...pendingRowKeys],
      ...(input.activeBoundaryId !== undefined
        ? { activeBoundaryId: input.activeBoundaryId }
        : {}),
      updatedAt: Date.now(),
    };
    this.set(nextFrame);
    const progressEventType =
      input.emitEventType ??
      (completedRowKeys.size > completedBefore ? 'map.progress' : null);
    if (progressEventType) {
      this.emitExecutionEvent({
        type: progressEventType,
        mapInvocationId: input.scope.mapInvocationId,
        mapNodeId: input.scope.mapNodeId ?? null,
        logicalNamespace: input.scope.logicalNamespace,
        artifactTableNamespace: input.scope.artifactTableNamespace,
        completedRows: nextFrame.completedRowKeys.length,
        failedRows: Math.max(
          0,
          input.totalRows -
            nextFrame.completedRowKeys.length -
            nextFrame.pendingRowKeys.length,
        ),
        totalRows: input.totalRows,
        at: Date.now(),
      } as PlayExecutionEvent);
    }
  }
}